Skip to content

Commit

Permalink
Add testing with vcr.py and a compose file
Browse files Browse the repository at this point in the history
  • Loading branch information
Askir committed Oct 25, 2024
1 parent 408148e commit 8eb554b
Show file tree
Hide file tree
Showing 9 changed files with 2,161 additions and 16 deletions.
33 changes: 28 additions & 5 deletions .github/workflows/test.yaml
Original file line number Diff line number Diff line change
@@ -1,8 +1,31 @@
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 pip install .
uv pip install .[dev]
- 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
12 changes: 12 additions & 0 deletions docker-compose.yaml
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -116,4 +116,5 @@ dev-dependencies = [
"pandas>=2.2.3",
"pytest-asyncio>=0.24.0",
"pyright>=1.1.386",
"vcrpy>=6.0.2",
]
14 changes: 11 additions & 3 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,20 @@

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"]


Expand Down
4 changes: 2 additions & 2 deletions tests/pg_vectorizer_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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"]:
Expand Down
56 changes: 56 additions & 0 deletions tests/utils.py
Original file line number Diff line number Diff line change
@@ -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,
)
Loading

0 comments on commit 8eb554b

Please sign in to comment.