def test_read_with_row_filter(df_type, backend: PostgresBackend, filter_expr,
expected_data):
"""Is the full dataset correctly read?"""
reset_test_table(True)
if df_type == pd.DataFrame:
df = backend.read_to_pandas("testdb", row_filter=filter_expr)
elif df_type == dict:
df = backend.read_to_dict("testdb", row_filter=filter_expr)
expected_data.assert_correct_and_equal(df)
def test_read_to_dict_sample_4(backend: PostgresBackend):
"""Is the full dataset correctly read?"""
reset_test_table(True)
df = backend.read_to_dict("testdb", sample=4)
df = pd.DataFrame.from_records(df)
SampleDataSchema.validate(df)
assert len(df) == 4
def reset_test_table(refill: bool = False):
"""Empty the test db table user1.testdb and maybe refill it with sample data
Args:
refill: Whether to fill the emptied table with sample data again
"""
with psycopg.connect(conninfo=CONNECTION_STRING) as conn:
with conn.cursor() as cursor:
cursor.execute("DELETE FROM user1.testdb")
if refill:
PostgresBackend(conninfo=CONNECTION_STRING).write_append(
"testdb",
SampleDataSet().dataframe())
def backend(prepare_test_db) -> PostgresBackend:
"""Initialized PostgresBackend pointing to the test database."""
return PostgresBackend(conninfo=CONNECTION_STRING)
def test_read_to_dict_drop_duplicates(backend: PostgresBackend):
"""Is the full dataset correctly read?"""
reset_test_table(True)
df = backend.read_to_dict("testdb", drop_duplicates=True)
sample_data = SampleDataSet()
sample_data.first_rows(len(sample_data) - 1).assert_correct_and_equal(df)
def test_read_to_pandas_sample_4(backend: PostgresBackend):
"""Is the full dataset correctly read?"""
reset_test_table(True)
df = backend.read_to_pandas("testdb", sample=4)
SampleDataSchema.validate(df)
assert len(df) == 4
def test_read_to_dict_top_3(backend: PostgresBackend):
"""Is the full dataset correctly read?"""
reset_test_table(True)
df = backend.read_to_dict("testdb", limit=3)
SampleDataSet().first_rows(3).assert_correct_and_equal(df)
def test_read_to_dict_columns(backend: PostgresBackend):
"""Is the full dataset correctly read?"""
reset_test_table(True)
df = backend.read_to_dict("testdb", columns=["col_int", "col_string"])
SampleDataSet().select_columns(
columns=["col_int", "col_string"]).assert_correct_and_equal(df)
def test_read_to_dict_empty_table(backend: PostgresBackend):
"""Try reading an empty table and see if the column names are fetched"""
df = backend.read_to_dict("empty")
assert df == dict(col1=[], col2=[])
def test_read_to_pandas_empty_table(backend: PostgresBackend):
"""Try reading an empty table and see if the column names are fetched"""
df = backend.read_to_pandas("empty")
pd.testing.assert_frame_equal(df, pd.DataFrame(columns=["col1", "col2"]))
def test_read_to_pandas_argchecks(backend: PostgresBackend, kwargs, exception):
"""Challenge the argument validation of the constructor"""
with pytest.raises(exception):
backend.read_to_pandas(**kwargs)