Class TestBigQuery

tests/test_bigquery.py:17–151

import kaggle_secrets

class TestBigQuery(unittest.TestCase):

    API_BASE_URL = "http://127.0.0.1:2121"
    def _test_integration(self, client):
        class HTTPHandler(BaseHTTPRequestHandler):
            called = False
            bearer_header_found = False

            def do_HEAD(self):
                self.send_response(200)

            def do_GET(self):
                HTTPHandler.called = True
                HTTPHandler.bearer_header_found = any(
                    k for k in self.headers if k == "authorization" and self.headers[k] == "Bearer secret")
                self.send_response(200)
                self.send_header("Content-type", "application/json")
                self.end_headers()
                sample_dataset = {
                    "id": "bigqueryproject:datasetname",
                    "datasetReference": {
                        "datasetId": "datasetname",
                        "projectId": "bigqueryproject"
                    }
                }
                self.wfile.write(json.dumps({"kind": "bigquery#datasetList", "datasets": [sample_dataset]}).encode("utf-8"))

        server_address = urlparse(self.API_BASE_URL)
        with HTTPServer((server_address.hostname, server_address.port), HTTPHandler) as httpd:
            threading.Thread(target=httpd.serve_forever).start()

            for dataset in client.list_datasets():
                self.assertEqual(dataset.dataset_id, "datasetname")    
            
            httpd.shutdown()
            self.assertTrue(
                    HTTPHandler.called, msg="Fake server was not called from the BQ client, but should have been.")

            self.assertTrue(
                HTTPHandler.bearer_header_found, msg="authorization header was missing from the BQ request.")
    
    def _setup_mocks(self, api_url_mock):
        api_url_mock.__str__.return_value = self.API_BASE_URL

    @patch.object(kaggle_secrets.UserSecretsClient, 'get_bigquery_access_token', return_value=('secret',1000))
    def test_project_with_connected_account(self, mock_access_token):
        env = EnvironmentVarGuard()
        env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
        with env:
            client = bigquery.Client(
                project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials(), client_options={"api_endpoint": TestBigQuery.API_BASE_URL})
            self._test_integration(client)

    @patch.object(kaggle_secrets.UserSecretsClient, 'get_bigquery_access_token', return_value=('secret',1000))
    def test_project_with_empty_integrations(self, mock_access_token):
        env = EnvironmentVarGuard()
        env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
        env.set('KAGGLE_KERNEL_INTEGRATIONS', '')