Method _test_client

tests/test_datasets.py:44–103

(self, client_func, expected_path, expected_body, is_tpu=True, success=True, iap_token=False)

    SERVER_ADDRESS = urlparse(os.getenv(_KAGGLE_URL_BASE_ENV_VAR_NAME, default="http://127.0.0.1:0"))

    def _test_client(self, client_func, expected_path, expected_body, is_tpu=True, success=True, iap_token=False):
        _request = {}

        class GetGcsPathHandler(GcsDatasetsHTTPHandler):

            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:
                    gcs_path = _TPU_GCS_BUCKET if is_tpu else _AUTOML_GCS_BUCKET
                    return {'result': {
                        'destinationBucket': gcs_path,
                        'destinationPath': None}, 'wasSuccessful': "true"}
                else:
                    return {'wasSuccessful': "false"}

        env = EnvironmentVarGuard()
        env.set(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME, _TEST_JWT)
        if is_tpu:
          env.set(_KAGGLE_TPU_NAME_ENV_VAR_NAME, 'FAKE_TPU')
        if iap_token:
          env.set(_KAGGLE_IAP_TOKEN_ENV_VAR_NAME, _TEST_IAP)
        with env:
            with HTTPServer((self.SERVER_ADDRESS.hostname, self.SERVER_ADDRESS.port), GetGcsPathHandler) as httpd:
                threading.Thread(target=httpd.serve_forever).start()

                try:
                    # Make sure this server is pointed to by clients (the port may have been dynamically chosen): 
                    env.set(_KAGGLE_URL_BASE_ENV_VAR_NAME, "http://"+httpd.server_address[0]+":"+str(httpd.server_address[1]))
                    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 the KaggleDatasets client.")
                self.assertEqual(
                    body,
                    expected_body,
                    msg="Fake server did not receive the right body from the KaggleDatasets client.")
                self.assertIn('Content-Type', headers.keys(),
                    msg="Fake server did not receive a Content-Type header from the KaggleDatasets client.")
                self.assertEqual('application/json', headers.get('Content-Type'),
                    msg="Fake server did not receive an application/json content type header from the KaggleDatasets client.")
                self.assertIn('X-Kaggle-Authorization', headers.keys(),
                    msg="Fake server did not receive an X-Kaggle-Authorization header from the KaggleDatasets client.")
                if iap_token:
                    self.assertEqual(f'Bearer {_TEST_IAP}', headers.get('Authorization'),
                        msg="Fake server did not receive an Authorization header from the KaggleDatasets client.")
                else:
                    self.assertNotIn('Authorization', headers.keys(),
                      msg="Fake server received an Authorization header from the KaggleDatasets client. It shouldn't.")