From 276618a5e6d2952774cad7454cdecd7e50558f8c Mon Sep 17 00:00:00 2001 From: Vincent Roseberry Date: Tue, 20 Oct 2020 20:25:23 +0000 Subject: [PATCH 1/2] Add new client to retrieve exportable .ipynb of current session. http://b/171318603 --- Dockerfile | 1 + patches/kaggle_session.py | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+) create mode 100644 patches/kaggle_session.py diff --git a/Dockerfile b/Dockerfile index 6fbcebb9..b79d3521 100644 --- a/Dockerfile +++ b/Dockerfile @@ -488,6 +488,7 @@ RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.c ENV PYTHONUSERBASE "/root/.local" ADD patches/kaggle_gcp.py /root/.local/lib/python3.7/site-packages/kaggle_gcp.py ADD patches/kaggle_secrets.py /root/.local/lib/python3.7/site-packages/kaggle_secrets.py +ADD patches/kaggle_session.py /root/.local/lib/python3.7/site-packages/kaggle_session.py ADD patches/kaggle_web_client.py /root/.local/lib/python3.7/site-packages/kaggle_web_client.py ADD patches/kaggle_datasets.py /root/.local/lib/python3.7/site-packages/kaggle_datasets.py ADD patches/log.py /root/.local/lib/python3.7/site-packages/log.py diff --git a/patches/kaggle_session.py b/patches/kaggle_session.py new file mode 100644 index 00000000..30679c86 --- /dev/null +++ b/patches/kaggle_session.py @@ -0,0 +1,26 @@ +""" +This library adds support for retrieving data related to the current user session. +""" + +import os + +from kaggle_web_client import KaggleWebClient + + +class UserSessionClient(): + GET_SOURCE_ENDPOINT = '/requests/GetKernelRunSourceForCaipRequest' + + def __init__(self): + self.web_client = KaggleWebClient() + + def get_exportable_ipynb(self): + """Fetch the .ipynb source of the current notebook session. + + If Kaggle datasets are attached to the notebook, the source will + include an additonnal cell with logic to download the datasets + outside the Kaggle platform. + """ + request_body = { + 'UseDraft': True, + } + return self.web_client.make_post_request(request_body, self.GET_SOURCE_ENDPOINT) From 2b0074e69ce07896065ca003785902f09c64cc11 Mon Sep 17 00:00:00 2001 From: Vincent Roseberry Date: Tue, 20 Oct 2020 22:09:12 +0000 Subject: [PATCH 2/2] Add unit tests --- tests/test_user_session.py | 104 +++++++++++++++++++++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 tests/test_user_session.py diff --git a/tests/test_user_session.py b/tests/test_user_session.py new file mode 100644 index 00000000..57f2de8c --- /dev/null +++ b/tests/test_user_session.py @@ -0,0 +1,104 @@ +import json +import os +import threading +import unittest + +from http.server import BaseHTTPRequestHandler, HTTPServer +from test.support import EnvironmentVarGuard +from urllib.parse import urlparse +from kaggle_session import UserSessionClient +from kaggle_web_client import (_KAGGLE_URL_BASE_ENV_VAR_NAME, + _KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME, + CredentialError, BackendError) + +class UserSessionHTTPHandler(BaseHTTPRequestHandler): + + def set_request(self): + raise NotImplementedError() + + def get_response(self): + raise NotImplementedError() + + def do_HEAD(s): + s.send_response(200) + + def do_POST(s): + s.set_request() + s.send_response(200) + s.send_header("Content-type", "application/json") + s.end_headers() + s.wfile.write(json.dumps(s.get_response()).encode("utf-8")) + +class TestUserSessionClient(unittest.TestCase): + SERVER_ADDRESS = urlparse(os.getenv(_KAGGLE_URL_BASE_ENV_VAR_NAME, default="http://127.0.0.1:8001")) + TEST_JWT = 'test-secrets-key' + + def _test_client(self, client_func, expected_path, expected_body, source=None, success=True): + _request = {} + + class GetKernelRunSourceForCaipHandler(UserSessionHTTPHandler): + def set_request(self): + _request['path'] = self.path + content_len = int(self.headers.get('Content-Length')) + _request['body'] = json.loads(self.rfile.read(content_len)) + _request['headers'] = self.headers + + def get_response(self): + if success: + return {'result': {'source': source}, 'wasSuccessful': 'true'} + return {'wasSuccessful': 'false'} + + env = EnvironmentVarGuard() + env.set(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME, self.TEST_JWT) + with env: + with HTTPServer((self.SERVER_ADDRESS.hostname, self.SERVER_ADDRESS.port), GetKernelRunSourceForCaipHandler) as httpd: + threading.Thread(target=httpd.serve_forever).start() + + try: + client_func() + finally: + httpd.shutdown() + + path, headers, body = _request['path'], _request['headers'], _request['body'] + self.assertEqual( + path, + expected_path, + msg="Fake server did not receive the right request from UserSessionClient.") + self.assertEqual( + body, + expected_body, + msg="Fake server did not receive the right body from UserSessionClient.") + + def test_no_token_fails(self): + env = EnvironmentVarGuard() + env.unset(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME) + with env: + with self.assertRaises(CredentialError): + client = UserSessionClient() + + def test_get_exportable_ipynb_succeeds(self): + source = "import foo" + + def call_get_ipynb(): + client = UserSessionClient() + response = client.get_exportable_ipynb() + self.assertEqual(source, response['source']) + + self._test_client( + call_get_ipynb, + '/requests/GetKernelRunSourceForCaipRequest', + {'UseDraft': True}, + source=source, + success=True) + + def test_get_exportable_ipynb_fails(self): + def call_get_ipynb(): + client = UserSessionClient() + with self.assertRaises(BackendError): + client.get_exportable_ipynb() + + self._test_client( + call_get_ipynb, + '/requests/GetKernelRunSourceForCaipRequest', + {'UseDraft': True}, + success=False) \ No newline at end of file