def test_delete_data(pg, init_schema): """Tests that delete_data removes all data from specified tables.""" path = pathlib.Path(__file__) / '..' / 'pg_data' / '001_test_table.csv' with psycopg2.connect(**pg.dsn()) as conn: lathorp.copy_data(conn, path) with conn.cursor() as cursor: cursor.execute('SELECT count(*) FROM test_table;') assert cursor.fetchone()[0] == 5 lathorp.delete_data(conn, {'test_table'}) with conn.cursor() as cursor: cursor.execute('SELECT count(*) FROM test_table;') assert cursor.fetchone()[0] == 0
def test_copy_data_file(pg, init_schema): """Tests that copy_data can load data from a file into the database.""" path = pathlib.Path(__file__) / '..' / 'pg_data' / '001_test_table.csv' assert path.resolve().is_file() with psycopg2.connect(**pg.dsn()) as conn: lathorp.copy_data(conn, path) with conn.cursor() as cursor: cursor.execute('SELECT num FROM test_table;') assert cursor.fetchall() == [(1001, ), (1002, ), (1003, ), (None, ), (1005, )] # Clean up for the next tests cursor.execute('TRUNCATE test_table CASCADE;')
def test_copy_data_raises_on_invalid_path(pg): """Tests copy_data fixture with exceptional conditions. The function is expected to raise a ValueError when given invalid paths, or paths that contain files whose names are in an unexpected format. """ with psycopg2.connect(**pg.dsn()) as conn: # The file name is in the wrong format: with pytest.raises(ValueError): lathorp.copy_data(conn, pathlib.Path(__file__)) # Missing path: path = pathlib.Path(__file__) / 'nosuchpath' assert not path.exists() with pytest.raises(ValueError): lathorp.copy_data(conn, path)
def test_copy_data_dir(pg, init_schema): """Tests that copy_data can load data from all files in a directory.""" path = pathlib.Path(__file__) / '..' / 'pg_data' assert path.resolve().is_dir() assert len(list(path.resolve().iterdir())) > 1 with psycopg2.connect(**pg.dsn()) as conn: lathorp.copy_data(conn, path) with conn.cursor() as cursor: cursor.execute('SELECT num FROM test_table;') assert cursor.fetchall() == [ (1001, ), (1002, ), (1003, ), (None, ), (1005, ), (2001, ), (2002, ), (2003, ), (None, ), (2005, ) ] # Clean up for the next tests cursor.execute('TRUNCATE test_table CASCADE;')
def connector(data_path=None, cursor_factory=None): dsn = psycopg2.extensions.make_dsn(**pg.dsn()) conn = psycopg2.connect(dsn, cursor_factory=cursor_factory) tables = [] if data_path: tables = lathorp.copy_data(conn, data_path) def finalize(): lathorp.delete_data(conn, tables) conn.close() request.addfinalizer(finalize) return conn