docker-python/tests/test_bigquery.py at master · HelWireless/docker-python · GitHub

History

86 lines (71 loc) · 3.51 KB

Raw

import unittest

import os

import threading

from test.support import EnvironmentVarGuard

from urllib.parse import urlparse

from http.server import BaseHTTPRequestHandler, HTTPServer

from google.cloud import bigquery

from google.auth.exceptions import DefaultCredentialsError

from kaggle_gcp import KaggleKernelCredentials, PublicBigqueryClient

class TestBigQuery(unittest.TestCase):

def _test_proxy(self, client, should_use_proxy):

class HTTPHandler(BaseHTTPRequestHandler):

called = False

header_found = False

def do_HEAD(s):

s.send_response(200)

def do_GET(s):

HTTPHandler.called = True

HTTPHandler.header_found = any(k for k in s.headers if k == "X-KAGGLE-PROXY-DATA" and s.headers[k] == "test-key")

s.send_response(200)

server_address = urlparse(os.getenv('KAGGLE_DATA_PROXY_URL'))

with HTTPServer((server_address.hostname, server_address.port), HTTPHandler) as httpd:

threading.Thread(target=httpd.serve_forever).start()

try:

for ds in client.list_datasets(): pass

except:

pass

httpd.shutdown()

if should_use_proxy:

self.assertTrue(HTTPHandler.called, msg="Fake server did not receive a request from the BQ client.")

self.assertTrue(HTTPHandler.header_found, msg="X-KAGGLE-PROXY-DATA header was missing from the BQ request.")

else:

self.assertFalse(HTTPHandler.called, msg="Fake server was called from the BQ client, but should not have been.")

def test_proxy_using_library(self):

env = EnvironmentVarGuard()

env.unset('KAGGLE_BQ_USER_JWT')

with env:

client = PublicBigqueryClient()

self._test_proxy(client, should_use_proxy=True)

def test_proxy_no_project(self):

env = EnvironmentVarGuard()

env.unset('KAGGLE_BQ_USER_JWT')

with env:

client = bigquery.Client()

self._test_proxy(client, should_use_proxy=True)

def test_project_with_connected_account(self):

env = EnvironmentVarGuard()

env.set('KAGGLE_BQ_USER_JWT', 'foobar')

with env:

client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())

self._test_proxy(client, should_use_proxy=False)

def test_simultaneous_clients(self):

env = EnvironmentVarGuard()

env.set('KAGGLE_BQ_USER_JWT', 'foobar')

with env:

proxy_client = bigquery.Client()

self._test_proxy(proxy_client, should_use_proxy=True)

bq_client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())

self._test_proxy(bq_client, should_use_proxy=False)

# Verify that proxy client is still going to proxy to ensure global Connection

# isn't being modified.

self._test_proxy(proxy_client, should_use_proxy=True)

def test_no_project_with_connected_account(self):

env = EnvironmentVarGuard()

env.set('KAGGLE_BQ_USER_JWT', 'foobar')

with env:

with self.assertRaises(DefaultCredentialsError):

# TODO(vimota): Handle this case, either default to Kaggle Proxy or use some default project

# by the user or throw a custom exception.

client = bigquery.Client(credentials=KaggleKernelCredentials())

self._test_proxy(client, should_use_proxy=False)

docker-python/tests/test_bigquery.py at master · HelWireless/docker-python · GitHub

Expand file tree

Breadcrumbs

test_bigquery.py

Latest commit

History

Breadcrumbs

test_bigquery.py

File metadata and controls