diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index 3301ca4431..16939c7a2d 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/src/google/adk/cli/cli_deploy.py @@ -63,7 +63,7 @@ def _ensure_agent_engine_dependency(requirements_txt_path: str) -> None: with open(requirements_txt_path, 'a', encoding='utf-8') as f: if requirements and not requirements.endswith('\n'): f.write('\n') - f.write('google-cloud-aiplatform[agent_engines]\n') + f.write(f'{_AGENT_ENGINE_REQUIREMENT}\n') f.write(f'google-adk=={__version__}\n') @@ -1017,7 +1017,7 @@ def to_agent_engine( if not os.path.exists(requirements_txt_path): click.echo(f'Creating {requirements_txt_path}...') with open(requirements_txt_path, 'w', encoding='utf-8') as f: - f.write('google-cloud-aiplatform[agent_engines]\n') + f.write(f'{_AGENT_ENGINE_REQUIREMENT}\n') f.write(f'google-adk=={__version__}\n') click.echo(f'Using google-adk=={__version__} in requirements') click.echo(f'Created {requirements_txt_path}') diff --git a/tests/unittests/cli/utils/test_cli_deploy.py b/tests/unittests/cli/utils/test_cli_deploy.py index 45c296cc85..650da437d7 100644 --- a/tests/unittests/cli/utils/test_cli_deploy.py +++ b/tests/unittests/cli/utils/test_cli_deploy.py @@ -303,6 +303,13 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: assert len(create_recorder.calls) == 1 assert str(rmtree_recorder.get_last_call_args()[0]) == str(tmp_dir) + requirements_file = tmp_dir / "agents" / "agent" / "requirements.txt" + assert requirements_file.is_file() + assert ( + "google-cloud-aiplatform[adk,agent_engines]" + in requirements_file.read_text() + ) + def test_to_agent_engine_raises_when_explicit_config_file_missing( monkeypatch: pytest.MonkeyPatch, @@ -685,3 +692,25 @@ def test_cli_deploy_agent_engine_artifact_service_uri(tmp_path: Path): mock_to_agent_engine.assert_called_once() _, kwargs = mock_to_agent_engine.call_args assert kwargs["artifact_service_uri"] == "gs://my-bucket" + + +def test_ensure_agent_engine_dependency(tmp_path: Path): + """Tests that _ensure_agent_engine_dependency appends correct extras.""" + requirements_file = tmp_path / "requirements.txt" + + # Case 1: raises FileNotFoundError when the file doesn't exist + with pytest.raises(FileNotFoundError): + cli_deploy._ensure_agent_engine_dependency(str(requirements_file)) + + # Case 2: appends both google-cloud-aiplatform with 'adk' and 'agent_engines' extras and versioned google-adk + requirements_file.write_text("") + cli_deploy._ensure_agent_engine_dependency(str(requirements_file)) + content = requirements_file.read_text() + assert "google-cloud-aiplatform[adk,agent_engines]\n" in content + assert f"google-adk=={cli_deploy.__version__}\n" in content + + # Case 3: does not append duplicate if google-cloud-aiplatform already exists + requirements_file.write_text("google-cloud-aiplatform[adk,agent_engines]\n") + cli_deploy._ensure_agent_engine_dependency(str(requirements_file)) + content = requirements_file.read_text() + assert content == "google-cloud-aiplatform[adk,agent_engines]\n"