atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_domain_errors.py

@@ -0,0 +1,329 @@
+"""Tests for domain errors module."""
+
+
+from atlas.domain.errors import (
+    AuthenticationError,
+    AuthorizationError,
+    ConfigurationError,
+    DataSourcePermissionError,
+    DomainError,
+    LLMAuthenticationError,
+    LLMConfigurationError,
+    LLMError,
+    LLMServiceError,
+    LLMTimeoutError,
+    MessageError,
+    PromptOverrideError,
+    RateLimitError,
+    SessionError,
+    SessionNotFoundError,
+    ToolAuthorizationError,
+    ToolError,
+    ValidationError,
+)
+
+
+class TestDomainError:
+    """Test suite for DomainError base class."""
+
+    def test_domain_error_with_message_only(self):
+        """Test DomainError creation with message only."""
+        message = "Something went wrong"
+        error = DomainError(message)
+
+        assert str(error) == message
+        assert error.message == message
+        assert error.code is None
+
+    def test_domain_error_with_message_and_code(self):
+        """Test DomainError creation with message and code."""
+        message = "Something went wrong"
+        code = "ERR_001"
+        error = DomainError(message, code)
+
+        assert str(error) == message
+        assert error.message == message
+        assert error.code == code
+
+    def test_domain_error_inheritance(self):
+        """Test that DomainError inherits from Exception."""
+        error = DomainError("test")
+        assert isinstance(error, Exception)
+
+    def test_domain_error_with_empty_message(self):
+        """Test DomainError with empty message."""
+        error = DomainError("")
+        assert error.message == ""
+        assert str(error) == ""
+
+    def test_domain_error_with_none_code(self):
+        """Test DomainError with explicitly None code."""
+        error = DomainError("test", None)
+        assert error.code is None
+
+
+class TestValidationError:
+    """Test suite for ValidationError."""
+
+    def test_validation_error_inheritance(self):
+        """Test that ValidationError inherits from DomainError."""
+        error = ValidationError("Invalid input")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+    def test_validation_error_with_code(self):
+        """Test ValidationError with error code."""
+        error = ValidationError("Invalid email format", "VALIDATION_001")
+        assert error.message == "Invalid email format"
+        assert error.code == "VALIDATION_001"
+
+
+class TestSessionError:
+    """Test suite for SessionError."""
+
+    def test_session_error_inheritance(self):
+        """Test that SessionError inherits from DomainError."""
+        error = SessionError("Session expired")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestMessageError:
+    """Test suite for MessageError."""
+
+    def test_message_error_inheritance(self):
+        """Test that MessageError inherits from DomainError."""
+        error = MessageError("Message processing failed")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestAuthenticationError:
+    """Test suite for AuthenticationError."""
+
+    def test_authentication_error_inheritance(self):
+        """Test that AuthenticationError inherits from DomainError."""
+        error = AuthenticationError("Invalid credentials")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestAuthorizationError:
+    """Test suite for AuthorizationError."""
+
+    def test_authorization_error_inheritance(self):
+        """Test that AuthorizationError inherits from DomainError."""
+        error = AuthorizationError("Access denied")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestConfigurationError:
+    """Test suite for ConfigurationError."""
+
+    def test_configuration_error_inheritance(self):
+        """Test that ConfigurationError inherits from DomainError."""
+        error = ConfigurationError("Invalid configuration")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestLLMError:
+    """Test suite for LLMError."""
+
+    def test_llm_error_inheritance(self):
+        """Test that LLMError inherits from DomainError."""
+        error = LLMError("LLM service failed")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestLLMServiceError:
+    """Test suite for LLMServiceError."""
+
+    def test_llm_service_error_inheritance(self):
+        """Test that LLMServiceError inherits from LLMError."""
+        error = LLMServiceError("Service unavailable")
+        assert isinstance(error, LLMError)
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestToolError:
+    """Test suite for ToolError."""
+
+    def test_tool_error_inheritance(self):
+        """Test that ToolError inherits from DomainError."""
+        error = ToolError("Tool execution failed")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestToolAuthorizationError:
+    """Test suite for ToolAuthorizationError."""
+
+    def test_tool_authorization_error_inheritance(self):
+        """Test that ToolAuthorizationError inherits from AuthorizationError."""
+        error = ToolAuthorizationError("Tool access denied")
+        assert isinstance(error, AuthorizationError)
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestDataSourcePermissionError:
+    """Test suite for DataSourcePermissionError."""
+
+    def test_data_source_permission_error_inheritance(self):
+        """Test that DataSourcePermissionError inherits from AuthorizationError."""
+        error = DataSourcePermissionError("Data source access denied")
+        assert isinstance(error, AuthorizationError)
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestLLMConfigurationError:
+    """Test suite for LLMConfigurationError."""
+
+    def test_llm_configuration_error_inheritance(self):
+        """Test that LLMConfigurationError inherits from ConfigurationError."""
+        error = LLMConfigurationError("Invalid LLM config")
+        assert isinstance(error, ConfigurationError)
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestSessionNotFoundError:
+    """Test suite for SessionNotFoundError."""
+
+    def test_session_not_found_error_inheritance(self):
+        """Test that SessionNotFoundError inherits from SessionError."""
+        error = SessionNotFoundError("Session not found")
+        assert isinstance(error, SessionError)
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestPromptOverrideError:
+    """Test suite for PromptOverrideError."""
+
+    def test_prompt_override_error_inheritance(self):
+        """Test that PromptOverrideError inherits from DomainError."""
+        error = PromptOverrideError("Prompt override failed")
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestRateLimitError:
+    """Test suite for RateLimitError."""
+
+    def test_rate_limit_error_inheritance(self):
+        """Test that RateLimitError inherits from LLMError."""
+        error = RateLimitError("Rate limit exceeded")
+        assert isinstance(error, LLMError)
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestLLMTimeoutError:
+    """Test suite for LLMTimeoutError."""
+
+    def test_llm_timeout_error_inheritance(self):
+        """Test that LLMTimeoutError inherits from LLMError."""
+        error = LLMTimeoutError("Request timed out")
+        assert isinstance(error, LLMError)
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestLLMAuthenticationError:
+    """Test suite for LLMAuthenticationError."""
+
+    def test_llm_authentication_error_inheritance(self):
+        """Test that LLMAuthenticationError inherits from AuthenticationError."""
+        error = LLMAuthenticationError("LLM authentication failed")
+        assert isinstance(error, AuthenticationError)
+        assert isinstance(error, DomainError)
+        assert isinstance(error, Exception)
+
+
+class TestErrorHierarchy:
+    """Test suite for error hierarchy and relationships."""
+
+    def test_all_errors_inherit_from_domain_error(self):
+        """Test that all custom errors inherit from DomainError."""
+        error_classes = [
+            ValidationError,
+            SessionError,
+            MessageError,
+            AuthenticationError,
+            AuthorizationError,
+            ConfigurationError,
+            LLMError,
+            LLMServiceError,
+            ToolError,
+            ToolAuthorizationError,
+            DataSourcePermissionError,
+            LLMConfigurationError,
+            SessionNotFoundError,
+            PromptOverrideError,
+            RateLimitError,
+            LLMTimeoutError,
+            LLMAuthenticationError,
+        ]
+
+        for error_class in error_classes:
+            error = error_class("test message")
+            assert isinstance(error, DomainError)
+            assert isinstance(error, Exception)
+
+    def test_error_message_and_code_preservation(self):
+        """Test that all error types preserve message and code correctly."""
+        error_classes = [
+            DomainError,
+            ValidationError,
+            SessionError,
+            MessageError,
+            AuthenticationError,
+            AuthorizationError,
+            ConfigurationError,
+            LLMError,
+            LLMServiceError,
+            ToolError,
+            ToolAuthorizationError,
+            DataSourcePermissionError,
+            LLMConfigurationError,
+            SessionNotFoundError,
+            PromptOverrideError,
+            RateLimitError,
+            LLMTimeoutError,
+            LLMAuthenticationError,
+        ]
+
+        test_message = "Test error message"
+        test_code = "TEST_001"
+
+        for error_class in error_classes:
+            error = error_class(test_message, test_code)
+            assert error.message == test_message
+            assert error.code == test_code
+            assert str(error) == test_message
+
+    def test_specific_inheritance_relationships(self):
+        """Test specific inheritance relationships between error types."""
+        # Test LLM-related errors
+        assert issubclass(LLMServiceError, LLMError)
+        assert issubclass(RateLimitError, LLMError)
+        assert issubclass(LLMTimeoutError, LLMError)
+
+        # Test authorization-related errors
+        assert issubclass(ToolAuthorizationError, AuthorizationError)
+        assert issubclass(DataSourcePermissionError, AuthorizationError)
+
+        # Test authentication-related errors
+        assert issubclass(LLMAuthenticationError, AuthenticationError)
+
+        # Test configuration-related errors
+        assert issubclass(LLMConfigurationError, ConfigurationError)
+
+        # Test session-related errors
+        assert issubclass(SessionNotFoundError, SessionError)

atlas/tests/test_domain_whitelist.py

@@ -0,0 +1,359 @@
+"""Tests for domain whitelist middleware."""
+
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+from fastapi import FastAPI
+
+from atlas.core.domain_whitelist import DomainWhitelistManager
+from atlas.core.domain_whitelist_middleware import DomainWhitelistMiddleware
+
+
+@pytest.fixture
+def temp_config():
+    """Create a temporary config file for testing."""
+    config_data = {
+        "version": "1.0",
+        "description": "Test config",
+        "domains": [
+            {"domain": "sandia.gov", "description": "Sandia National Labs"},
+            {"domain": "doe.gov", "description": "DOE"},
+            {"domain": "example.org", "description": "Example"},
+        ],
+        "subdomain_matching": True
+    }
+
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
+        json.dump(config_data, f)
+        temp_path = Path(f.name)
+
+    yield temp_path
+
+    # Cleanup
+    if temp_path.exists():
+        temp_path.unlink()
+
+
+
+class TestDomainWhitelistManager:
+    """Test the domain whitelist manager."""
+
+    def test_load_config(self, temp_config):
+        """Test loading configuration from file."""
+        manager = DomainWhitelistManager(config_path=temp_config)
+
+        assert "sandia.gov" in manager.get_domains()
+        assert "doe.gov" in manager.get_domains()
+        assert "example.org" in manager.get_domains()
+        assert len(manager.get_domains()) == 3
+
+    def test_missing_config_file(self):
+        """Test that missing config file allows all domains (fail open)."""
+        non_existent_path = Path("/tmp/nonexistent_whitelist_config_12345.json")
+        manager = DomainWhitelistManager(config_path=non_existent_path)
+
+        # Config should not be loaded
+        assert manager.config_loaded is False
+        assert len(manager.get_domains()) == 0
+
+        # But should allow all domains (fail open)
+        assert manager.is_domain_allowed("user@gmail.com") is True
+        assert manager.is_domain_allowed("user@any-domain.com") is True
+        assert manager.is_domain_allowed("test@example.org") is True
+
+    def test_invalid_json_config(self):
+        """Test that invalid JSON config allows all domains (fail open)."""
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
+            f.write("{ invalid json content ]}")
+            temp_path = Path(f.name)
+
+        try:
+            manager = DomainWhitelistManager(config_path=temp_path)
+
+            # Config should not be loaded
+            assert manager.config_loaded is False
+            assert len(manager.get_domains()) == 0
+
+            # Should allow all domains (fail open)
+            assert manager.is_domain_allowed("user@gmail.com") is True
+            assert manager.is_domain_allowed("test@example.org") is True
+        finally:
+            if temp_path.exists():
+                temp_path.unlink()
+
+    def test_empty_domains_list(self):
+        """Test config with empty domains list."""
+        config_data = {
+            "version": "1.0",
+            "description": "Empty config",
+            "domains": [],
+            "subdomain_matching": True
+        }
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
+            json.dump(config_data, f)
+            temp_path = Path(f.name)
+
+        try:
+            manager = DomainWhitelistManager(config_path=temp_path)
+
+            # Config should be loaded successfully even with empty domains
+            assert manager.config_loaded is True
+            assert len(manager.get_domains()) == 0
+
+            # Should block all domains when config is valid but empty
+            assert manager.is_domain_allowed("user@gmail.com") is False
+            assert manager.is_domain_allowed("user@sandia.gov") is False
+        finally:
+            if temp_path.exists():
+                temp_path.unlink()
+
+    def test_domain_matching(self, temp_config):
+        """Test domain matching logic."""
+        manager = DomainWhitelistManager(config_path=temp_config)
+
+        # Exact matches
+        assert manager.is_domain_allowed("user@sandia.gov") is True
+        assert manager.is_domain_allowed("user@doe.gov") is True
+
+        # Subdomain matches
+        assert manager.is_domain_allowed("user@mail.sandia.gov") is True
+        assert manager.is_domain_allowed("user@sub.doe.gov") is True
+
+        # Invalid domains
+        assert manager.is_domain_allowed("user@gmail.com") is False
+        assert manager.is_domain_allowed("user@sandia.com") is False  # Wrong TLD
+
+    def test_invalid_email(self, temp_config):
+        """Test handling of invalid email addresses."""
+        manager = DomainWhitelistManager(config_path=temp_config)
+
+        assert manager.is_domain_allowed("notanemail") is False
+        assert manager.is_domain_allowed("") is False
+        assert manager.is_domain_allowed("no-at-sign.com") is False
+
+
+@pytest.fixture
+def create_middleware():
+    """Factory fixture to create middleware with custom config."""
+    from starlette.middleware.base import BaseHTTPMiddleware
+
+    def _create(config_path):
+        app = FastAPI()
+
+        # Monkey-patch to use custom config
+        original_init = DomainWhitelistMiddleware.__init__
+        def patched_init(self, app, auth_redirect_url="/auth"):
+            BaseHTTPMiddleware.__init__(self, app)
+            self.auth_redirect_url = auth_redirect_url
+            self.whitelist_manager = DomainWhitelistManager(config_path=config_path)
+
+        DomainWhitelistMiddleware.__init__ = patched_init
+        middleware = DomainWhitelistMiddleware(app)
+        DomainWhitelistMiddleware.__init__ = original_init
+
+        return middleware
+
+    return _create
+
+
+class TestDomainWhitelistMiddleware:
+    """Test domain whitelist middleware."""
+
+    def test_middleware_with_allowed_domain(self, temp_config, create_middleware):
+        """Test that allowed domains pass through."""
+        from starlette.requests import Request
+        from starlette.responses import Response
+
+        middleware = create_middleware(temp_config)
+
+        async def call_next(request):
+            return Response("OK", status_code=200)
+
+        async def test_request():
+            scope = {
+                "type": "http",
+                "method": "GET",
+                "path": "/api/test",
+                "query_string": b"",
+                "headers": [],
+                "state": {},
+            }
+            request = Request(scope)
+            request.state.user_email = "test@sandia.gov"
+
+            response = await middleware.dispatch(request, call_next)
+            assert response.status_code == 200
+
+        import asyncio
+        asyncio.run(test_request())
+
+    def test_middleware_with_disallowed_domain(self, temp_config, create_middleware):
+        """Test that disallowed domains are blocked."""
+        from starlette.requests import Request
+        from starlette.responses import Response
+
+        middleware = create_middleware(temp_config)
+
+        async def call_next(request):
+            return Response("OK", status_code=200)
+
+        async def test_request():
+            scope = {
+                "type": "http",
+                "method": "GET",
+                "path": "/api/test",
+                "query_string": b"",
+                "headers": [],
+                "state": {},
+            }
+            request = Request(scope)
+            request.state.user_email = "test@gmail.com"
+
+            response = await middleware.dispatch(request, call_next)
+            assert response.status_code == 403
+
+        import asyncio
+        asyncio.run(test_request())
+
+
+    def test_health_endpoint_bypass(self, temp_config, create_middleware):
+        """Test that health endpoint bypasses whitelist check."""
+        from starlette.requests import Request
+        from starlette.responses import Response
+
+        middleware = create_middleware(temp_config)
+
+        async def call_next(request):
+            return Response("OK", status_code=200)
+
+        async def test_request():
+            scope = {
+                "type": "http",
+                "method": "GET",
+                "path": "/api/health",
+                "query_string": b"",
+                "headers": [],
+                "state": {},
+            }
+            request = Request(scope)
+            # No email - should still pass for health check
+
+            response = await middleware.dispatch(request, call_next)
+            assert response.status_code == 200
+
+        import asyncio
+        asyncio.run(test_request())
+
+    def test_middleware_with_missing_config(self, create_middleware):
+        """Test that middleware with missing config allows all domains."""
+        from starlette.requests import Request
+        from starlette.responses import Response
+
+        non_existent_path = Path("/tmp/nonexistent_whitelist_config_12345.json")
+        middleware = create_middleware(non_existent_path)
+
+        async def call_next(request):
+            return Response("OK", status_code=200)
+
+        async def test_request():
+            scope = {
+                "type": "http",
+                "method": "GET",
+                "path": "/api/test",
+                "query_string": b"",
+                "headers": [],
+                "state": {},
+            }
+            request = Request(scope)
+            request.state.user_email = "test@gmail.com"
+
+            # Should pass even though config is missing (fail open)
+            response = await middleware.dispatch(request, call_next)
+            assert response.status_code == 200
+
+        import asyncio
+        asyncio.run(test_request())
+
+    def test_middleware_with_invalid_config(self, create_middleware):
+        """Test that middleware with invalid config allows all domains."""
+        from starlette.requests import Request
+        from starlette.responses import Response
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
+            f.write("{ invalid json }")
+            temp_path = Path(f.name)
+
+        try:
+            middleware = create_middleware(temp_path)
+
+            async def call_next(request):
+                return Response("OK", status_code=200)
+
+            async def test_request():
+                scope = {
+                    "type": "http",
+                    "method": "GET",
+                    "path": "/api/test",
+                    "query_string": b"",
+                    "headers": [],
+                    "state": {},
+                }
+                request = Request(scope)
+                request.state.user_email = "test@anydomain.com"
+
+                # Should pass even though config is invalid (fail open)
+                response = await middleware.dispatch(request, call_next)
+                assert response.status_code == 200
+
+            import asyncio
+            asyncio.run(test_request())
+        finally:
+            if temp_path.exists():
+                temp_path.unlink()
+
+    def test_middleware_with_empty_domains(self, create_middleware):
+        """Test that middleware with empty domains list blocks all."""
+        from starlette.requests import Request
+        from starlette.responses import Response
+
+        config_data = {
+            "version": "1.0",
+            "description": "Empty config",
+            "domains": [],
+            "subdomain_matching": True
+        }
+
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
+            json.dump(config_data, f)
+            temp_path = Path(f.name)
+
+        try:
+            middleware = create_middleware(temp_path)
+
+            async def call_next(request):
+                return Response("OK", status_code=200)
+
+            async def test_request():
+                scope = {
+                    "type": "http",
+                    "method": "GET",
+                    "path": "/api/test",
+                    "query_string": b"",
+                    "headers": [],
+                    "state": {},
+                }
+                request = Request(scope)
+                request.state.user_email = "test@anydomain.com"
+
+                # Should block because empty domains is a valid config
+                response = await middleware.dispatch(request, call_next)
+                assert response.status_code == 403
+
+            import asyncio
+            asyncio.run(test_request())
+        finally:
+            if temp_path.exists():
+                temp_path.unlink()