Test your sql functions!
unknown
python
2 years ago
1.5 kB
12
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) == 1
Editor is loading...