tracer-cloud--opensre/tests/tools/conftest.py

"""Shared test infrastructure for all tool tests.

Provides:
- BaseToolContract: a pytest mixin that any tool test can inherit from. Provides
  contract tests for metadata, is_available, and extract_params.
- Shared fixture factories: mock_agent_state, mock_boto3_client, mock_http_response.
"""

from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock

# ---------------------------------------------------------------------------
# Shared fixture factories (plain functions, called directly from tests)
# ---------------------------------------------------------------------------


def mock_agent_state(overrides: dict | None = None) -> dict[str, Any]:
    """Return a realistic AgentState dict covering common source integrations."""
    state: dict[str, Any] = {
        "cloudwatch": {
            "log_group": "/aws/lambda/my-function",
            "log_stream": "2024/01/01/[$LATEST]abc123",
            "correlation_id": "req-123",
        },
        "datadog": {
            "connection_verified": True,
            "api_key": "dd_api_key_test",
            "app_key": "dd_app_key_test",
            "site": "datadoghq.com",
            "default_query": "service:my-service",
            "time_range_minutes": 60,
            "node_ip": "10.0.1.42",
        },
        "groundcover": {
            "connection_verified": True,
            "api_key": "gc_test_token",
            "mcp_url": "https://mcp.groundcover.com/api/mcp",
            "tenant_uuid": "",
            "backend_id": "",
            "timezone": "UTC",
            "default_query": "level:error | fields _time, workload, instance, content | limit 50",
            "period": "PT1H",
        },
        "grafana": {
            "connection_verified": True,
            "grafana_endpoint": "https://grafana.example.com",
            "grafana_api_key": "glsa_test",
            "service_name": "my-service",
            "pipeline_name": "my-pipeline",
            "time_range_minutes": 60,
        },
        "eks": {
            "connection_verified": True,
            "cluster_name": "my-cluster",
            "namespace": "default",
            "pod_name": "my-pod-abc",
            "deployment": "my-deployment",
            "role_arn": "arn:aws:iam::123456789012:role/eks-role",
            "region": "us-east-1",
        },
        "sentry": {
            "connection_verified": True,
            "organization_slug": "my-org",
            "sentry_token": "sntryu_test",
            "sentry_url": "https://sentry.io",
            "project_slug": "my-project",
            "issue_id": "12345",
        },
        "github": {
            "connection_verified": True,
            "owner": "my-org",
            "repo": "my-repo",
            "path": "src/main.py",
            "github_url": "http://github.example.com/mcp",
            "github_mode": "streamable-http",
            "github_token": "ghp_test",
        },
        "opensearch": {
            "connection_verified": True,
            "url": "http://localhost:9200",
            "api_key": None,
            "default_query": "*",
            "time_range_minutes": 60,
            "index_pattern": "logs-*",
        },
        "honeycomb": {
            "connection_verified": True,
            "service_name": "my-service",
            "trace_id": "abc123trace",
            "dataset": "__all__",
            "honeycomb_api_key": "hc_test_key",
        },
        "coralogix": {
            "connection_verified": True,
            "coralogix_api_key": "cx_test_key",
            "default_query": "source logs | limit 50",
            "time_range_minutes": 60,
        },
        "s3": {
            "bucket": "my-bucket",
            "key": "my-key.json",
            "prefix": "my-prefix/",
        },
        "lambda": {
            "function_name": "my-lambda-function",
        },
        "tracer_web": {
            "trace_id": "trace-abc-123",
        },
        "prefect": {
            "connection_verified": True,
            "api_url": "http://localhost:4200/api",
            "api_key": "",
            "account_id": "",
            "workspace_id": "",
        },
        "google_docs": {
            "configured": True,
            "credentials_file": "/path/to/credentials.json",
            "folder_id": "abc123folder",
        },
        "mongodb": {
            "connection_string": "mongodb://localhost:27017",
        },
        "redis": {
            "host": "localhost",
            "port": 6379,
            "username": "",
            "password": "",
            "db": 0,
            "ssl": False,
        },
        "postgresql": {
            "host": "localhost",
            "port": 5432,
            "database": "opensre_test",
            "username": "opensre",
            "password": "test123",
            "ssl_mode": "prefer",
        },
        "mongodb_atlas": {
            "api_public_key": "test-pub-key",
            "api_private_key": "test-priv-key",
            "project_id": "test-project-id",
            "base_url": "https://cloud.mongodb.com/api/atlas/v2",
        },
        "mariadb": {
            "host": "localhost",
            "port": 3306,
            "database": "testdb",
            "username": "testuser",
            "password": "",
        },
        "mysql": {
            "host": "localhost",
            "port": 3306,
            "database": "opensre_test",
            "username": "opensre",
            "password": "test123",
            "ssl_mode": "preferred",
        },
        "splunk": {
            "connection_verified": True,
            "base_url": "https://splunk.test.corp.com:8089",
            "token": "splunk_test_bearer_token",
            "index": "main",
            "verify_ssl": False,
            "ca_bundle": "/etc/ssl/certs/corp-ca.pem",
            "default_query": 'index=main "NullPointerException" | head 50',
            "time_range_minutes": 60,
        },
        "temporal": {
            "base_url": "http://localhost:7233",
            "api_key": "",
            "namespace": "default",
        },
    }
    if overrides:
        for key, value in overrides.items():
            if isinstance(value, dict) and isinstance(state.get(key), dict):
                state[key] = {**state.get(key, {}), **value}
            else:
                state[key] = value
    return state


def mock_boto3_client(service: str, responses: dict | None = None) -> MagicMock:
    """Return a MagicMock that mimics a boto3 client with preset responses."""
    client = MagicMock()
    if responses:
        for method, return_value in responses.items():
            getattr(client, method).return_value = return_value
    return client


def mock_http_response(status: int, json_body: Any) -> MagicMock:
    """Return a MagicMock that mimics a requests.Response."""
    response = MagicMock()
    response.status_code = status
    response.ok = status < 400
    response.json.return_value = json_body
    response.text = str(json_body)
    return response


class MockHttpxResponse:
    """Typed stand-in for an httpx.Response — only the surface tests touch.

    Lets a single helper cover both response shapes the tools must handle:
    a successful JSON body returned via ``json()``, OR an HTTP error raised
    inside ``raise_for_status()`` (the tools wrap the call in ``try/except``
    and treat both as the same failure path).

    Use this from any tool test that mocks ``httpx.post``:

        def _fake_post(url, headers, json, timeout):
            return MockHttpxResponse({"data": []})

        monkeypatch.setattr("tools.MyTool.httpx.post", _fake_post)
    """

    def __init__(
        self,
        payload: Any,
        *,
        raise_for_status_error: Exception | None = None,
    ) -> None:
        self._payload = payload
        self._error = raise_for_status_error

    def raise_for_status(self) -> None:
        if self._error is not None:
            raise self._error

    def json(self) -> Any:
        return self._payload


# ---------------------------------------------------------------------------
# BaseToolContract mixin
# ---------------------------------------------------------------------------


class BaseToolContract:
    """Mixin that provides shared contract tests for any tool.

    Subclasses must implement ``get_tool_under_test()`` which returns either:
    - A BaseTool instance (class-based tools), or
    - A RegisteredTool (function-based tools, accessed via func.__opensre_registered_tool__)

    The mixin handles both patterns uniformly.

    Example for class-based tool::

        class TestMyTool(BaseToolContract):
            def get_tool_under_test(self):
                return MyTool()

    Example for function-based tool::

        from tools.MyTool import my_func

        class TestMyTool(BaseToolContract):
            def get_tool_under_test(self):
                return my_func.__opensre_registered_tool__
    """

    def get_tool_under_test(self) -> Any:
        raise NotImplementedError("Subclasses must implement get_tool_under_test()")

    def _tool(self) -> Any:
        """Return the tool under test, resolving RegisteredTool if needed."""
        return self.get_tool_under_test()

    def test_metadata_has_valid_name(self) -> None:
        tool = self._tool()
        assert isinstance(tool.name, str), "name must be a string"
        assert tool.name.strip(), "name must be non-empty"

    def test_metadata_has_valid_description(self) -> None:
        tool = self._tool()
        assert isinstance(tool.description, str), "description must be a string"
        assert tool.description.strip(), "description must be non-empty"

    def test_metadata_has_input_schema(self) -> None:
        tool = self._tool()
        schema = tool.input_schema
        assert isinstance(schema, dict), "input_schema must be a dict"
        assert "type" in schema, "input_schema must have 'type'"
        assert "properties" in schema, "input_schema must have 'properties'"

    def test_metadata_has_valid_source(self) -> None:
        tool = self._tool()
        assert isinstance(tool.source, str), "source must be a string"
        assert tool.source.strip(), "source must be non-empty"

    def test_is_available_returns_bool(self) -> None:
        tool = self._tool()
        result = tool.is_available({})
        assert isinstance(result, bool), "is_available({}) must return a bool"

    def test_extract_params_returns_dict(self) -> None:
        tool = self._tool()
        # Some tools key directly into their source dict without checking for existence.
        # Provide a full mock state so extract_params can always run without KeyError.
        sources = mock_agent_state()
        try:
            result = tool.extract_params(sources)
        except (KeyError, TypeError):
            # If extract_params requires specific keys not in the state, skip silently.
            return
        assert isinstance(result, dict), "extract_params(sources) must return a dict"