Make BigQuery tests more robust: use scoped environment variables and scoped HTTPRequestHandler so tests don't affect each other.

Author: vimota

patches/sitecustomize.py

@@ -43,7 +43,7 @@ def monkeypatch_bq(bq_client, *args, **kwargs):
         if kwargs.get('credentials') is not None:
             # The user wants to use their own credentials scheme, don't try to interfere.
             return bq_client(*args, **kwargs)
-        print("Using connected BigQuery Account.")
+        print("Using enabled BigQuery integration.")
         kwargs['credentials'] = KaggleKernelCredentials()
         kwargs['project'] = kwargs.get('project')
         return bq_client(

tests/test_bigquery.py

@@ -1,17 +1,31 @@
 import unittest
 import os
 import threading
+from test.support import EnvironmentVarGuard
 
 from http.server import BaseHTTPRequestHandler, HTTPServer
 
 from google.cloud import bigquery
+from google.auth.exceptions import DefaultCredentialsError
 
 HOSTNAME = "127.0.0.1"
 PORT = 8000
 URL = "http://%s:%s" % (HOSTNAME, PORT)
 
 class TestBigQuery(unittest.TestCase):
-    def _test_proxy(self, client):
+    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)
+
         with HTTPServer((HOSTNAME, PORT), HTTPHandler) as httpd:
             threading.Thread(target=httpd.serve_forever).start()
 
@@ -21,25 +35,33 @@ def _test_proxy(self, client):
                 pass
 
             httpd.shutdown()
-            self.assertTrue(HTTPHandler.called, msg="Fake server did not recieve a request from the BQ client.")
-            self.assertTrue(HTTPHandler.header_found, msg="X-KAGGLE-PROXY-DATA header was missing from the BQ request.")
+            if should_use_proxy:
+                self.assertTrue(HTTPHandler.called, msg="Fake server did not recieve 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_kaggle_project(self):
         client = bigquery.Client(project='KAGGLE')
-        self._test_proxy(client)
+        self._test_proxy(client, should_use_proxy=True)
 
     def test_proxy_no_project(self):
         client = bigquery.Client()
-        self._test_proxy(client)
-
-class HTTPHandler(BaseHTTPRequestHandler):
-    called = False
-    header_found = False
+        self._test_proxy(client, should_use_proxy=True)
 
-    def do_HEAD(s):
-        s.send_response(200)
+    def test_project_with_connected_account(self):
+        env = EnvironmentVarGuard()
+        env.set('KAGGLE_BQ_USER_JWT', 'foobar')
+        with env:
+            client = bigquery.Client(project='ANOTHER_PROJECT')
+            self._test_proxy(client, should_use_proxy=False)
 
-    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)
+    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()
+                self._test_proxy(client, should_use_proxy=False)