replypier/backend/tests/knowledge/test_optimized_pgvector.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import pytest
from unittest.mock import MagicMock, patch
import numpy as np

from agno.document import Document
from agno.vectordb.pgvector.pgvector import PgVector
from app.knowledge.optimized_pgvector import OptimizedPgVector


class TestOptimizedPgVector:
    """Test cases for OptimizedPgVector"""
    
    @pytest.fixture
    def mock_embedder(self):
        """Create a mock embedder"""
        mock = MagicMock()
        mock.get_embedding.return_value = np.array([0.1, 0.2, 0.3])
        mock.dimensions = 3
        return mock
    
    @pytest.fixture
    def mock_session(self):
        """Create a mock database session"""
        mock = MagicMock()
        # Configure session to return a context manager
        mock.return_value.__enter__.return_value = mock
        return mock
    
    @patch('sqlalchemy.engine.create_engine')
    def test_init(self, mock_create_engine, mock_embedder):
        """Test initialization"""
        # Arrange
        mock_engine = MagicMock()
        mock_create_engine.return_value = mock_engine
        
        # Act
        db = OptimizedPgVector(
            table_name="test_table",
            db_url="postgresql+psycopg://user:pass@localhost/testdb",
            embedder=mock_embedder
        )
        
        # Assert
        assert db.table_name == "test_table"
        assert db.embedder == mock_embedder
        assert db.dimensions == 3
    
    @patch('sqlalchemy.engine.create_engine')
    def test_upsert_with_pre_embedded_documents(self, mock_create_engine, mock_embedder, mock_session):
        """Test upserting pre-embedded documents"""
        # Arrange
        mock_engine = MagicMock()
        mock_create_engine.return_value = mock_engine
        
        db = OptimizedPgVector(
            table_name="test_table",
            db_url="postgresql+psycopg://user:pass@localhost/testdb",
            embedder=mock_embedder
        )
        db.Session = mock_session
        db._clean_content = lambda x: x  # Mock the content cleaning function
        
        # Create a document with an existing embedding
        pre_embedded_doc = Document(
            id="test1",
            name="Test Document 1",
            content="This is a test document",
            embedding=np.array([0.5, 0.6, 0.7])
        )
        
        # Create a document without an embedding
        unembedded_doc = Document(
            id="test2",
            name="Test Document 2",
            content="This is another test document"
        )
        
        # Mock the document's embed method
        with patch.object(Document, 'embed') as mock_embed:
            # Act
            with patch('sqlalchemy.dialects.postgresql.insert') as mock_insert:
                mock_insert.return_value.on_conflict_do_update.return_value = "upsert_stmt"
                db.upsert([pre_embedded_doc, unembedded_doc])
            
            # Assert
            # The document's embed method should be called for the unembedded document
            mock_embed.assert_called_once_with(embedder=mock_embedder)
            
            # Both documents should be passed to the upsert statement
            assert mock_session.execute.call_count == 1
        
    @patch('sqlalchemy.engine.create_engine')
    def test_insert_with_pre_embedded_documents(self, mock_create_engine, mock_embedder, mock_session):
        """Test inserting pre-embedded documents"""
        # Arrange
        mock_engine = MagicMock()
        mock_create_engine.return_value = mock_engine
        
        db = OptimizedPgVector(
            table_name="test_table",
            db_url="postgresql+psycopg://user:pass@localhost/testdb",
            embedder=mock_embedder
        )
        db.Session = mock_session
        db._clean_content = lambda x: x  # Mock the content cleaning function
        
        # Create a document with an existing embedding
        pre_embedded_doc = Document(
            id="test1",
            name="Test Document 1",
            content="This is a test document",
            embedding=np.array([0.5, 0.6, 0.7])
        )
        
        # Create a document without an embedding
        unembedded_doc = Document(
            id="test2",
            name="Test Document 2",
            content="This is another test document"
        )
        
        # Mock the document's embed method
        with patch.object(Document, 'embed') as mock_embed:
            # Act
            with patch('sqlalchemy.dialects.postgresql.insert') as mock_insert:
                db.insert([pre_embedded_doc, unembedded_doc])
            
            # Assert
            # The document's embed method should be called for the unembedded document
            mock_embed.assert_called_once_with(embedder=mock_embedder)
            
            # Both documents should be passed to the insert statement
            assert mock_session.execute.call_count == 1


if __name__ == "__main__":
    pytest.main(["-v", "test_optimized_pgvector.py"])


class _PgError(Exception):
    """Stand-in for a psycopg error, which carries the SQLSTATE as .sqlstate."""

    def __init__(self, sqlstate):
        super().__init__(sqlstate)
        self.sqlstate = sqlstate


class TestConcurrentTableCreate:
    """The per-org vector table is created by whichever queue item gets there first.

    agno's PgVector.create() does table_exists() and then CREATE TABLE, which is
    not atomic. The processor runs two items at once, so an organisation whose
    first two sources are queued together had both workers create the table and
    the loser died with
    `duplicate key value violates unique constraint "pg_type_typname_nsp_index"`,
    losing that source (production queue item 372, goldmaidsteam.com).
    """

    def _vector_db(self):
        return OptimizedPgVector.__new__(OptimizedPgVector)

    @pytest.mark.parametrize("sqlstate", ["42P07", "23505"])
    def test_losing_the_create_race_is_not_an_error(self, sqlstate):
        db = self._vector_db()
        db.table_name = "d_org"
        db.schema = "ai"

        with patch.object(PgVector, "create", side_effect=_PgError(sqlstate)):
            db.create()  # the other worker created it; that is the table we wanted

    def test_a_real_create_failure_still_raises(self):
        db = self._vector_db()
        db.table_name = "d_org"
        db.schema = "ai"

        with patch.object(PgVector, "create", side_effect=_PgError("42501")):
            with pytest.raises(_PgError):
                db.create()

    def test_duplicate_detected_through_sqlalchemy_wrapping(self):
        """SQLAlchemy wraps the driver error; the SQLSTATE lives on .orig."""
        wrapped = RuntimeError("duplicate")
        wrapped.orig = _PgError("23505")

        with patch.object(PgVector, "create", side_effect=wrapped):
            db = self._vector_db()
            db.table_name = "d_org"
            db.schema = "ai"
            db.create()