Test your sql functions!
unknown
python
3 years ago
1.5 kB
13
Indexable
import os
import pandas as pd
from sqlalchemy import (
BigInteger,
Column,
DateTime,
Float,
MetaData,
String,
Text,
)
BASE_PATH = os.path.split(os.path.abspath(__file__))[0] + "/parquets"
def load_records(file_path):
df = pd.read_parquet(file_path)
df = df.replace({np.nan: None})
return df.to_dict(orient="records")
class Student:
__tablename__ = 'student'
id = Column("id", BigInteger, primary_key=True)
first_name = Column('first_name', String)
last_name = Column('last_name', String)
@pytest.fixture(scope="module")
def database():
with testing.postgresql.Postgresql() as postgresql:
meta = MetaData()
engine = create_engine(postgresql.url()) # <- spins up a dummy engine
meta.create_all(engine) # <- will create all tables as defined in ORM if it's available in the file
Session = sessionmaker(engine)
with Session() as session:
table_names = [
"student"
]
for tn in table_names:
session.execute(
meta.tables[tn].insert().values(load_records(os.path.join(BASE_PATH, f"{tn}/{tn}.parquet")))
)
session.commit()
yield engine
-----
def test_sql_function(database):
Session = sessionmaker(engine)
with Session() as session:
res = session.execute("select * from students where first_name = 'Jon'").all()
assert len(res) == 1Editor is loading...