test_user_session.py

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)