agent-starter-pack 0.2.3__py3-none-any.whl → 0.3.0__py3-none-any.whl

src/base_template/tests/unit/test_utils/test_tracing_exporter.py

@@ -1,140 +0,0 @@
-# Copyright 2025 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from collections.abc import Generator
-from typing import Any
-from unittest.mock import Mock, patch
-
-import pytest
-from google.cloud import logging as google_cloud_logging
-from google.cloud import storage
-from opentelemetry.sdk.trace import ReadableSpan
-
-from app.utils.tracing import CloudTraceLoggingSpanExporter
-
-
-@pytest.fixture
-def mock_logging_client() -> Mock:
-    """Create a mock logging client."""
-    return Mock(spec=google_cloud_logging.Client)
-
-
-@pytest.fixture
-def mock_storage_client() -> Mock:
-    """Create a mock storage client."""
-    return Mock(spec=storage.Client)
-
-
-@pytest.fixture
-def mock_credentials() -> Any:
-    """Create mock credentials."""
-    return Mock()
-
-
-@pytest.fixture
-def patch_auth(mock_credentials: Any) -> Generator[Mock, None, None]:
-    """Patch the google.auth.default function."""
-    with patch(
-        "google.auth.default", return_value=(mock_credentials, "project")
-    ) as mock_auth:
-        yield mock_auth
-
-
-@pytest.fixture
-def patch_clients(
-    mock_logging_client: Mock, mock_storage_client: Mock
-) -> Generator[None, None, None]:
-    """Patch the logging and storage clients."""
-    with patch("google.cloud.logging.Client", return_value=mock_logging_client):
-        with patch("google.cloud.storage.Client", return_value=mock_storage_client):
-            yield
-
-
-@pytest.fixture
-def exporter(
-    mock_logging_client: Mock,
-    mock_storage_client: Mock,
-    patch_auth: Any,
-    mock_credentials: Any,
-    patch_clients: Any,
-) -> CloudTraceLoggingSpanExporter:
-    """Create a CloudTraceLoggingSpanExporter instance for testing."""
-    exporter = CloudTraceLoggingSpanExporter(
-        project_id="test-project",
-        logging_client=mock_logging_client,
-        storage_client=mock_storage_client,
-        bucket_name="test-bucket",
-    )
-    return exporter
-
-
-def test_init(exporter: CloudTraceLoggingSpanExporter) -> None:
-    """Test the initialization of CloudTraceLoggingSpanExporter."""
-    assert exporter.project_id == "test-project"
-    assert exporter.bucket_name == "test-bucket"
-    assert exporter.debug is False
-
-
-def test_store_in_gcs(exporter: CloudTraceLoggingSpanExporter) -> None:
-    """Test the store_in_gcs method of CloudTraceLoggingSpanExporter."""
-    span_id = "test-span-id"
-    content = "test-content"
-    uri = exporter.store_in_gcs(content, span_id)
-    assert uri == f"gs://test-bucket/spans/{span_id}.json"
-    exporter.bucket.blob.assert_called_once_with(f"spans/{span_id}.json")
-
-
-@patch("json.dumps")
-def test_process_large_attributes_small_payload(
-    mock_json_dumps: Mock, exporter: CloudTraceLoggingSpanExporter
-) -> None:
-    """Test processing of small payload attributes."""
-    mock_json_dumps.return_value = "a" * 100  # Small payload
-    span_dict = {"attributes": {"key": "value"}}
-    result = exporter._process_large_attributes(span_dict, "span-id")
-    assert result == span_dict
-
-
-@patch("json.dumps")
-def test_process_large_attributes_large_payload(
-    mock_json_dumps: Mock, exporter: CloudTraceLoggingSpanExporter
-) -> None:
-    """Test processing of large payload attributes."""
-    mock_json_dumps.return_value = "a" * (400 * 1024 + 1)  # Large payload
-    span_dict = {
-        "attributes": {
-            "key1": "value1",
-        }
-    }
-    result = exporter._process_large_attributes(span_dict, "span-id")
-    assert "uri_payload" in result["attributes"]
-    assert "url_payload" in result["attributes"]
-
-
-@patch.object(CloudTraceLoggingSpanExporter, "_process_large_attributes")
-def test_export(
-    mock_process_large_attributes: Mock, exporter: CloudTraceLoggingSpanExporter
-) -> None:
-    """Test the export method of CloudTraceLoggingSpanExporter."""
-    mock_span = Mock(spec=ReadableSpan)
-    mock_span.get_span_context.return_value.trace_id = 123
-    mock_span.get_span_context.return_value.span_id = 456
-    mock_span.to_json.return_value = '{"key": "value"}'
-
-    mock_process_large_attributes.return_value = {"processed": "data"}
-
-    exporter.export([mock_span])
-
-    mock_process_large_attributes.assert_called_once()
-    exporter.logger.log_struct.assert_called_once()

src/deployment_targets/cloud_run/tests/unit/test_server.py

@@ -1,124 +0,0 @@
-# Copyright 2025 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-import logging
-import os
-import sys
-from collections.abc import Generator
-from unittest.mock import MagicMock, patch
-
-import pytest
-from fastapi.testclient import TestClient
-from google.auth.credentials import Credentials
-from langchain_core.messages import HumanMessage
-
-from app.utils.typing import InputChat
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-@pytest.fixture(autouse=True)
-def mock_google_cloud_credentials() -> Generator[None, None, None]:
-    """Mock Google Cloud credentials for testing."""
-    with patch.dict(
-        os.environ,
-        {
-            "GOOGLE_APPLICATION_CREDENTIALS": "/path/to/mock/credentials.json",
-            "GOOGLE_CLOUD_PROJECT_ID": "mock-project-id",
-        },
-    ):
-        yield
-
-
-@pytest.fixture(autouse=True)
-def mock_google_auth_default() -> Generator[None, None, None]:
-    """Mock the google.auth.default function for testing."""
-    mock_credentials = MagicMock(spec=Credentials)
-    mock_project = "mock-project-id"
-
-    with patch("google.auth.default", return_value=(mock_credentials, mock_project)):
-        yield
-
-
-@pytest.fixture
-def sample_input_chat() -> InputChat:
-    """
-    Fixture to create a sample input chat for testing.
-    """
-    return InputChat(
-        messages=[HumanMessage(content="What is the meaning of life?")],
-    )
-
-
-def test_redirect_root_to_docs() -> None:
-    """
-    Test that the root endpoint (/) redirects to the Swagger UI documentation.
-    """
-    # Mock the agent module before importing server
-    mock_agent = MagicMock()
-    with patch.dict(sys.modules, {"app.agent": mock_agent}):
-        # Now import server after the mock is in place
-        from app.server import app
-
-        client = TestClient(app)
-        response = client.get("/")
-        assert response.status_code == 200
-        assert "Swagger UI" in response.text
-
-
-@pytest.mark.asyncio
-async def test_stream_chat_events() -> None:
-    """
-    Test the stream endpoint to ensure it correctly handles
-    streaming responses and generates the expected events.
-    """
-    input_data = {
-        "input": {
-            "messages": [
-                {"type": "human", "content": "Hello, AI!"},
-                {"type": "ai", "content": "Hello!"},
-                {"type": "human", "content": "What cooking recipes do you suggest?"},
-            ],
-        },
-        "config": {"metadata": {"user_id": "test-user", "session_id": "test-session"}},
-    }
-
-    mock_events = [{"content": "Mocked response"}, {"content": "Additional response"}]
-
-    # Create a mock agent module
-    mock_agent_module = MagicMock()
-    mock_agent_module.agent = MagicMock()
-    mock_agent_module.agent.stream.return_value = mock_events
-
-    # Patch the module import
-    with patch.dict(sys.modules, {"app.agent": mock_agent_module}):
-        # Import server after the mock is in place
-        from app.server import app
-
-        client = TestClient(app)
-        response = client.post("/stream_messages", json=input_data)
-
-        assert response.status_code == 200
-
-        events = []
-        for line in response.iter_lines():
-            if line:
-                events.append(json.loads(line))
-
-        assert len(events) == 2
-        assert events[0]["content"] == "Mocked response"
-        assert events[1]["content"] == "Additional response"