class TestPostgresDB(unittest.TestCase):
def setUp(self):
self.temp_schema = TEMP_SCHEMA
self.pg = Postgres()
self.tbl = Table([['ID', 'Name'], [1, 'Jim'], [2, 'John'],
[3, 'Sarah']])
# Create a schema, create a table, create a view
setup_sql = f"""
drop schema if exists {self.temp_schema} cascade;
create schema {self.temp_schema};
"""
other_sql = f"""
create table {self.temp_schema}.test (id smallint,name varchar(5));
create view {self.temp_schema}.test_view as (select * from {self.temp_schema}.test);
""" # noqa: E501
self.pg.query(setup_sql)
self.pg.query(other_sql)
def tearDown(self):
# Drop the view, the table and the schema
teardown_sql = f"""
drop schema if exists {self.temp_schema} cascade;
"""
self.pg.query(teardown_sql)
def test_query(self):
# Check that query sending back expected result
r = self.pg.query('select 1')
self.assertEqual(r[0]['?column?'], 1)
def test_query_with_parameters(self):
table_name = f"{self.temp_schema}.test"
self.pg.copy(self.tbl, f"{self.temp_schema}.test", if_exists='append')
sql = f"select * from {table_name} where name = %s"
name = 'Sarah'
r = self.pg.query(sql, parameters=[name])
self.assertEqual(r[0]['name'], name)
sql = f"select * from {table_name} where name in (%s, %s)"
names = ['Sarah', 'John']
r = self.pg.query(sql, parameters=names)
self.assertEqual(r.num_rows, 2)
def test_copy(self):
# Copy a table and ensure table exists
self.pg.copy(self.tbl,
f'{self.temp_schema}.test_copy',
if_exists='drop')
r = self.pg.query(
f"select * from {self.temp_schema}.test_copy where name='Jim'")
self.assertEqual(r[0]['id'], 1)
# Copy table and ensure truncate works.
self.pg.copy(self.tbl,
f'{self.temp_schema}.test_copy',
if_exists='truncate')
tbl = self.pg.query(
f"select count(*) from {self.temp_schema}.test_copy")
self.assertEqual(tbl.first, 3)
# Copy table and ensure that drop works.
self.pg.copy(self.tbl,
f'{self.temp_schema}.test_copy',
if_exists='drop')
tbl = self.pg.query(
f"select count(*) from {self.temp_schema}.test_copy")
self.assertEqual(tbl.first, 3)
# Copy table and ensure that append works.
self.pg.copy(self.tbl,
f'{self.temp_schema}.test_copy',
if_exists='append')
tbl = self.pg.query(
f"select count(*) from {self.temp_schema}.test_copy")
self.assertEqual(tbl.first, 6)
# Try to copy the table and ensure that default fail works.
self.assertRaises(ValueError, self.pg.copy, self.tbl,
f'{self.temp_schema}.test_copy')
# Try to copy the table and ensure that explicit fail works.
self.assertRaises(ValueError,
self.pg.copy,
self.tbl,
f'{self.temp_schema}.test_copy',
if_exists='fail')
def test_to_postgres(self):
self.tbl.to_postgres(f'{self.temp_schema}.test_copy')
r = self.pg.query(
f"select * from {self.temp_schema}.test_copy where name='Jim'")
self.assertEqual(r[0]['id'], 1)
def test_from_postgres(self):
tbl = Table([['id', 'name'], [1, 'Jim'], [2, 'John'], [3, 'Sarah']])
self.pg.copy(self.tbl,
f'{self.temp_schema}.test_copy',
if_exists='drop')
out_tbl = self.tbl.from_postgres(
f"SELECT * FROM {self.temp_schema}.test_copy")
assert_matching_tables(out_tbl, tbl)
