FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Addressed mostly nits from code review and made tests safer but unset… · feibyte/docker-python@ba8e462 · GitHub

/ docker-python Public
forked from Kaggle/docker-python

Commit ba8e462

Browse files
vimota
committed
Addressed mostly nits from code review and made tests safer but unsetting env var in tests that expected none set.
1 parent ddf0170 commit ba8e462

2 files changed

Lines changed: 18 additions & 10 deletions

File tree

‎patches/sitecustomize.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
import os
21
from google.auth import credentials
32
from google.cloud import bigquery
43
from google.cloud.bigquery._http import Connection
4+
import os
5+
56

67
class KaggleKernelCredentials(credentials.Credentials):
78
def refresh(self, request):
@@ -11,16 +12,18 @@ def refresh(self, request):
1112

1213
kaggle_proxy_token = os.getenv("KAGGLE_DATA_PROXY_TOKEN")
1314
CONNECTION_BASE_URL = Connection.API_BASE_URL
15+
16+
1417
def monkeypatch_bq(bq_client, *args, **kwargs):
1518
data_proxy_project = os.getenv("KAGGLE_DATA_PROXY_PROJECT")
1619
bq_user_jwt = os.getenv("KAGGLE_BQ_USER_JWT")
17-
specified_project = kwargs.get('project').lower() if 'project' in kwargs else None
20+
specified_project = kwargs['project'] if 'project' in kwargs else None
1821
# Use Data Proxy if user has specified to use the Kaggle project, or if
1922
# there are no connected GCP accounts (to maintain backwards compatibility).
20-
if bq_user_jwt is None and specified_project and specified_project != 'kaggle':
23+
if bq_user_jwt is None and specified_project and specified_project.lower() != 'kaggle':
2124
raise Exception("In order to query a private BigQuery project, please connect a GCP account. "
2225
"Otherwise specify 'kaggle' as the project to use Kaggle's public dataset BigQuery integration.")
23-
use_data_proxy = specified_project == 'kaggle' or bq_user_jwt is None
26+
use_data_proxy = (specified_project and specified_project.lower() == 'kaggle') or bq_user_jwt is None
2427
if use_data_proxy:
2528
if data_proxy_project is None or kaggle_proxy_token is None:
2629
# We don't have the data proxy info so leave the bq client unmodified.
@@ -43,7 +46,6 @@ def monkeypatch_bq(bq_client, *args, **kwargs):
4346
return bq_client(*args, **kwargs)
4447
print("Using enabled BigQuery integration.")
4548
kwargs['credentials'] = KaggleKernelCredentials()
46-
kwargs['project'] = kwargs.get('project')
4749
return bq_client(
4850
*args,
4951
**kwargs)

‎tests/test_bigquery.py‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,18 +36,24 @@ def do_GET(s):
3636

3737
httpd.shutdown()
3838
if should_use_proxy:
39-
self.assertTrue(HTTPHandler.called, msg="Fake server did not recieve a request from the BQ client.")
39+
self.assertTrue(HTTPHandler.called, msg="Fake server did not receive a request from the BQ client.")
4040
self.assertTrue(HTTPHandler.header_found, msg="X-KAGGLE-PROXY-DATA header was missing from the BQ request.")
4141
else:
4242
self.assertFalse(HTTPHandler.called, msg="Fake server was called from the BQ client, but should not have been.")
4343

4444
def test_proxy_kaggle_project(self):
45-
client = bigquery.Client(project='KAGGLE')
46-
self._test_proxy(client, should_use_proxy=True)
45+
env = EnvironmentVarGuard()
46+
env.unset('KAGGLE_BQ_USER_JWT')
47+
with env:
48+
client = bigquery.Client(project='KAGGLE')
49+
self._test_proxy(client, should_use_proxy=True)
4750

4851
def test_proxy_no_project(self):
49-
client = bigquery.Client()
50-
self._test_proxy(client, should_use_proxy=True)
52+
env = EnvironmentVarGuard()
53+
env.unset('KAGGLE_BQ_USER_JWT')
54+
with env:
55+
client = bigquery.Client()
56+
self._test_proxy(client, should_use_proxy=True)
5157

5258
def test_project_with_connected_account(self):
5359
env = EnvironmentVarGuard()

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL