Function test_app_with_gemini_enterprise

tests/unittests/cli/test_fast_api.py:930–988

Create a TestClient with gemini_enterprise_app_name set.

(
    mock_session_service,
    mock_artifact_service,
    mock_memory_service,
    mock_agent_loader,
    mock_eval_sets_manager,
    mock_eval_set_results_manager,
    monkeypatch,
)

@pytest.fixture
def test_app_with_gemini_enterprise(
    mock_session_service,
    mock_artifact_service,
    mock_memory_service,
    mock_agent_loader,
    mock_eval_sets_manager,
    mock_eval_set_results_manager,
    monkeypatch,
):
  """Create a TestClient with gemini_enterprise_app_name set."""
  monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project")
  mock_agent_loader.list_agents = MagicMock(
      return_value=["test_app", "gemini_app"]
  )

  mock_adk_app_instance = MagicMock()
  mock_adk_app_instance._tmpl_attrs = {}

  async def get_session_impl(**kwargs):
    return {"result": "success", "kwargs": kwargs}

  mock_adk_app_instance.get_session = get_session_impl

  async def stream_query_impl(**kwargs):
    yield {"chunk": 1, "kwargs": kwargs}
    await asyncio.sleep(0)
    yield {"chunk": 2, "kwargs": kwargs}

  mock_adk_app_instance.stream_query = stream_query_impl

  with (
      patch("google.auth.default", return_value=(MagicMock(), "test-project")),
      patch("vertexai.init", new_callable=MagicMock) as mock_vertexai_init,
      patch(
          "vertexai.agent_engines.AdkApp", return_value=mock_adk_app_instance
      ) as mock_adk_app_cls,
      patch("google.adk.agents.Agent", new_callable=MagicMock),
      patch(
          "google.adk.telemetry._agent_engine.TopSpanProcessor",
          new_callable=MagicMock,
      ),
      patch(
          "google.adk.telemetry._agent_engine.get_propagated_context",
          new_callable=MagicMock,
      ),
  ):
    client = _create_test_client(
        mock_session_service,
        mock_artifact_service,
        mock_memory_service,
        mock_agent_loader,
        mock_eval_sets_manager,
        mock_eval_set_results_manager,
        gemini_enterprise_app_name="gemini_app",
    )
    client.mock_vertexai_init = mock_vertexai_init
    client.mock_adk_app_cls = mock_adk_app_cls
    client.mock_adk_app_instance = mock_adk_app_instance