| 24 | self.end_headers() |
| 25 | |
| 26 | class TestGoogleGenerativeAiPatch(unittest.TestCase): |
| 27 | endpoint = "http://127.0.0.1:80" |
| 28 | |
| 29 | def test_proxy_enabled(self): |
| 30 | env = EnvironmentVarGuard() |
| 31 | secrets_token = "secrets_token" |
| 32 | proxy_token = "proxy_token" |
| 33 | env.set("KAGGLE_USER_SECRETS_TOKEN", secrets_token) |
| 34 | env.set("KAGGLE_DATA_PROXY_TOKEN", proxy_token) |
| 35 | env.set("KAGGLE_DATA_PROXY_URL", self.endpoint) |
| 36 | env.set("KAGGLE_GRPC_DATA_PROXY_URL", "http://127.0.0.1:50001") |
| 37 | env.set("KAGGLE_GOOGLE_GENERATIVE_AI_USE_REST_ONLY", "True") |
| 38 | server_address = urlparse(self.endpoint) |
| 39 | with env: |
| 40 | with HTTPServer((server_address.hostname, server_address.port), HTTPHandler) as httpd: |
| 41 | threading.Thread(target=httpd.serve_forever).start() |
| 42 | import google.generativeai as palm |
| 43 | api_key = "NotARealAPIKey" |
| 44 | palm.configure(api_key = api_key) |
| 45 | try: |
| 46 | for _ in palm.list_models(): |
| 47 | pass |
| 48 | except: |
| 49 | pass |
| 50 | httpd.shutdown() |
| 51 | self.assertTrue(HTTPHandler.called) |
| 52 | self.assertIn("/palmapi", HTTPHandler.path) |
| 53 | self.assertEqual(proxy_token, HTTPHandler.headers["x-kaggle-proxy-data"]) |
| 54 | self.assertEqual("Bearer {}".format(secrets_token), HTTPHandler.headers["x-kaggle-authorization"]) |
| 55 | self.assertEqual(api_key, HTTPHandler.headers["x-goog-api-key"]) |
nothing calls this directly
no outgoing calls
no test coverage detected