Add testing with vcr.py and a compose file

Author: Askir

.github/workflows/ruff.yaml

@@ -19,8 +19,7 @@ jobs:
       run: uv venv
     - name: Install dependencies
       run: |
-        uv pip install .
-        uv pip install .[dev]
+        uv sync
     - name: Run Ruff linter
       run: uv run ruff check .
     - name: Run Ruff formatter

.github/workflows/test.yaml

@@ -1,8 +1,30 @@
-name: CI
-on:  [workflow_dispatch, pull_request, push]
+name: Tests
+
+on:
+  pull_request:
+    branches: [ main ]
 
 jobs:
-  test:
-    if: false
+  pytest:
     runs-on: ubuntu-latest
-    steps: [uses: fastai/workflows/nbdev-ci@master]
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+    - name: Install uv
+      run: pip install uv
+    - name: Create venv
+      run: uv venv
+    - name: Install dependencies
+      run: |
+        uv sync
+    - name: Start docker-compose 
+      run: docker compose up -d
+    - name: Run Test
+      run: uv run pytest
+    - name: Logs
+      run: docker compose logs
+    - name: Stop docker-compose
+      run: docker compose down

docker-compose.yaml

@@ -0,0 +1,12 @@
+services:
+  db:
+    image: timescale/timescaledb-ha:pg16
+    ports:
+      - "5432:5432"
+    environment:
+      - POSTGRES_PASSWORD=postgres
+      - POSTGRES_USER=postgres
+      - POSTGRES_DB=postgres
+      - TIMESCALEDB_TELEMETRY=off
+    volumes:
+      - ./data:/var/lib/postgresql/data

pyproject.toml

@@ -116,4 +116,5 @@ dev-dependencies = [
     "pandas>=2.2.3",
     "pytest-asyncio>=0.24.0",
     "pyright>=1.1.386",
+    "vcrpy>=6.0.2",
 ]

tests/conftest.py

@@ -2,19 +2,28 @@
 
 import psycopg2
 import pytest
-from dotenv import find_dotenv, load_dotenv
+
+# from dotenv import find_dotenv, load_dotenv
+
+
+@pytest.fixture(scope="module")
+def setup_env_variables() -> None:
+    os.environ.clear()
+    os.environ["TIMESCALE_SERVICE_URL"] = "postgres://postgres:postgres@localhost:5432/postgres"
+    os.environ["OPENAI_API_KEY"] = "fake key"
 
 
 @pytest.fixture(scope="module")
-def service_url() -> str:
-    _ = load_dotenv(find_dotenv(), override=True)
+def service_url(setup_env_variables: None) -> str:  # noqa: ARG001
+    # _ = load_dotenv(find_dotenv(), override=True)
     return os.environ["TIMESCALE_SERVICE_URL"]
 
 
 @pytest.fixture(scope="module", autouse=True)
-def create_temp_schema(service_url: str) -> None:
+def setup_db(service_url: str) -> None:
     conn = psycopg2.connect(service_url)
     with conn.cursor() as cursor:
+        cursor.execute("CREATE EXTENSION IF NOT EXISTS ai CASCADE;")
         cursor.execute("CREATE SCHEMA IF NOT EXISTS temp;")
     conn.commit()
     conn.close()

tests/pg_vectorizer_test.py

@@ -2,12 +2,12 @@
 from typing import Any
 
 import psycopg2
-import pytest
 from langchain.docstore.document import Document
 from langchain.text_splitter import CharacterTextSplitter
 from langchain_community.vectorstores.timescalevector import TimescaleVector
 from langchain_openai import OpenAIEmbeddings
 
+from tests.utils import http_recorder
 from timescale_vector import client
 from timescale_vector.pgvectorizer import Vectorize
 
@@ -31,7 +31,7 @@ def get_document(blog: dict[str, Any]) -> list[Document]:
     return docs
 
 
-@pytest.mark.skip(reason="requires OpenAI API key")
+@http_recorder.use_cassette("pg_vectorizer.yaml")
 def test_pg_vectorizer(service_url: str) -> None:
     with psycopg2.connect(service_url) as conn, conn.cursor() as cursor:
         for item in ["blog", "blog_embedding_work_queue", "blog_embedding"]:

tests/utils.py

@@ -0,0 +1,56 @@
+import os
+from typing import Any
+
+import vcr
+
+vcr_cassette_path = os.path.join(os.path.dirname(__file__), "vcr_cassettes")
+
+
+def remove_set_cookie_header(response: dict[str, Any]):
+    """
+    Removes the Set-Cookie header from a VCR.py response object to improve cassette consistency.
+
+    This function can be used as a before_record callback in your VCR configuration
+    to ensure that Set-Cookie headers are stripped from responses before they are
+    recorded to cassettes.
+
+    Args:
+        response (vcr.request.Response): The VCR.py response object to modify
+
+    Returns:
+        vcr.request.Response: The modified response object with Set-Cookie headers removed
+
+    Example:
+        import vcr
+
+        # Configure VCR with the callback
+        vcr = vcr.VCR(
+            before_record_response=remove_set_cookie_header,
+            match_on=['uri', 'method']
+        )
+
+        with vcr.use_cassette('tests/fixtures/my_cassette.yaml'):
+            # Make your HTTP requests here
+            pass
+    """
+
+    # Get the headers from the response
+    headers = response["headers"]
+
+    # Headers to remove (case-insensitive)
+    headers_to_remove = ["set-cookie", "Set-Cookie"]
+
+    # Remove Set-Cookie headers if they exist
+    for header in headers_to_remove:
+        if header in headers:
+            del headers[header]
+
+    return response
+
+
+http_recorder = vcr.VCR(
+    cassette_library_dir=vcr_cassette_path,
+    record_mode="once",
+    filter_headers=["authorization", "cookie"],
+    before_record_response=remove_set_cookie_header,
+)