OSVauco/opax-mcp/test_phase2b_firestore_store.py

import unittest
from unittest.mock import MagicMock, AsyncMock, patch
from datetime import datetime, timezone
from pathlib import Path
import sys

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "opax-mcp"))

from contracts.common import (
    EmmaConversation, MemoryRecord, AuditEvent, ToolRiskLevel
)
from persistence.memory_store import InMemoryMemoryStore
from persistence.firestore_store import FirestoreMemoryStore
from persistence.memory_scope import calculate_effective_memory_scope
from policy.caller_context import derive_caller_context


class TestPhase2BFirestoreStore(unittest.IsolatedAsyncioTestCase):

    def setUp(self):
        # Patch the transactional decorator to simply execute the inner function
        self.transaction_patcher = patch("persistence.firestore_store.async_transactional", new=lambda f: f)
        self.transaction_patcher.start()

        self.mock_firestore_client = MagicMock()

        # Mocks for save_conversation
        self.mock_transaction = MagicMock()
        self.mock_transaction.create = MagicMock()
        self.mock_transaction.update = MagicMock()
        self.mock_firestore_client.transaction.return_value = self.mock_transaction

        # Mocks for create_conversation
        self.mock_firestore_client.collection.return_value.document.return_value.set = AsyncMock()

        # Mocks for append_message
        self.mock_firestore_client.collection.return_value.document.return_value.collection.return_value.add = AsyncMock()
        self.mock_firestore_client.collection.return_value.document.return_value.update = AsyncMock()

        # Mocks for record_audit_event
        self.mock_firestore_client.collection.return_value.document.return_value.create = AsyncMock()

        self.store = FirestoreMemoryStore(project_id="test-project", client=self.mock_firestore_client)
    def tearDown(self):
        self.transaction_patcher.stop()

    async def test_firestore_client_is_injected(self):
        """(1) Verifies the Firestore client is used when provided."""
        self.assertIs(self.store._client, self.mock_firestore_client)

    async def test_create_conversation_uses_correct_path(self):
        """(2) Verifies conversation is stored under the correct collection/document ID."""
        convo = await self.store.create_conversation("owner1", "ws1", "creator1")
        
        self.mock_firestore_client.collection.assert_called_with("emma_conversations")
        collection_mock = self.mock_firestore_client.collection.return_value

        collection_mock.document.assert_called_with(convo.conversation_id)
        document_mock = collection_mock.document.return_value
        document_mock.set.assert_awaited_once()


    async def test_append_message_uses_subcollection_and_updates_timestamp(self):
        """(3) Verifies messages are stored as a subcollection and the parent is updated."""
        convo_id = "convo-123"
        message = {"role": "user", "content": "hello"}

        # Configure the mock chain for this specific test
        mock_convo_collection = self.mock_firestore_client.collection.return_value
        mock_doc = mock_convo_collection.document.return_value
        mock_messages_collection = mock_doc.collection.return_value

        await self.store.append_message(convo_id, message)

        # Check path to messages subcollection
        self.mock_firestore_client.collection.assert_called_with("emma_conversations")
        mock_convo_collection.document.assert_called_with(convo_id)
        mock_doc.collection.assert_called_with("messages")

        # Check that add was called on the subcollection
        mock_messages_collection.add.assert_awaited_once_with(message)

        # Check that the parent conversation's timestamp was updated
        mock_doc.update.assert_awaited_once()
        # Ensure the argument contains the last_updated_at field.
        # We don't care about the exact time, just that it's a datetime.
        update_call_args = mock_doc.update.call_args[0][0]
        self.assertIn("last_updated_at", update_call_args)
        self.assertIsInstance(update_call_args["last_updated_at"], datetime)

    async def test_save_conversation_creates_new_if_not_exists(self):
        """Tests that save_conversation creates a document if none exists."""
        convo = EmmaConversation(owner_id="owner1", workspace_id="ws1", created_by="c1")
        mock_doc_ref = self.mock_firestore_client.collection.return_value.document.return_value

        # Mock the snapshot to show the document doesn't exist
        mock_snapshot = MagicMock()
        mock_snapshot.exists = False
        mock_doc_ref.get = AsyncMock(return_value=mock_snapshot)

        await self.store.save_conversation(convo)

        # Verify that create was called on the transaction
        self.mock_transaction.create.assert_called_once_with(mock_doc_ref, convo.model_dump(mode="json"))
        self.mock_transaction.update.assert_not_called()

    async def test_save_conversation_updates_existing_with_same_owner(self):
        """Tests that save_conversation updates a document if it exists with the same owner."""
        convo = EmmaConversation(owner_id="owner1", workspace_id="ws1", created_by="c1")
        mock_doc_ref = self.mock_firestore_client.collection.return_value.document.return_value

        # Mock the snapshot to show a matching existing document
        mock_snapshot = MagicMock()
        mock_snapshot.exists = True
        mock_snapshot.to_dict.return_value = {"owner_id": "owner1", "workspace_id": "ws1"}
        mock_doc_ref.get = AsyncMock(return_value=mock_snapshot)

        await self.store.save_conversation(convo)

        self.mock_transaction.update.assert_called_once()
        self.mock_transaction.create.assert_not_called()

    async def test_save_conversation_raises_error_for_different_owner(self):
        """Tests that save_conversation fails if the owner_id mismatches."""
        convo = EmmaConversation(owner_id="owner1", workspace_id="ws1", created_by="c1")
        mock_doc_ref = self.mock_firestore_client.collection.return_value.document.return_value

        # Mock snapshot with a different owner
        mock_snapshot = MagicMock()
        mock_snapshot.exists = True
        mock_snapshot.to_dict.return_value = {"owner_id": "DIFFERENT_OWNER", "workspace_id": "ws1"}
        mock_doc_ref.get = AsyncMock(return_value=mock_snapshot)

        with self.assertRaises(PermissionError):
            await self.store.save_conversation(convo)

        self.mock_transaction.create.assert_not_called()
        self.mock_transaction.update.assert_not_called()

    async def test_save_conversation_raises_error_for_different_workspace(self):
        """Tests that save_conversation fails if the workspace_id mismatches."""
        convo = EmmaConversation(owner_id="owner1", workspace_id="ws1", created_by="c1")
        mock_doc_ref = self.mock_firestore_client.collection.return_value.document.return_value

        # Mock snapshot with a different workspace
        mock_snapshot = MagicMock()
        mock_snapshot.exists = True
        mock_snapshot.to_dict.return_value = {"owner_id": "owner1", "workspace_id": "DIFFERENT_WORKSPACE"}
        mock_doc_ref.get = AsyncMock(return_value=mock_snapshot)

        with self.assertRaises(PermissionError):
            await self.store.save_conversation(convo)

        self.mock_transaction.create.assert_not_called()
        self.mock_transaction.update.assert_not_called()

    async def test_audit_event_is_append_only(self):
        """(E) Verifies audit events use create() to be append-only."""
        event = AuditEvent(event_type="test", caller_id="c1", details={}, workspace_id="ws1")
        await self.store.record_audit_event(event)

        self.mock_firestore_client.collection.assert_called_with("emma_audit_events")
        collection_mock = self.mock_firestore_client.collection.return_value
        collection_mock.document.assert_called_with(event.event_id)
        document_mock = collection_mock.document.return_value
        document_mock.create.assert_awaited_once() # Must use create(), not set()
        document_mock.set.assert_not_called()


class TestPhase2BMemoryScope(unittest.TestCase):

    def test_effective_scope_logic(self):
        """(D) Tests the effective memory scope calculation."""
        admin_caller = derive_caller_context("user:chris.christiansen@vauco.no")
        operator_caller = derive_caller_context("agent:perplexity")

        convo = MagicMock()

        # Admin requests all, gets all allowed for admin
        requested = {"current_conversation", "owner_private_memory", "workspace_operational_memory"}
        effective = calculate_effective_memory_scope(requested, admin_caller, convo)
        self.assertEqual(effective, {"current_conversation", "owner_private_memory", "workspace_operational_memory"})

        # Operator requests all, gets only what operator is allowed
        effective = calculate_effective_memory_scope(requested, operator_caller, convo)
        self.assertEqual(effective, {"current_conversation", "workspace_operational_memory"})

        # Operator requests only conversation, gets only conversation
        requested = {"current_conversation"}
        effective = calculate_effective_memory_scope(requested, operator_caller, convo)
        self.assertEqual(effective, {"current_conversation"})

if __name__ == '__main__':
    unittest.main()