atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_atlas_rag_integration.py

@@ -0,0 +1,224 @@
+"""Integration tests for AtlasRAGClient with the mock service.
+
+These tests require the atlas-rag-api-mock service to be running.
+They can be skipped if the mock service is not available.
+"""
+
+import os
+import subprocess
+import sys
+import time
+
+import httpx
+import pytest
+
+# Add paths for imports
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from atlas.modules.rag.atlas_rag_client import AtlasRAGClient
+
+MOCK_URL = "http://localhost:8002"
+MOCK_TOKEN = "test-atlas-rag-token"
+MOCK_STARTUP_TIMEOUT = 10
+
+
+def is_mock_running() -> bool:
+    """Check if the mock service is running."""
+    try:
+        response = httpx.get(f"{MOCK_URL}/health", timeout=2.0)
+        return response.status_code == 200
+    except Exception:
+        return False
+
+
+@pytest.fixture(scope="module")
+def mock_service():
+    """Start the mock service if not already running."""
+    if is_mock_running():
+        yield MOCK_URL
+        return
+
+    # Try to start the mock service
+    mock_path = os.path.join(
+        os.path.dirname(__file__), "..", "..", "mocks", "atlas-rag-api-mock", "main.py"
+    )
+    mock_path = os.path.abspath(mock_path)
+
+    if not os.path.exists(mock_path):
+        pytest.skip(f"Mock service not found at {mock_path}")
+
+    process = subprocess.Popen(
+        [sys.executable, mock_path],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+
+    # Wait for service to start
+    start_time = time.time()
+    while time.time() - start_time < MOCK_STARTUP_TIMEOUT:
+        if is_mock_running():
+            break
+        time.sleep(0.5)
+    else:
+        process.terminate()
+        pytest.skip("Could not start mock service")
+
+    yield MOCK_URL
+
+    # Cleanup
+    process.terminate()
+    process.wait(timeout=5)
+
+
+@pytest.fixture
+def client():
+    """Create an AtlasRAGClient configured for the mock service."""
+    return AtlasRAGClient(
+        base_url=MOCK_URL,
+        bearer_token=MOCK_TOKEN,
+        default_model="test-model",
+        top_k=4,
+    )
+
+
+class TestAtlasRAGIntegration:
+    """Integration tests for AtlasRAGClient with the mock service."""
+
+    @pytest.mark.asyncio
+    async def test_discover_data_sources_success(self, mock_service, client):
+        """Test discovering data sources for a known user."""
+        sources = await client.discover_data_sources("test@test.com")
+
+        assert len(sources) > 0
+        # test@test.com has employee, engineering, devops, admin groups - sees all sources
+        source_names = [s.name for s in sources]
+        assert "company-policies" in source_names
+        assert "technical-docs" in source_names
+        assert "product-knowledge" in source_names  # Public
+
+        # Check compliance levels are returned
+        for source in sources:
+            assert source.compliance_level in ["Internal", "Public"]
+
+    @pytest.mark.asyncio
+    async def test_discover_data_sources_unknown_user(self, mock_service, client):
+        """Test discovering data sources for an unknown user returns only public sources."""
+        sources = await client.discover_data_sources("unknown@example.com")
+        # Unknown users get public sources only
+        source_names = [s.name for s in sources]
+        assert "product-knowledge" in source_names
+        assert "company-policies" not in source_names
+        assert "technical-docs" not in source_names
+
+    @pytest.mark.asyncio
+    async def test_discover_data_sources_limited_access(self, mock_service, client):
+        """Test that users only see corpora they have access to."""
+        # bob@example.com has employee and sales groups
+        sources = await client.discover_data_sources("bob@example.com")
+
+        source_names = [s.name for s in sources]
+        assert "company-policies" in source_names  # requires employee
+        assert "product-knowledge" in source_names  # Public
+        # Should NOT have access to technical-docs (requires engineering or devops)
+        assert "technical-docs" not in source_names
+
+    @pytest.mark.asyncio
+    async def test_query_rag_success(self, mock_service, client):
+        """Test successful RAG query."""
+        messages = [{"role": "user", "content": "What is the API authentication?"}]
+
+        response = await client.query_rag(
+            user_name="test@test.com",
+            data_source="technical-docs",
+            messages=messages,
+        )
+
+        assert response.content is not None
+        assert len(response.content) > 0
+        # Should contain information about API or authentication
+        assert "API" in response.content or "authentication" in response.content.lower()
+
+        # Check metadata
+        assert response.metadata is not None
+        assert response.metadata.query_processing_time_ms >= 0
+        assert response.metadata.data_source_name == "technical-docs"
+        assert response.metadata.retrieval_method == "keyword-search"
+        assert len(response.metadata.documents_found) > 0
+
+    @pytest.mark.asyncio
+    async def test_query_rag_with_metadata(self, mock_service, client):
+        """Test that RAG query returns document metadata."""
+        messages = [{"role": "user", "content": "Tell me about deployment pipeline"}]
+
+        response = await client.query_rag(
+            user_name="charlie@example.com",  # Has employee, engineering, devops groups
+            data_source="technical-docs",
+            messages=messages,
+        )
+
+        assert response.metadata is not None
+        assert len(response.metadata.documents_found) > 0
+
+        # Check document metadata structure
+        doc = response.metadata.documents_found[0]
+        assert doc.source is not None
+        assert doc.confidence_score > 0
+        assert doc.content_type is not None
+
+    @pytest.mark.asyncio
+    async def test_query_rag_access_denied(self, mock_service, client):
+        """Test that RAG query returns 403 for unauthorized access."""
+        messages = [{"role": "user", "content": "Show me technical docs"}]
+
+        with pytest.raises(Exception) as exc_info:
+            await client.query_rag(
+                user_name="bob@example.com",  # Has employee, sales - no engineering/devops
+                data_source="technical-docs",  # Requires engineering or devops
+                messages=messages,
+            )
+
+        # Should raise HTTPException with 403
+        assert "403" in str(exc_info.value) or "access" in str(exc_info.value).lower()
+
+    @pytest.mark.asyncio
+    async def test_query_rag_corpus_not_found(self, mock_service, client):
+        """Test that RAG query returns 404 for non-existent corpus."""
+        messages = [{"role": "user", "content": "Search something"}]
+
+        with pytest.raises(Exception) as exc_info:
+            await client.query_rag(
+                user_name="test@test.com",
+                data_source="non-existent-corpus",
+                messages=messages,
+            )
+
+        # Should raise HTTPException with 404
+        assert "404" in str(exc_info.value) or "not found" in str(exc_info.value).lower()
+
+
+class TestAtlasRAGAuthFailures:
+    """Test authentication failure scenarios."""
+
+    @pytest.mark.asyncio
+    async def test_missing_token(self, mock_service):
+        """Test that requests without token fail with 401."""
+        client = AtlasRAGClient(
+            base_url=MOCK_URL,
+            bearer_token=None,  # No token
+        )
+
+        # Discovery should return empty list on auth failure (graceful degradation)
+        sources = await client.discover_data_sources("test@test.com")
+        assert sources == []
+
+    @pytest.mark.asyncio
+    async def test_invalid_token(self, mock_service):
+        """Test that requests with invalid token fail with 401."""
+        client = AtlasRAGClient(
+            base_url=MOCK_URL,
+            bearer_token="invalid-token",
+        )
+
+        # Discovery should return empty list on auth failure (graceful degradation)
+        sources = await client.discover_data_sources("test@test.com")
+        assert sources == []

atlas/tests/test_attach_file_flow.py

@@ -0,0 +1,287 @@
+import base64
+import uuid
+
+import pytest
+
+from atlas.application.chat.service import ChatService
+from atlas.modules.file_storage.manager import FileManager
+from atlas.modules.file_storage.mock_s3_client import MockS3StorageClient
+
+
+class FakeLLM:
+    async def call_plain(self, model_name, messages, temperature=0.7):
+        return "ok"
+
+    async def call_with_tools(self, model_name, messages, tools_schema, tool_choice="auto", temperature=0.7):
+        from atlas.interfaces.llm import LLMResponse
+        return LLMResponse(content="ok", tool_calls=None, model_used=model_name)
+
+    async def call_with_rag(self, model_name, messages, data_sources, user_email, temperature=0.7):
+        return "ok"
+
+    async def call_with_rag_and_tools(self, model_name, messages, data_sources, tools_schema, user_email, tool_choice="auto", temperature=0.7):
+        from atlas.interfaces.llm import LLMResponse
+        return LLMResponse(content="ok", tool_calls=None, model_used=model_name)
+
+
+@pytest.fixture
+def file_manager():
+    # Use in-process mock S3 for deterministic tests
+    return FileManager(s3_client=MockS3StorageClient())
+
+
+@pytest.fixture
+def chat_service(file_manager):
+    # Minimal ChatService wiring for file/session operations
+    return ChatService(llm=FakeLLM(), file_manager=file_manager)
+
+
+@pytest.mark.asyncio
+async def test_handle_attach_file_success_creates_session_and_emits_update(chat_service, file_manager):
+    user_email = "user1@example.com"
+    session_id = uuid.uuid4()
+
+    # Seed a file into the mock storage for this user
+    filename = "report.txt"
+    content_b64 = base64.b64encode(b"hello world").decode()
+    upload_meta = await file_manager.s3_client.upload_file(
+        user_email=user_email,
+        filename=filename,
+        content_base64=content_b64,
+        content_type="text/plain",
+        tags={"source": "user"},
+        source_type="user",
+    )
+    s3_key = upload_meta["key"]
+
+    updates = []
+
+    async def capture_update(msg):
+        updates.append(msg)
+
+    # Act: attach the file to a brand new session (auto-creates session)
+    resp = await chat_service.handle_attach_file(
+        session_id=session_id,
+        s3_key=s3_key,
+        user_email=user_email,
+        update_callback=capture_update,
+    )
+
+    # Assert: success response and files_update emitted
+    assert resp.get("type") == "file_attach"
+    assert resp.get("success") is True
+    assert resp.get("filename") == filename
+
+    assert any(
+        u.get("type") == "intermediate_update" and u.get("update_type") == "files_update"
+        for u in updates
+    ), "Expected a files_update intermediate update to be emitted"
+
+    # Session context should include the file by filename
+    session = chat_service.sessions.get(session_id)
+    assert session is not None
+    assert filename in session.context.get("files", {})
+    assert session.context["files"][filename]["key"] == s3_key
+
+
+@pytest.mark.asyncio
+async def test_handle_attach_file_not_found_returns_error(chat_service):
+    user_email = "user1@example.com"
+    session_id = uuid.uuid4()
+
+    # Non-existent S3 key for the same user
+    bad_key = f"users/{user_email}/uploads/does_not_exist_12345.txt"
+    resp = await chat_service.handle_attach_file(
+        session_id=session_id,
+        s3_key=bad_key,
+        user_email=user_email,
+        update_callback=None,
+    )
+
+    assert resp.get("type") == "file_attach"
+    assert resp.get("success") is False
+    assert "File not found" in resp.get("error", "")
+
+
+@pytest.mark.asyncio
+async def test_handle_attach_file_unauthorized_other_user_key(chat_service, file_manager):
+    # Upload under user1
+    owner_email = "owner@example.com"
+    other_email = "other@example.com"
+    session_id = uuid.uuid4()
+
+    filename = "secret.pdf"
+    content_b64 = base64.b64encode(b"top-secret").decode()
+    upload_meta = await file_manager.s3_client.upload_file(
+        user_email=owner_email,
+        filename=filename,
+        content_base64=content_b64,
+        content_type="application/pdf",
+        tags={"source": "user"},
+        source_type="user",
+    )
+    s3_key = upload_meta["key"]
+
+    # Attempt to attach with a different user should fail
+    resp = await chat_service.handle_attach_file(
+        session_id=session_id,
+        s3_key=s3_key,
+        user_email=other_email,
+        update_callback=None,
+    )
+
+    assert resp.get("type") == "file_attach"
+    assert resp.get("success") is False
+    assert "Access denied" in resp.get("error", "")
+
+
+@pytest.mark.asyncio
+async def test_handle_reset_session_reinitializes(chat_service):
+    user_email = "user1@example.com"
+    session_id = uuid.uuid4()
+
+    # Create a session first
+    await chat_service.create_session(session_id, user_email)
+    assert chat_service.sessions.get(session_id) is not None
+
+    # Reset the session
+    resp = await chat_service.handle_reset_session(session_id=session_id, user_email=user_email)
+
+    assert resp.get("type") == "session_reset"
+    # After reset, a fresh active session should exist for the same id
+    new_session = chat_service.sessions.get(session_id)
+    assert new_session is not None
+    assert new_session.active is True
+
+
+@pytest.mark.asyncio
+async def test_handle_download_file_success_after_attach(chat_service, file_manager):
+    user_email = "user1@example.com"
+    session_id = uuid.uuid4()
+
+    # Upload and then attach to session
+    filename = "notes.md"
+    content_bytes = b"### Title\nSome content."
+    content_b64 = base64.b64encode(content_bytes).decode()
+    upload_meta = await file_manager.s3_client.upload_file(
+        user_email=user_email,
+        filename=filename,
+        content_base64=content_b64,
+        content_type="text/markdown",
+        tags={"source": "user"},
+        source_type="user",
+    )
+    s3_key = upload_meta["key"]
+
+    await chat_service.handle_attach_file(
+        session_id=session_id,
+        s3_key=s3_key,
+        user_email=user_email,
+        update_callback=None,
+    )
+
+    # Act: download by filename (from session context)
+    resp = await chat_service.handle_download_file(
+        session_id=session_id,
+        filename=filename,
+        user_email=user_email,
+    )
+
+    assert resp.get("type") is not None
+    # content_base64 should match uploaded content
+    returned_b64 = resp.get("content_base64")
+    assert isinstance(returned_b64, str) and len(returned_b64) > 0
+    assert base64.b64decode(returned_b64) == content_bytes
+
+
+@pytest.mark.asyncio
+async def test_handle_download_file_not_in_session_returns_error(chat_service):
+    user_email = "user1@example.com"
+    session_id = uuid.uuid4()
+    filename = "missing.txt"
+
+    # No attach performed; should error that file isn't in session
+    resp = await chat_service.handle_download_file(
+        session_id=session_id,
+        filename=filename,
+        user_email=user_email,
+    )
+
+    assert resp.get("error") == "Session or file manager not available" or resp.get("error") == "File not found in session"
+
+
+@pytest.mark.asyncio
+async def test_upload_file_with_spaces_in_filename(file_manager):
+    """Files with spaces in their names should upload successfully after sanitization."""
+    user_email = "user1@example.com"
+    filename_with_spaces = "my report file.txt"
+    content_b64 = base64.b64encode(b"some content").decode()
+
+    result = await file_manager.upload_file(
+        user_email=user_email,
+        filename=filename_with_spaces,
+        content_base64=content_b64,
+        source_type="user",
+    )
+
+    # Filename should be sanitized (spaces replaced with underscores)
+    assert result["filename"] == "my_report_file.txt"
+    assert "my_report_file.txt" in result["key"]
+    assert " " not in result["key"]
+
+
+@pytest.mark.asyncio
+async def test_upload_multiple_files_with_spaces(file_manager):
+    """upload_multiple_files should sanitize filenames containing spaces."""
+    user_email = "user1@example.com"
+    files = {
+        "my document.pdf": base64.b64encode(b"pdf bytes").decode(),
+        "another file.txt": base64.b64encode(b"text bytes").decode(),
+    }
+
+    uploaded = await file_manager.upload_multiple_files(
+        user_email=user_email,
+        files=files,
+        source_type="user",
+    )
+
+    assert "my_document.pdf" in uploaded
+    assert "another_file.txt" in uploaded
+    for key in uploaded.values():
+        assert " " not in key
+
+
+@pytest.mark.asyncio
+async def test_attach_file_with_spaces_end_to_end(chat_service, file_manager):
+    """Full flow: upload a file with spaces, attach it, verify sanitized name in session."""
+    user_email = "user1@example.com"
+    session_id = uuid.uuid4()
+    filename_with_spaces = "test report.txt"
+    content_b64 = base64.b64encode(b"hello spaces").decode()
+
+    upload_meta = await file_manager.upload_file(
+        user_email=user_email,
+        filename=filename_with_spaces,
+        content_base64=content_b64,
+        source_type="user",
+    )
+    s3_key = upload_meta["key"]
+
+    resp = await chat_service.handle_attach_file(
+        session_id=session_id,
+        s3_key=s3_key,
+        user_email=user_email,
+        update_callback=None,
+    )
+
+    assert resp.get("success") is True
+    assert " " not in s3_key
+
+    # Verify the session stores the sanitized filename (no spaces)
+    session = chat_service.sessions.get(session_id)
+    assert session is not None
+    session_files = session.context.get("files", {})
+    # The filename key in the session should have underscores, not spaces
+    stored_names = list(session_files.keys())
+    for name in stored_names:
+        assert " " not in name, f"Session stored filename with spaces: {name}"

atlas/tests/test_auth_utils.py

@@ -0,0 +1,165 @@
+"""Tests for auth_utils module."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from atlas.core.authorization_manager import AuthorizationManager, create_authorization_manager
+
+
+class TestAuthorizationManager:
+    """Test suite for AuthorizationManager class."""
+
+    @pytest.fixture
+    def mock_auth_check_func(self):
+        """Create a mock auth check function."""
+        return AsyncMock()
+
+    @pytest.fixture
+    def mock_app_settings(self):
+        """Create mock app settings."""
+        settings = MagicMock()
+        settings.admin_group = "admin"
+        return settings
+
+    @pytest.fixture
+    def auth_manager(self, mock_auth_check_func, mock_app_settings):
+        """Create an AuthorizationManager instance with mocked dependencies."""
+        with patch('atlas.core.authorization_manager.get_app_settings', return_value=mock_app_settings):
+            return AuthorizationManager(mock_auth_check_func)
+
+    @pytest.mark.asyncio
+    async def test_is_admin_returns_true_for_admin_user(self, auth_manager, mock_auth_check_func):
+        """Test that is_admin returns True when auth check function returns True."""
+        # Arrange
+        user_email = "admin@example.com"
+        mock_auth_check_func.return_value = True
+
+        # Act
+        result = await auth_manager.is_admin(user_email)
+
+        # Assert
+        assert result is True
+        mock_auth_check_func.assert_called_once_with(user_email, "admin")
+
+    @pytest.mark.asyncio
+    async def test_is_admin_returns_false_for_non_admin_user(self, auth_manager, mock_auth_check_func):
+        """Test that is_admin returns False when auth check function returns False."""
+        # Arrange
+        user_email = "user@example.com"
+        mock_auth_check_func.return_value = False
+
+        # Act
+        result = await auth_manager.is_admin(user_email)
+
+        # Assert
+        assert result is False
+        mock_auth_check_func.assert_called_once_with(user_email, "admin")
+
+    @pytest.mark.asyncio
+    async def test_is_admin_uses_correct_admin_group(self, mock_auth_check_func):
+        """Test that is_admin uses the admin group from app settings."""
+        # Arrange
+        custom_admin_group = "super_admin"
+        mock_settings = MagicMock()
+        mock_settings.admin_group = custom_admin_group
+
+        with patch('atlas.core.authorization_manager.get_app_settings', return_value=mock_settings):
+            auth_manager = AuthorizationManager(mock_auth_check_func)
+
+        user_email = "admin@example.com"
+        mock_auth_check_func.return_value = True
+
+        # Act
+        await auth_manager.is_admin(user_email)
+
+        # Assert
+        mock_auth_check_func.assert_called_once_with(user_email, custom_admin_group)
+
+    @pytest.mark.asyncio
+    async def test_is_admin_handles_auth_check_exception(self, auth_manager, mock_auth_check_func):
+        """Test that is_admin properly propagates exceptions from auth check function."""
+        # Arrange
+        user_email = "user@example.com"
+        mock_auth_check_func.side_effect = Exception("Auth service unavailable")
+
+        # Act & Assert
+        with pytest.raises(Exception, match="Auth service unavailable"):
+            await auth_manager.is_admin(user_email)
+
+    @pytest.mark.asyncio
+    async def test_is_admin_with_empty_email(self, auth_manager, mock_auth_check_func):
+        """Test that is_admin handles empty email addresses."""
+        # Arrange
+        user_email = ""
+        mock_auth_check_func.return_value = False
+
+        # Act
+        result = await auth_manager.is_admin(user_email)
+
+        # Assert
+        assert result is False
+        mock_auth_check_func.assert_called_once_with("", "admin")
+
+    @pytest.mark.asyncio
+    async def test_is_admin_with_none_email(self, auth_manager, mock_auth_check_func):
+        """Test that is_admin handles None email addresses."""
+        # Arrange
+        user_email = None
+        mock_auth_check_func.return_value = False
+
+        # Act
+        result = await auth_manager.is_admin(user_email)
+
+        # Assert
+        assert result is False
+        mock_auth_check_func.assert_called_once_with(None, "admin")
+
+
+class TestCreateAuthorizationManager:
+    """Test suite for create_authorization_manager factory function."""
+
+    def test_create_authorization_manager_returns_instance(self):
+        """Test that factory function returns an AuthorizationManager instance."""
+        # Arrange
+        mock_auth_check_func = AsyncMock()
+
+        # Act
+        with patch('atlas.core.authorization_manager.get_app_settings'):
+            result = create_authorization_manager(mock_auth_check_func)
+
+        # Assert
+        assert isinstance(result, AuthorizationManager)
+        assert result.auth_check_func is mock_auth_check_func
+
+    def test_create_authorization_manager_initializes_app_settings(self):
+        """Test that factory function properly initializes app settings."""
+        # Arrange
+        mock_auth_check_func = AsyncMock()
+        mock_settings = MagicMock()
+        mock_settings.admin_group = "test_admin"
+
+        # Act
+        with patch('atlas.core.authorization_manager.get_app_settings', return_value=mock_settings) as mock_get_settings:
+            result = create_authorization_manager(mock_auth_check_func)
+
+        # Assert
+        mock_get_settings.assert_called_once()
+        assert result.app_settings is mock_settings
+
+    @pytest.mark.asyncio
+    async def test_created_manager_works_correctly(self):
+        """Test that the manager created by factory function works correctly."""
+        # Arrange
+        mock_auth_check_func = AsyncMock(return_value=True)
+        mock_settings = MagicMock()
+        mock_settings.admin_group = "admin"
+
+        # Act
+        with patch('atlas.core.authorization_manager.get_app_settings', return_value=mock_settings):
+            manager = create_authorization_manager(mock_auth_check_func)
+            result = await manager.is_admin("test@example.com")
+
+        # Assert
+        assert result is True
+        mock_auth_check_func.assert_called_once_with("test@example.com", "admin")