"""Unit tests for the Grafana Tempo mixin.""" from __future__ import annotations from unittest.mock import Mock, patch from integrations.grafana.tempo import TempoMixin class FakeGrafanaClient(TempoMixin): """Fake Grafana client to test the TempoMixin.""" def __init__(self, is_configured: bool = True): self.is_configured = is_configured self.account_id = "test-account-123" self.tempo_datasource_uid = "tempo-uid-abc" def _build_datasource_url(self, uid: str, path: str) -> str: return f"https://grafana.fake/api/datasources/uid/{uid}{path}" def _make_request(self, url: str, params: dict | None = None) -> dict: del url, params # To be mocked in tests return {} def _get_auth_headers(self) -> dict: return {"Authorization": "Bearer fake-token"} class TestTempoMixin: """Tests for the Tempo trace query capabilities.""" def test_query_tempo_not_configured(self): """Test behavior when the client is not configured.""" client = FakeGrafanaClient(is_configured=False) result = client.query_tempo(service_name="auth-service") assert result["success"] is False assert "not configured for account" in result["error"] assert result["traces"] == [] def test_query_tempo_general_exception(self): """Test general exception handling during a query.""" client = FakeGrafanaClient() client._make_request = Mock(side_effect=Exception("Connection timeout")) result = client.query_tempo(service_name="auth-service") assert result["success"] is False assert result["error"] == "Connection timeout" assert result["response"] == "" assert result["traces"] == [] def test_query_tempo_http_exception_with_response(self): """Test exception handling when the exception contains a response object.""" client = FakeGrafanaClient() class MockResponse: status_code = 403 text = "Permission denied for this datasource" class MockException(Exception): response = MockResponse() client._make_request = Mock(side_effect=MockException("HTTP Error")) result = client.query_tempo(service_name="auth-service") assert result["success"] is False assert result["error"] == "Tempo query failed: 403" assert "Permission denied" in result["response"] @patch("integrations.grafana.tempo.requests.get") def test_query_tempo_successful_trace_parsing(self, mock_requests_get): """Test a successful trace query and the subsequent span parsing.""" client = FakeGrafanaClient() # Mock the search response (from self._make_request) client._make_request = Mock( return_value={ "traces": [ { "traceID": "trace-123", "rootServiceName": "auth-service", "durationMs": 150, "spanCount": 2, } ] } ) # Mock the trace details response (from requests.get in _get_trace_details) mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "batches": [ { "scopeSpans": [ { "spans": [ { "name": "DB Query", "attributes": [ { "key": "db.system", "value": {"stringValue": "postgresql"}, }, {"key": "http.status_code", "value": {"intValue": 200}}, ], } ] } ] } ] } mock_requests_get.return_value = mock_response # Execute result = client.query_tempo(service_name="auth-service") # Assert Search result parsing assert result["success"] is True assert result["total_traces"] == 1 assert result["service_name"] == "auth-service" assert len(result["traces"]) == 1 # Assert enriched trace structure enriched_trace = result["traces"][0] assert enriched_trace["trace_id"] == "trace-123" assert enriched_trace["root_service"] == "auth-service" assert enriched_trace["duration_ms"] == 150 assert enriched_trace["span_count"] == 2 # Assert span parsing and attribute extraction assert len(enriched_trace["spans"]) == 1 span = enriched_trace["spans"][0] assert span["name"] == "DB Query" assert span["attributes"]["db.system"] == "postgresql" assert span["attributes"]["http.status_code"] == 200 @patch("integrations.grafana.tempo.requests.get") def test_get_trace_details_network_failure(self, mock_requests_get): """Test _get_trace_details graceful degradation on network error.""" client = FakeGrafanaClient() mock_requests_get.side_effect = Exception("Requests connection error") result = client._get_trace_details(trace_id="trace-123") # Should catch the error and return empty spans gracefully assert result == {"spans": []} def test_extract_span_attributes_edge_cases(self): """Test extraction of various attribute types.""" client = FakeGrafanaClient() mock_span = { "attributes": [ {"key": "valid_string", "value": {"stringValue": "test"}}, {"key": "valid_int", "value": {"intValue": 42}}, {"key": "valid_bool", "value": {"boolValue": True}}, {"key": "valid_double", "value": {"doubleValue": 1.5}}, {"key": "empty_value", "value": {}}, {"value": {"stringValue": "missing_key"}}, # Should be skipped! ] } attributes = client._extract_span_attributes(mock_span) assert attributes.get("valid_string") == "test" assert attributes.get("valid_int") == 42 assert attributes.get("valid_bool") is True assert attributes.get("valid_double") == 1.5 assert "empty_value" not in attributes assert "" not in attributes @patch("integrations.grafana.tempo.requests.get") def test_get_trace_details_non_200_status(self, mock_requests_get): """Test _get_trace_details when the API returns a non-200 status.""" client = FakeGrafanaClient() # Setup mock to return a 404 status code mock_response = Mock() mock_response.status_code = 404 mock_requests_get.return_value = mock_response result = client._get_trace_details(trace_id="trace-123") # Assert it safely falls back to empty spans assert result == {"spans": []}