Skip to content

Commit

Permalink
Updated tests to support SQLAlchemy 1
Browse files Browse the repository at this point in the history
  • Loading branch information
ankane committed Oct 2, 2024
1 parent 8879d95 commit 89ec21d
Showing 1 changed file with 10 additions and 2 deletions.
12 changes: 10 additions & 2 deletions tests/test_sqlalchemy.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,17 @@
from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer
from sqlalchemy.exc import StatementError
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.orm import declarative_base, mapped_column, Session
from sqlalchemy.orm import declarative_base, Session
from sqlalchemy.sql import func

try:
from sqlalchemy.orm import mapped_column
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
sqlalchemy_version = 2
except ImportError:
mapped_column = Column
sqlalchemy_version = 1

engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test')
with Session(engine) as session:
session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
Expand Down Expand Up @@ -418,6 +425,7 @@ def test_automap(self):
assert item.embedding.tolist() == [1, 2, 3]

@pytest.mark.asyncio
@pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
async def test_async(self):
engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test')
async_session = async_sessionmaker(engine, expire_on_commit=False)
Expand Down

0 comments on commit 89ec21d

Please sign in to comment.