microdao-daarion/tools/observability_tool/tests/test_observability_tool.py

"""
Tests for Observability Tool
"""

import pytest
import os
import sys
from unittest.mock import AsyncMock, patch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))

from services.router.tool_manager import ToolManager, ToolResult


class TestObservabilityToolMetrics:
    """Test metrics query functionality"""
    
    @pytest.mark.asyncio
    async def test_metrics_query_valid_promql(self):
        """Test that valid PromQL queries work"""
        tool_mgr = ToolManager({})
        
        with patch.object(tool_mgr.http_client, 'get') as mock_get:
            mock_response = AsyncMock()
            mock_response.status_code = 200
            mock_response.json.return_value = {
                "status": "success",
                "data": {
                    "result": [
                        {"metric": {"service": "gateway"}, "value": [1234567890, "0.5"]}
                    ]
                }
            }
            mock_get.return_value = mock_response
            
            result = await tool_mgr._observability_tool({
                "action": "metrics_query",
                "params": {
                    "query": "rate(http_requests_total[5m])",
                    "datasource": "prometheus"
                }
            })
        
        assert result.success is True
    
    @pytest.mark.asyncio
    async def test_metrics_query_invalid_promql_blocked(self):
        """Test that invalid PromQL is blocked"""
        tool_mgr = ToolManager({})
        
        result = await tool_mgr._observability_tool({
            "action": "metrics_query",
            "params": {
                "query": "group(*) by (service)",  # Not in allowlist
                "datasource": "prometheus"
            }
        })
        
        assert result.success is False
        assert "allowed" in result.error.lower()


class TestObservabilityToolLogs:
    """Test log query functionality"""
    
    @pytest.mark.asyncio
    async def test_logs_query_time_limit(self):
        """Test that time window limit is enforced"""
        tool_mgr = ToolManager({})
        
        # Try to query with 48h window (exceeds 24h limit)
        result = await tool_mgr._observability_tool({
            "action": "logs_query",
            "params": {
                "query": "{service=\"gateway\"}",
                "time_range": {
                    "from": "2024-01-01T00:00:00Z",
                    "to": "2024-01-03T00:00:00Z"
                }
            }
        })
        
        assert result.success is False
        assert "limit" in result.error.lower()


class TestObservabilityToolTraces:
    """Test trace functionality"""
    
    @pytest.mark.asyncio
    async def test_traces_query_by_id(self):
        """Test getting trace by ID"""
        tool_mgr = ToolManager({})
        
        with patch.object(tool_mgr.http_client, 'get') as mock_get:
            mock_response = AsyncMock()
            mock_response.status_code = 200
            mock_response.json.return_value = {"batches": []}
            mock_get.return_value = mock_response
            
            result = await tool_mgr._observability_tool({
                "action": "traces_query",
                "params": {
                    "trace_id": "abc123"
                }
            })
        
        assert result.success is True


class TestObservabilityToolServiceOverview:
    """Test service overview"""
    
    @pytest.mark.asyncio
    async def test_service_overview_default_time(self):
        """Test service overview with default time range"""
        tool_mgr = ToolManager({})
        
        with patch.object(tool_mgr.http_client, 'get') as mock_get:
            mock_response = AsyncMock()
            mock_response.status_code = 200
            mock_response.json.return_value = {
                "status": "success",
                "data": {
                    "result": [
                        {"value": [1234567890, "0.1"]}
                    ]
                }
            }
            mock_get.return_value = mock_response
            
            result = await tool_mgr._observability_tool({
                "action": "service_overview",
                "params": {
                    "service": "gateway"
                }
            })
        
        assert result.success is True
        assert "metrics" in result.result


class TestObservabilityToolSecurity:
    """Test security features"""
    
    @pytest.mark.asyncio
    async def test_query_timeout(self):
        """Test that queries have timeout"""
        tool_mgr = ToolManager({})
        
        # The tool should have timeout protection
        result = await tool_mgr._observability_tool({
            "action": "metrics_query",
            "params": {
                "query": "rate(http_requests_total[5m])"
            }
        })
        
        # Should either succeed or fail gracefully
        assert result is not None


if __name__ == "__main__":
    pytest.main([__file__, "-v"])