atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_unified_rag_service.py

@@ -0,0 +1,398 @@
+"""Unit tests for UnifiedRAGService.
+
+Tests the unified RAG service that aggregates HTTP and MCP RAG sources.
+"""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from atlas.domain.unified_rag_service import UnifiedRAGService
+from atlas.modules.config.config_manager import RAGSourceConfig, RAGSourcesConfig
+from atlas.modules.rag.client import DataSource, RAGResponse
+
+
+@pytest.fixture
+def mock_config_manager():
+    """Create a mock config manager with test RAG sources."""
+    config_manager = MagicMock()
+
+    # Create test RAG sources config
+    http_source = RAGSourceConfig(
+        type="http",
+        display_name="Test HTTP RAG",
+        description="Test HTTP RAG source",
+        url="http://test-rag.example.com",
+        bearer_token="test-token",
+        groups=["users"],
+        compliance_level="Internal",
+        enabled=True,
+    )
+
+    mcp_source = RAGSourceConfig(
+        type="mcp",
+        display_name="Test MCP RAG",
+        description="Test MCP RAG source",
+        command=["python", "test_mcp.py"],
+        groups=["admin"],
+        compliance_level="SOC2",
+        enabled=True,
+    )
+
+    disabled_source = RAGSourceConfig(
+        type="http",
+        display_name="Disabled RAG",
+        url="http://disabled.example.com",
+        enabled=False,
+    )
+
+    config_manager.rag_sources_config = RAGSourcesConfig(
+        sources={
+            "test_http": http_source,
+            "test_mcp": mcp_source,
+            "disabled": disabled_source,
+        }
+    )
+
+    return config_manager
+
+
+@pytest.fixture
+def mock_auth_check():
+    """Create a mock auth check function."""
+    async def auth_check(username: str, group: str) -> bool:
+        # test@test.com is in "users" group only
+        if username == "test@test.com":
+            return group == "users"
+        # admin@test.com is in both "users" and "admin" groups
+        if username == "admin@test.com":
+            return group in ["users", "admin"]
+        return False
+
+    return auth_check
+
+
+@pytest.fixture
+def unified_rag_service(mock_config_manager, mock_auth_check):
+    """Create a UnifiedRAGService instance for testing."""
+    return UnifiedRAGService(
+        config_manager=mock_config_manager,
+        mcp_manager=None,
+        auth_check_func=mock_auth_check,
+    )
+
+
+class TestUnifiedRAGServiceInit:
+    """Tests for UnifiedRAGService initialization."""
+
+    def test_init_with_all_params(self, mock_config_manager, mock_auth_check):
+        """Test initialization with all parameters."""
+        service = UnifiedRAGService(
+            config_manager=mock_config_manager,
+            mcp_manager=MagicMock(),
+            auth_check_func=mock_auth_check,
+        )
+
+        assert service.config_manager == mock_config_manager
+        assert service.auth_check_func == mock_auth_check
+        assert service._http_clients == {}
+
+    def test_init_without_optional_params(self, mock_config_manager):
+        """Test initialization without optional parameters."""
+        service = UnifiedRAGService(config_manager=mock_config_manager)
+
+        assert service.mcp_manager is None
+        assert service.auth_check_func is None
+
+
+class TestHTTPClientCaching:
+    """Tests for HTTP client caching logic."""
+
+    def test_get_http_client_creates_new_client(self, unified_rag_service, mock_config_manager):
+        """Test that _get_http_client creates a new client when not cached."""
+        source_config = mock_config_manager.rag_sources_config.sources["test_http"]
+
+        with patch("atlas.domain.unified_rag_service.resolve_env_var", side_effect=lambda v, **kw: v):
+            client = unified_rag_service._get_http_client("test_http", source_config)
+
+        assert client is not None
+        assert "test_http" in unified_rag_service._http_clients
+        assert unified_rag_service._http_clients["test_http"] == client
+
+    def test_get_http_client_returns_cached_client(self, unified_rag_service, mock_config_manager):
+        """Test that _get_http_client returns cached client on second call."""
+        source_config = mock_config_manager.rag_sources_config.sources["test_http"]
+
+        with patch("atlas.domain.unified_rag_service.resolve_env_var", side_effect=lambda v, **kw: v):
+            client1 = unified_rag_service._get_http_client("test_http", source_config)
+            client2 = unified_rag_service._get_http_client("test_http", source_config)
+
+        assert client1 is client2
+
+
+class TestUserAuthorization:
+    """Tests for user authorization logic."""
+
+    @pytest.mark.asyncio
+    async def test_is_user_authorized_no_groups(self, unified_rag_service):
+        """Test authorization when no groups are required."""
+        result = await unified_rag_service._is_user_authorized("anyone@test.com", [])
+        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_is_user_authorized_user_in_group(self, unified_rag_service):
+        """Test authorization when user is in required group."""
+        result = await unified_rag_service._is_user_authorized("test@test.com", ["users"])
+        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_is_user_authorized_user_not_in_group(self, unified_rag_service):
+        """Test authorization when user is not in required group."""
+        result = await unified_rag_service._is_user_authorized("test@test.com", ["admin"])
+        assert result is False
+
+    @pytest.mark.asyncio
+    async def test_is_user_authorized_no_auth_func(self, mock_config_manager):
+        """Test authorization when no auth check function is provided."""
+        service = UnifiedRAGService(config_manager=mock_config_manager)
+        result = await service._is_user_authorized("anyone@test.com", ["admin"])
+        # Should return True when no auth function (permissive by default)
+        assert result is True
+
+
+class TestDiscoverDataSources:
+    """Tests for data source discovery."""
+
+    @pytest.mark.asyncio
+    async def test_discover_skips_disabled_sources(self, unified_rag_service):
+        """Test that disabled sources are skipped during discovery."""
+        with patch.object(unified_rag_service, "_discover_http_source", new_callable=AsyncMock) as mock_discover:
+            mock_discover.return_value = {"server": "test", "sources": []}
+
+            await unified_rag_service.discover_data_sources("test@test.com")
+
+            # Should not be called for disabled source
+            call_args = [call[0][0] for call in mock_discover.call_args_list]
+            assert "disabled" not in call_args
+
+    @pytest.mark.asyncio
+    async def test_discover_filters_by_authorization(self, unified_rag_service):
+        """Test that sources are filtered by user authorization."""
+        with patch.object(unified_rag_service, "_discover_http_source", new_callable=AsyncMock) as mock_discover:
+            mock_discover.return_value = {"server": "test", "sources": []}
+
+            # test@test.com is only in "users" group, not "admin"
+            await unified_rag_service.discover_data_sources("test@test.com")
+
+            # Should only discover test_http (users group), not test_mcp (admin group)
+            call_args = [call[0][0] for call in mock_discover.call_args_list]
+            assert "test_http" in call_args
+            # test_mcp requires admin group, which test@test.com doesn't have
+
+    @pytest.mark.asyncio
+    async def test_discover_includes_admin_sources_for_admin(self, unified_rag_service):
+        """Test that admin user can see admin-only sources."""
+        with patch.object(unified_rag_service, "_discover_http_source", new_callable=AsyncMock) as mock_discover:
+            mock_discover.return_value = {"server": "test", "sources": []}
+
+            # admin@test.com is in both "users" and "admin" groups
+            await unified_rag_service.discover_data_sources("admin@test.com")
+
+            # Should discover test_http (users group)
+            call_args = [call[0][0] for call in mock_discover.call_args_list]
+            assert "test_http" in call_args
+
+
+class TestDiscoverHTTPSource:
+    """Tests for HTTP source discovery."""
+
+    @pytest.mark.asyncio
+    async def test_discover_http_source_success(self, unified_rag_service, mock_config_manager):
+        """Test successful HTTP source discovery."""
+        source_config = mock_config_manager.rag_sources_config.sources["test_http"]
+
+        mock_client = AsyncMock()
+        mock_client.discover_data_sources.return_value = [
+            DataSource(name="corpus1", compliance_level="Internal"),
+            DataSource(name="corpus2", compliance_level="Public"),
+        ]
+
+        with patch.object(unified_rag_service, "_get_http_client", return_value=mock_client):
+            result = await unified_rag_service._discover_http_source(
+                "test_http", source_config, "test@test.com"
+            )
+
+        assert result is not None
+        assert result["server"] == "test_http"
+        assert result["displayName"] == "Test HTTP RAG"
+        assert len(result["sources"]) == 2
+        assert result["sources"][0]["id"] == "corpus1"
+        assert result["sources"][1]["id"] == "corpus2"
+
+    @pytest.mark.asyncio
+    async def test_discover_http_source_empty(self, unified_rag_service, mock_config_manager):
+        """Test HTTP source discovery with no data sources."""
+        source_config = mock_config_manager.rag_sources_config.sources["test_http"]
+
+        mock_client = AsyncMock()
+        mock_client.discover_data_sources.return_value = []
+
+        with patch.object(unified_rag_service, "_get_http_client", return_value=mock_client):
+            result = await unified_rag_service._discover_http_source(
+                "test_http", source_config, "test@test.com"
+            )
+
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_discover_http_source_error(self, unified_rag_service, mock_config_manager):
+        """Test HTTP source discovery handles errors gracefully."""
+        source_config = mock_config_manager.rag_sources_config.sources["test_http"]
+
+        mock_client = AsyncMock()
+        mock_client.discover_data_sources.side_effect = Exception("Connection failed")
+
+        with patch.object(unified_rag_service, "_get_http_client", return_value=mock_client):
+            result = await unified_rag_service._discover_http_source(
+                "test_http", source_config, "test@test.com"
+            )
+
+        assert result is None
+
+
+class TestQueryRAG:
+    """Tests for RAG query routing."""
+
+    @pytest.mark.asyncio
+    async def test_query_rag_with_qualified_source(self, unified_rag_service, mock_config_manager):
+        """Test querying RAG with qualified source (server:source_id)."""
+        mock_client = AsyncMock()
+        mock_client.query_rag.return_value = RAGResponse(
+            content="Test response",
+            metadata=None,
+        )
+
+        with patch.object(unified_rag_service, "_get_http_client", return_value=mock_client):
+            result = await unified_rag_service.query_rag(
+                username="test@test.com",
+                qualified_data_source="test_http:corpus1",
+                messages=[{"role": "user", "content": "test query"}],
+            )
+
+        assert result.content == "Test response"
+        mock_client.query_rag.assert_called_once_with(
+            "test@test.com",
+            "corpus1",
+            [{"role": "user", "content": "test query"}],
+        )
+
+    @pytest.mark.asyncio
+    async def test_query_rag_unknown_server(self, unified_rag_service):
+        """Test querying RAG with unknown server raises error."""
+        with pytest.raises(ValueError, match="RAG source not found"):
+            await unified_rag_service.query_rag(
+                username="test@test.com",
+                qualified_data_source="unknown_server:corpus1",
+                messages=[],
+            )
+
+    @pytest.mark.asyncio
+    async def test_query_rag_mcp_source_without_service_raises(self, unified_rag_service):
+        """Test querying MCP source without RAGMCPService raises ValueError."""
+        # The unified_rag_service fixture has no rag_mcp_service configured
+        with pytest.raises(ValueError, match="RAGMCPService not configured"):
+            await unified_rag_service.query_rag(
+                username="admin@test.com",
+                qualified_data_source="test_mcp:corpus1",
+                messages=[],
+            )
+
+    @pytest.mark.asyncio
+    async def test_query_rag_mcp_source_routes_to_mcp_service(self, mock_config_manager, mock_auth_check):
+        """Test that MCP source queries are routed to RAGMCPService."""
+        # Create a mock RAGMCPService
+        mock_rag_mcp_service = MagicMock()
+        mock_rag_mcp_service.synthesize = AsyncMock(return_value={
+            "results": {
+                "answer": "Test answer from MCP RAG",
+                "citations": [],
+            },
+            "meta_data": {
+                "providers": {
+                    "test_mcp": {"used_synth": True, "error": None}
+                },
+                "fallback_used": False,
+            },
+        })
+
+        # Create service with rag_mcp_service
+        service = UnifiedRAGService(
+            config_manager=mock_config_manager,
+            mcp_manager=None,
+            auth_check_func=mock_auth_check,
+            rag_mcp_service=mock_rag_mcp_service,
+        )
+
+        messages = [{"role": "user", "content": "What is the fleet info?"}]
+        result = await service.query_rag(
+            username="admin@test.com",
+            qualified_data_source="test_mcp:corpus1",
+            messages=messages,
+        )
+
+        # Verify RAGMCPService.synthesize was called
+        mock_rag_mcp_service.synthesize.assert_called_once_with(
+            username="admin@test.com",
+            query="What is the fleet info?",
+            sources=["test_mcp:corpus1"],
+        )
+
+        # Verify response format
+        assert isinstance(result, RAGResponse)
+        assert result.content == "Test answer from MCP RAG"
+        assert result.metadata is not None
+        assert result.metadata.data_source_name == "test_mcp"
+        assert result.metadata.retrieval_method == "mcp_synthesis"
+
+
+class TestSourceFiltering:
+    """Tests for source filtering methods."""
+
+    def test_get_http_sources(self, unified_rag_service):
+        """Test getting only HTTP sources."""
+        sources = unified_rag_service.get_http_sources()
+
+        assert "test_http" in sources
+        assert "test_mcp" not in sources
+        assert "disabled" not in sources  # Disabled sources are excluded
+
+    def test_get_mcp_sources(self, unified_rag_service):
+        """Test getting only MCP sources."""
+        sources = unified_rag_service.get_mcp_sources()
+
+        assert "test_mcp" in sources
+        assert "test_http" not in sources
+        assert "disabled" not in sources
+
+
+class TestFindServerForSource:
+    """Tests for server lookup by source ID."""
+
+    def test_find_server_returns_none(self, unified_rag_service):
+        """Test that _find_server_for_source returns None (unimplemented)."""
+        result = unified_rag_service._find_server_for_source("corpus1")
+        assert result is None
+
+
+class TestQueryRAGWithoutQualification:
+    """Tests for querying RAG without server prefix."""
+
+    @pytest.mark.asyncio
+    async def test_query_rag_without_prefix_raises(self, unified_rag_service):
+        """Test querying without server prefix raises error."""
+        with pytest.raises(ValueError, match="Could not find server"):
+            await unified_rag_service.query_rag(
+                username="test@test.com",
+                qualified_data_source="corpus1",  # No server prefix
+                messages=[],
+            )

atlas/tests/test_username_override_in_approval.py

@@ -0,0 +1,258 @@
+"""Tests for username override security in tool approval flow."""
+
+
+from unittest.mock import Mock
+
+from atlas.application.chat.utilities.tool_executor import _filter_args_to_schema, inject_context_into_args
+
+
+class TestUsernameOverrideInApproval:
+    """Test that username override cannot be bypassed through approval argument editing."""
+
+    def test_username_override_after_user_edit(self):
+        """Test that username is re-injected even after user edits it during approval."""
+        # Setup session context with authenticated user
+        session_context = {
+            "user_email": "alice@example.com",
+            "files": {}
+        }
+
+        # Simulate user editing username to a different value during approval
+        user_edited_args = {
+            "username": "malicious@example.com",  # User tried to change this
+            "data": "test data"
+        }
+
+        # Mock tool manager that indicates tool accepts username
+        mock_tool_manager = Mock()
+        mock_tool_manager.get_tools_schema.return_value = [{
+            "function": {
+                "name": "create_record",
+                "parameters": {
+                    "properties": {
+                        "username": {"type": "string"},
+                        "data": {"type": "string"}
+                    }
+                }
+            }
+        }]
+
+        # Re-inject context (simulating what should happen after user approval)
+        re_injected_args = inject_context_into_args(
+            user_edited_args,
+            session_context,
+            "create_record",
+            mock_tool_manager
+        )
+
+        # Verify username was overridden back to authenticated user
+        assert re_injected_args["username"] == "alice@example.com"
+        assert re_injected_args["data"] == "test data"
+
+        # Re-filter to schema to simulate complete flow
+        filtered_args = _filter_args_to_schema(
+            re_injected_args,
+            "create_record",
+            mock_tool_manager
+        )
+
+        # Final result should have correct username
+        assert filtered_args["username"] == "alice@example.com"
+
+    def test_username_override_with_tool_that_doesnt_accept_username(self):
+        """Test that username is not injected for tools that don't accept it."""
+        session_context = {
+            "user_email": "alice@example.com",
+            "files": {}
+        }
+
+        user_edited_args = {
+            "query": "test query"
+        }
+
+        # Mock tool manager that indicates tool does NOT accept username
+        mock_tool_manager = Mock()
+        mock_tool_manager.get_tools_schema.return_value = [{
+            "function": {
+                "name": "search",
+                "parameters": {
+                    "properties": {
+                        "query": {"type": "string"}
+                    }
+                }
+            }
+        }]
+
+        # Inject context
+        re_injected_args = inject_context_into_args(
+            user_edited_args,
+            session_context,
+            "search",
+            mock_tool_manager
+        )
+
+        # Verify username was NOT injected
+        assert "username" not in re_injected_args
+        assert re_injected_args["query"] == "test query"
+
+    def test_username_override_with_no_tool_manager(self):
+        """Test username injection when no tool manager is available (fallback)."""
+        session_context = {
+            "user_email": "bob@example.com",
+            "files": {}
+        }
+
+        user_edited_args = {
+            "data": "some data"
+        }
+
+        # Inject context with no tool manager (fallback mode)
+        re_injected_args = inject_context_into_args(
+            user_edited_args,
+            session_context,
+            "some_tool",
+            None  # No tool manager
+        )
+
+        # Should still inject username in fallback mode
+        assert re_injected_args["username"] == "bob@example.com"
+        assert re_injected_args["data"] == "some data"
+
+    def test_multiple_security_injections_after_edit(self):
+        """Test that multiple security-critical parameters are protected."""
+        session_context = {
+            "user_email": "secure_user@example.com",
+            "files": {
+                "test.pdf": {"key": "file_key_123"}
+            }
+        }
+
+        # User tries to edit both username and filename details
+        user_edited_args = {
+            "username": "hacked@example.com",  # Should be overridden
+            "filename": "test.pdf",  # Valid filename
+            "data": "edited data"
+        }
+
+        mock_tool_manager = Mock()
+        mock_tool_manager.get_tools_schema.return_value = [{
+            "function": {
+                "name": "process_file",
+                "parameters": {
+                    "properties": {
+                        "username": {"type": "string"},
+                        "filename": {"type": "string"},
+                        "data": {"type": "string"}
+                    }
+                }
+            }
+        }]
+
+        re_injected_args = inject_context_into_args(
+            user_edited_args,
+            session_context,
+            "process_file",
+            mock_tool_manager
+        )
+
+        # Username should be corrected
+        assert re_injected_args["username"] == "secure_user@example.com"
+        # File handling should work normally
+        assert "original_filename" in re_injected_args
+        assert re_injected_args["original_filename"] == "test.pdf"
+        # Data remains as user edited
+        assert re_injected_args["data"] == "edited data"
+
+    def test_prevented_impersonation_attack(self):
+        """Test specific impersonation attack scenario from vulnerability."""
+        session_context = {
+            "user_email": "alice@example.com",
+            "files": {}
+        }
+
+        # User (alice) tries to impersonate admin via approval dialog
+        user_edited_args = {
+            "username": "admin@example.com",  # Impersonation attempt
+            "action": "delete_all_data"
+        }
+
+        mock_tool_manager = Mock()
+        mock_tool_manager.get_tools_schema.return_value = [{
+            "function": {
+                "name": "admin_action",
+                "parameters": {
+                    "properties": {
+                        "username": {"type": "string"},
+                        "action": {"type": "string"}
+                    }
+                }
+            }
+        }]
+
+        # Re-inject context (the security fix)
+        re_injected_args = inject_context_into_args(
+            user_edited_args,
+            session_context,
+            "admin_action",
+            mock_tool_manager
+        )
+
+        # Re-filter for complete security
+        filtered_args = _filter_args_to_schema(
+            re_injected_args,
+            "admin_action",
+            mock_tool_manager
+        )
+
+        # Security enforced: attack prevented
+        assert filtered_args["username"] == "alice@example.com"  # Not admin
+        assert filtered_args["action"] == "delete_all_data"  # Non-security param unchanged
+
+    def test_schema_filtering_preserves_security_injection(self):
+        """Test that schema filtering works correctly with re-injected arguments."""
+        session_context = {
+            "user_email": "secure@example.com",
+            "files": {}
+        }
+
+        # User tries to add schema-violating parameters
+        user_edited_args = {
+            "username": "hacked@example.com",
+            "data": "legitimate data",
+            "extra_param": "should_be_removed"  # Not in schema
+        }
+
+        mock_tool_manager = Mock()
+        mock_tool_manager.get_tools_schema.return_value = [{
+            "function": {
+                "name": "limited_tool",
+                "parameters": {
+                    "properties": {
+                        "username": {"type": "string"},
+                        "data": {"type": "string"}
+                        # extra_param is NOT in schema
+                    }
+                }
+            }
+        }]
+
+        # Re-inject and re-filter (complete security flow)
+        re_injected_args = inject_context_into_args(
+            user_edited_args,
+            session_context,
+            "limited_tool",
+            mock_tool_manager
+        )
+
+        filtered_args = _filter_args_to_schema(
+            re_injected_args,
+            "limited_tool",
+            mock_tool_manager
+        )
+
+        # Correct username enforced
+        assert filtered_args["username"] == "secure@example.com"
+        # Legitimate data preserved
+        assert filtered_args["data"] == "legitimate data"
+        # Schema violation removed
+        assert "extra_param" not in filtered_args