|
18 | 18 | import google.genai |
19 | 19 | import pytest |
20 | 20 | from testing_support.fixture.event_loop import event_loop as loop |
| 21 | +from testing_support.fixture.vcr import * # noqa: F403 |
21 | 22 | from testing_support.fixtures import ( |
22 | 23 | collector_agent_registration_fixture, |
23 | 24 | collector_available_fixture, |
|
46 | 47 |
|
47 | 48 |
|
48 | 49 | @pytest.fixture |
49 | | -def gemini_client(replay_id, is_vertex): |
| 50 | +def gemini_client(vcr_recording, is_vertex): |
50 | 51 | """ |
51 | | - This configures the Gemini client to use a ReplayApiClient which will either record or replay responses depending |
52 | | - on the mode. The mode can be controlled by setting NEW_RELIC_TESTING_RECORD_GEMINI_RESPONSES=1 as an environment |
53 | | - variable to run using the real Gemini backend. (Default: mocking) |
54 | | - """ |
55 | | - from newrelic.core.config import _environ_as_bool |
| 52 | + This configures the Gemini client to use a fake API key when replaying responses through VCR. |
56 | 53 |
|
57 | | - record_mode = _environ_as_bool("NEW_RELIC_TESTING_RECORD_GEMINI_RESPONSES", False) |
58 | | - # Auto mode will record any missing files, but still replay existing files. |
59 | | - # This allows us to add new tests without having to re-record everything. |
60 | | - # If you need to re-record everything, you can delete the existing replays completely. |
61 | | - replay_client_mode = "auto" if record_mode else "replay" |
| 54 | + To record new responses, set a valid API key and run pytest with the flag --record-mode=new_episodes. |
| 55 | + """ |
62 | 56 |
|
63 | | - if record_mode: |
| 57 | + if vcr_recording: |
64 | 58 | google_api_key = os.environ.get("GOOGLE_API_KEY") |
65 | 59 | if not google_api_key: |
66 | 60 | raise RuntimeError("GOOGLE_API_KEY environment variable required.") |
67 | 61 | else: |
68 | | - google_api_key = os.environ["GOOGLE_API_KEY"] = "GEMINI_API_KEY" |
69 | | - |
70 | | - # Set the replay directory to a location in this test suite. |
71 | | - replay_dir = Path(__file__).parent / "replays" |
72 | | - replay_dir.mkdir(exist_ok=True) # Recreate this directory if it's missing for recording purposes |
73 | | - os.environ["GOOGLE_GENAI_REPLAYS_DIRECTORY"] = str(replay_dir) |
| 62 | + google_api_key = os.environ["GOOGLE_API_KEY"] = "FAKE_GEMINI_API_KEY" |
74 | 63 |
|
75 | | - # Monkeypatch the Gemini client to use the replay client which will either record or replay responses depending on the mode. |
76 | | - replay_client = google.genai._replay_api_client.ReplayApiClient(mode=replay_client_mode, replay_id=replay_id) |
77 | | - google.genai.client.Client._get_api_client = lambda self, *args, **kwargs: replay_client |
78 | 64 | gemini_client = google.genai.Client(api_key=google_api_key, vertexai=is_vertex) |
79 | 65 |
|
80 | 66 | yield gemini_client |
81 | 67 |
|
82 | | - gemini_client._api_client.close() |
83 | 68 | gemini_client.close() |
84 | 69 |
|
85 | 70 |
|
86 | | -@pytest.fixture |
87 | | -def replay_id(): |
88 | | - pytest_name = os.environ.get("PYTEST_CURRENT_TEST").split("::") |
89 | | - test_module = Path(pytest_name[0]).with_suffix("").name |
90 | | - test_name, test_params = pytest_name[-1].split(" ")[0].split("[") |
91 | | - |
92 | | - # Don't use the actual parameter string in the replay ID to avoid having different replays for certain parameters |
93 | | - # like is_async and to avoid issues with sorting. We also can't directly use the fixtures or we will cause the |
94 | | - # embedding tests to run despite not having streaming or chat parameters. Instead we will just infer the info |
95 | | - # from the test names and reproduce it in a standardized format. |
96 | | - is_streaming = "stream" in test_params |
97 | | - is_chat = "chat" in test_params |
98 | | - is_vertex = "vertex" in test_params |
99 | | - test_params_suffix = f"{'streaming' if is_streaming else 'invoke'}-{'chat' if is_chat else 'model'}-{'vertex' if is_vertex else 'standard'}" |
100 | | - |
101 | | - return f"{test_module}/{test_name}/{test_params_suffix}" |
102 | | - |
103 | | - |
104 | 71 | @pytest.fixture(scope="session", params=["standard", "vertex"]) |
105 | 72 | def is_vertex(request): |
106 | 73 | return request.param == "vertex" |
|
0 commit comments