python-sdk/pkg/hanzo-memory/tests/test_infinity_client.py

"""Test InfinityDB client."""

import uuid


class TestInfinityClient:
    """Test InfinityDB client."""

    def test_client_initialization(self, db_client):
        """Test client initialization."""
        assert db_client is not None
        assert hasattr(db_client, "infinity")

    def test_create_project(self, db_client, sample_user_id):
        """Test project creation."""
        # Ensure projects table exists
        db_client.create_projects_table()

        project_id = f"test_project_{uuid.uuid4().hex[:8]}"

        project = db_client.create_project(
            project_id=project_id,
            user_id=sample_user_id,
            name="Test Project",
            description="A test project",
            metadata={"test": True},
        )

        assert project["project_id"] == project_id
        assert project["user_id"] == sample_user_id
        assert project["name"] == "Test Project"
        assert "created_at" in project

    def test_memory_operations(self, db_client, sample_user_id, sample_project_id):
        """Test memory CRUD operations."""
        # Create memories table
        db_client.create_memories_table(sample_user_id)

        # Add a memory
        memory_id = f"mem_{uuid.uuid4().hex[:8]}"
        embedding = [0.1] * 384  # Mock embedding

        memory = db_client.add_memory(
            memory_id=memory_id,
            user_id=sample_user_id,
            project_id=sample_project_id,
            content="Test memory content",
            embedding=embedding,
            metadata={"source": "test"},
            importance=7.5,
        )

        assert memory["memory_id"] == memory_id
        assert memory["content"] == "Test memory content"
        assert memory["importance"] == 7.5

    def test_memory_search(self, db_client, sample_user_id, sample_project_id):
        """Test memory search functionality."""
        # Create memories table
        db_client.create_memories_table(sample_user_id)

        # Add some memories
        memories_data = [
            ("Python programming is fun", [0.8, 0.2] + [0.1] * 382),
            ("JavaScript is also great", [0.2, 0.8] + [0.1] * 382),
            ("I love coding", [0.5, 0.5] + [0.1] * 382),
        ]

        for i, (content, embedding) in enumerate(memories_data):
            db_client.add_memory(
                memory_id=f"mem_{i}",
                user_id=sample_user_id,
                project_id=sample_project_id,
                content=content,
                embedding=embedding,
            )

        # Search for Python-related memories
        query_embedding = [0.9, 0.1] + [0.1] * 382
        results = db_client.search_memories(
            user_id=sample_user_id,
            query_embedding=query_embedding,
            limit=2,
        )

        assert not results.is_empty()
        assert len(results) <= 2

    def test_knowledge_base_operations(
        self, db_client, sample_user_id, sample_project_id
    ):
        """Test knowledge base operations."""
        # Ensure knowledge bases table exists
        db_client.create_knowledge_bases_table()

        kb_id = f"kb_{uuid.uuid4().hex[:8]}"

        # Create knowledge base
        kb = db_client.create_knowledge_base(
            kb_id=kb_id,
            user_id=sample_user_id,
            project_id=sample_project_id,
            name="Test Knowledge Base",
            description="Testing KB",
        )

        assert kb["kb_id"] == kb_id
        assert kb["name"] == "Test Knowledge Base"

        # Add a fact
        fact_id = f"fact_{uuid.uuid4().hex[:8]}"
        embedding = [0.1] * 384

        fact = db_client.add_fact(
            fact_id=fact_id,
            kb_id=kb_id,
            content="Python was created by Guido van Rossum",
            embedding=embedding,
            metadata={"category": "history"},
        )

        assert fact["fact_id"] == fact_id
        assert fact["content"] == "Python was created by Guido van Rossum"

    def test_chat_operations(self, db_client, sample_user_id, sample_project_id):
        """Test chat operations."""
        # Create chats table
        db_client.create_chats_table(sample_user_id)

        # Add chat messages
        session_id = f"session_{uuid.uuid4().hex[:8]}"
        messages = [
            ("user", "Hello, how are you?"),
            ("assistant", "I'm doing well, thank you!"),
        ]

        for i, (role, content) in enumerate(messages):
            embedding = [0.1] * 384
            chat = db_client.add_chat_message(
                chat_id=f"chat_{i}",
                user_id=sample_user_id,
                project_id=sample_project_id,
                session_id=session_id,
                role=role,
                content=content,
                embedding=embedding,
            )

            assert chat["role"] == role
            assert chat["content"] == content

        # Get chat history
        history = db_client.get_chat_history(
            user_id=sample_user_id,
            session_id=session_id,
        )

        assert not history.is_empty()