🔧 chore: Update test configuration and add integration tests for pgvector

- Modify `pytest.ini` to exclude integration tests by default, requiring explicit invocation.
- Replace `pytest-postgresql` with `testcontainers[postgres]` in `test_requirements.txt` for better container management.
- Update CI workflow to run unit and integration tests separately, generating distinct test result files.
- Introduce shared fixtures and integration tests for pgvector, ensuring proper SQL query performance and index usage in a real PostgreSQL environment.

Author: danny-avila

.github/workflows/ci.yml

@@ -32,14 +32,21 @@ jobs:
           pip install -r requirements.txt
           pip install -r test_requirements.txt
 
-      - name: Run tests and generate report
+      - name: Run unit tests
         env:
           OPENAI_API_KEY: test_key
         run: |
-          pytest --maxfail=1 --disable-warnings --junitxml=test-results.xml
+          pytest --maxfail=1 --disable-warnings --junitxml=test-results-unit.xml
+
+      - name: Run integration tests (pgvector container)
+        env:
+          OPENAI_API_KEY: test_key
+        run: |
+          pytest -m integration --maxfail=3 --disable-warnings --junitxml=test-results-integration.xml
 
       - name: Upload test results
+        if: always()
         uses: actions/upload-artifact@v4
         with:
           name: test-results
-          path: test-results.xml
+          path: test-results-*.xml

pytest.ini

@@ -13,8 +13,8 @@ markers =
     integration: marks tests as integration tests (may require more resources)
     slow: marks tests as slow running
 
-# Default options
-addopts = -v --tb=short
+# Default options — exclude integration tests (require Docker) unless explicitly requested
+addopts = -v --tb=short -m "not integration"
 
 # Filter warnings
 filterwarnings =

test_requirements.txt

@@ -1,5 +1,5 @@
 pytest==8.3.4
 pytest-asyncio==0.26.0
-pytest-postgresql==7.0.1
+testcontainers[postgres]>=4.14.0
 mongomock==4.3.0
 httpx==0.27.0

tests/integration/__init__.py

@@ -0,0 +1,114 @@
+"""Shared fixtures for integration tests that use a real pgvector PostgreSQL container.
+
+Equivalent to mongodb-memory-server in Node.js: spins up a real, ephemeral
+PostgreSQL instance with pgvector for production-parity testing.
+"""
+
+import pytest
+import sqlalchemy
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+from testcontainers.postgres import PostgresContainer
+
+PGVECTOR_IMAGE = "pgvector/pgvector:pg16"
+
+
+@pytest.fixture(scope="session")
+def pg_container():
+    """Start a pgvector PostgreSQL container once for the entire test session."""
+    with PostgresContainer(PGVECTOR_IMAGE, driver="psycopg2") as pg:
+        yield pg
+
+
+@pytest.fixture(scope="session")
+def pg_url(pg_container):
+    """SQLAlchemy connection URL for the test container."""
+    return pg_container.get_connection_url()
+
+
+@pytest.fixture(scope="session")
+def engine(pg_url):
+    """Session-scoped SQLAlchemy engine connected to the test container."""
+    eng = sqlalchemy.create_engine(pg_url)
+    with eng.begin() as conn:
+        conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
+    yield eng
+    eng.dispose()
+
+
+@pytest.fixture(scope="session")
+def _create_tables(engine):
+    """Create the langchain tables and indexes once for the session.
+
+    Mirrors the production schema created by LangChain PGVector + our
+    ensure_vector_indexes() startup logic.
+    """
+    with engine.begin() as conn:
+        conn.execute(
+            text(
+                """
+            CREATE TABLE IF NOT EXISTS langchain_pg_collection (
+                uuid UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+                name VARCHAR NOT NULL UNIQUE,
+                cmetadata JSONB
+            )
+        """
+            )
+        )
+        conn.execute(
+            text(
+                """
+            CREATE TABLE IF NOT EXISTS langchain_pg_embedding (
+                uuid UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+                collection_id UUID REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE,
+                embedding vector(3),
+                document VARCHAR,
+                cmetadata JSONB,
+                custom_id VARCHAR
+            )
+        """
+            )
+        )
+        conn.execute(
+            text(
+                """
+            CREATE INDEX IF NOT EXISTS idx_langchain_pg_embedding_file_id
+            ON langchain_pg_embedding ((cmetadata->>'file_id'))
+        """
+            )
+        )
+        conn.execute(
+            text(
+                """
+            CREATE INDEX IF NOT EXISTS ix_cmetadata_gin
+            ON langchain_pg_embedding
+            USING gin (cmetadata jsonb_path_ops)
+        """
+            )
+        )
+
+
+@pytest.fixture(scope="session")
+def collection_id(engine, _create_tables):
+    """Insert a test collection and return its UUID."""
+    with engine.begin() as conn:
+        row = conn.execute(
+            text(
+                "INSERT INTO langchain_pg_collection (name, cmetadata) "
+                "VALUES (:name, :meta) RETURNING uuid"
+            ),
+            {"name": "test_collection", "meta": "{}"},
+        ).fetchone()
+    return row[0]
+
+
+@pytest.fixture()
+def db_session(engine, _create_tables):
+    """Per-test session that rolls back after each test for isolation."""
+    conn = engine.connect()
+    trans = conn.begin()
+    session = Session(bind=conn)
+    yield session, conn
+    session.close()
+    trans.rollback()
+    conn.close()