Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_iter_rows_dict_factory(pgtestdb_test_tables, pgtestdb_conn):
sql = "SELECT * FROM src"
result = iter_rows(sql, pgtestdb_conn, row_factory=dict_rowfactory)
expected = [
{'id': 1, 'value': 1.234, 'simple_text': 'text', 'utf8_text': 'Öæ°\nz',
'day': datetime.date(2018, 12, 7),
'date_time': datetime.datetime(2018, 12, 7, 13, 1, 59)},
{'id': 2, 'value': 2.234, 'simple_text': 'text', 'utf8_text': 'Öæ°\nz',
'day': datetime.date(2018, 12, 8),
'date_time': datetime.datetime(2018, 12, 8, 13, 1, 59)},
{'id': 3, 'value': 2.234, 'simple_text': 'text', 'utf8_text': 'Öæ°\nz',
'day': datetime.date(2018, 12, 9),
'date_time': datetime.datetime(2018, 12, 9, 13, 1, 59)},
]
assert list(result) == expected
def test_copy_rows_transform(pgtestdb_conn, pgtestdb_test_tables, my_transform):
# Arrange
select_sql = "SELECT * FROM src"
insert_sql = "INSERT INTO dest (id) VALUES (%s)"
expected = [(2, None, None, None, None, None),
(3, None, None, None, None, None)]
# Act
copy_rows(select_sql, pgtestdb_conn, insert_sql, pgtestdb_conn,
transform=my_transform)
# Assert
sql = "SELECT * FROM dest"
result = iter_rows(sql, pgtestdb_conn)
assert list(result) == expected
def test_iter_rows_no_results(pgtestdb_test_tables, pgtestdb_conn):
sql = "SELECT * FROM src WHERE id = -1"
result = iter_rows(sql, pgtestdb_conn)
assert list(result) == []
def test_iter_rows_transform(pgtestdb_test_tables, pgtestdb_conn,
test_table_data):
# Arrange
sql = "SELECT * FROM src"
def my_transform(rows):
# Simple transform function that changes size and number of rows
return [row.id for row in rows if row.id > 1]
# Act
result = iter_rows(sql, pgtestdb_conn, transform=my_transform)
# Assert
expected = [row[0] for row in test_table_data if row[0] > 1]
assert list(result) == expected
def test_copy_rows_happy_path(pgtestdb_conn, pgtestdb_test_tables,
pgtestdb_insert_sql, test_table_data):
# Arrange and act
select_sql = "SELECT * FROM src"
insert_sql = pgtestdb_insert_sql.replace('src', 'dest')
copy_rows(select_sql, pgtestdb_conn, insert_sql, pgtestdb_conn)
# Assert
sql = "SELECT * FROM dest"
result = iter_rows(sql, pgtestdb_conn)
assert list(result) == test_table_data
def test_iter_rows_happy_path(pgtestdb_test_tables, pgtestdb_conn,
test_table_data):
sql = "SELECT * FROM src"
result = iter_rows(sql, pgtestdb_conn)
assert list(result) == test_table_data
def test_iter_rows_bad_query(pgtestdb_test_tables, pgtestdb_conn):
sql = "SELECT * FROM public.this_does_not_exist"
with pytest.raises(ETLHelperExtractError):
result = iter_rows(sql, pgtestdb_conn)
list(result) # Call list to activate returned generator