|
| 1 | +# Copyright 2026 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Tests for what the recordings plugin writes to the recording file.""" |
| 16 | + |
| 17 | +from google.adk.agents.callback_context import CallbackContext |
| 18 | +from google.adk.cli.plugins.recordings_plugin import RecordingsPlugin |
| 19 | +from google.adk.models.llm_request import LlmRequest |
| 20 | +from google.adk.models.llm_response import LlmResponse |
| 21 | +from google.genai import types |
| 22 | + |
| 23 | +from ... import testing_utils |
| 24 | + |
| 25 | + |
| 26 | +async def test_after_run_omits_http_options_from_the_recording_file(tmp_path): |
| 27 | + """http_options must not reach the recording file. |
| 28 | +
|
| 29 | + The plugin keeps the live LlmRequest and writes it once the run ends, so |
| 30 | + whatever the model layer put on config.http_options is still there. headers |
| 31 | + commonly holds an Authorization bearer token, and base_url and extra_body |
| 32 | + carry caller-supplied credentials too; the recording file is committed as a |
| 33 | + conformance fixture. |
| 34 | + """ |
| 35 | + plugin = RecordingsPlugin() |
| 36 | + invocation_context = await testing_utils.create_invocation_context( |
| 37 | + testing_utils.create_test_agent(name='agent_a') |
| 38 | + ) |
| 39 | + invocation_context.session.state['_adk_recordings_config'] = { |
| 40 | + 'dir': str(tmp_path), |
| 41 | + 'user_message_index': 0, |
| 42 | + 'streaming_mode': 'none', |
| 43 | + } |
| 44 | + callback_context = CallbackContext(invocation_context) |
| 45 | + llm_request = LlmRequest( |
| 46 | + model='fake-model', |
| 47 | + contents=[ |
| 48 | + types.Content(role='user', parts=[types.Part(text='roll a die')]) |
| 49 | + ], |
| 50 | + config=types.GenerateContentConfig( |
| 51 | + temperature=0.5, |
| 52 | + http_options=types.HttpOptions( |
| 53 | + headers={'Authorization': 'Bearer test-bearer-token'}, |
| 54 | + base_url='https://proxy.example/?sig=test-signature', |
| 55 | + extra_body={'api_key': 'test-extra-body-key'}, |
| 56 | + ), |
| 57 | + ), |
| 58 | + ) |
| 59 | + |
| 60 | + await plugin.before_run_callback(invocation_context=invocation_context) |
| 61 | + await plugin.before_model_callback( |
| 62 | + callback_context=callback_context, llm_request=llm_request |
| 63 | + ) |
| 64 | + await plugin.after_model_callback( |
| 65 | + callback_context=callback_context, |
| 66 | + llm_response=LlmResponse( |
| 67 | + content=types.Content( |
| 68 | + role='model', parts=[types.Part(text='rolled a 4')] |
| 69 | + ) |
| 70 | + ), |
| 71 | + ) |
| 72 | + await plugin.after_run_callback(invocation_context=invocation_context) |
| 73 | + |
| 74 | + written = (tmp_path / 'generated-recordings.yaml').read_text(encoding='utf-8') |
| 75 | + assert 'http_options' not in written |
| 76 | + assert 'test-bearer-token' not in written |
| 77 | + assert 'test-signature' not in written |
| 78 | + assert 'test-extra-body-key' not in written |
| 79 | + # after_run_callback swallows write failures, so confirm the recording was |
| 80 | + # written at all and that only http_options was dropped from the config. |
| 81 | + assert 'fake-model' in written |
| 82 | + assert 'roll a die' in written |
| 83 | + assert 'rolled a 4' in written |
| 84 | + assert 'temperature: 0.5' in written |
0 commit comments