Test your sql functions!

mail@pastecode.io avatar
unknown
python
a year ago
1.5 kB
8
Indexable
Never
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