|
| 1 | +import json |
| 2 | +import unittest |
| 3 | +import threading |
| 4 | + |
| 5 | +from test.support.os_helper import EnvironmentVarGuard |
| 6 | +from urllib.parse import urlparse |
| 7 | + |
| 8 | +from http.server import BaseHTTPRequestHandler, HTTPServer |
| 9 | + |
| 10 | +class HTTPHandler(BaseHTTPRequestHandler): |
| 11 | + called = False |
| 12 | + path = None |
| 13 | + headers = {} |
| 14 | + |
| 15 | + def do_HEAD(self): |
| 16 | + self.send_response(200) |
| 17 | + |
| 18 | + def do_POST(self): |
| 19 | + HTTPHandler.path = self.path |
| 20 | + HTTPHandler.headers = self.headers |
| 21 | + HTTPHandler.called = True |
| 22 | + self.send_response(200) |
| 23 | + self.send_header("Content-type", "application/json") |
| 24 | + self.end_headers() |
| 25 | + |
| 26 | +class TestGoogleGenAiPatch(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 | + server_address = urlparse(self.endpoint) |
| 37 | + with env: |
| 38 | + with HTTPServer((server_address.hostname, server_address.port), HTTPHandler) as httpd: |
| 39 | + threading.Thread(target=httpd.serve_forever).start() |
| 40 | + from google import genai |
| 41 | + api_key = "NotARealAPIKey" |
| 42 | + client = genai.Client(api_key = api_key) |
| 43 | + try: |
| 44 | + client.models.generate_content( |
| 45 | + model="gemini-2.0-flash-exp", |
| 46 | + contents="What's the largest planet in our solar system?" |
| 47 | + ) |
| 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"]) |
0 commit comments