class PostgresExplorerTest(CommonExplorerTestCases.CommonExplorerTests):
pii_db_drop = """
DROP TABLE full_pii;
DROP TABLE partial_pii;
DROP TABLE no_pii;
"""
@staticmethod
def execute_script(cursor, script):
for query in script.split(';'):
if len(query.strip()) > 0:
cursor.execute(query)
@pytest.fixture(scope="class")
def create_tables(self, request):
self.conn = psycopg2.connect(host="127.0.0.1",
user="******",
password="******",
database="piidb")
self.conn.autocommit = True
with self.conn.cursor() as cursor:
self.execute_script(cursor, pii_data_script)
cursor.close()
def drop_tables():
with self.conn.cursor() as d_cursor:
d_cursor.execute(self.pii_db_drop)
logging.info("Executed drop script")
d_cursor.close()
self.conn.close()
request.addfinalizer(drop_tables)
def setUp(self):
self.explorer = PostgreSQLExplorer(
Namespace(host="127.0.0.1",
user="******",
password="******",
database="piidb",
catalog=None))
def tearDown(self):
self.explorer.get_connection().close()
def test_schema(self):
names = [sch.get_name() for sch in self.explorer.get_schemas()]
self.assertEqual(['public'], names)
def get_test_schema(self):
return "public"
class PostgresExplorerTest(CommonExplorerTestCases.CommonExplorerTests):
pii_db_drop = """
DROP TABLE full_pii;
DROP TABLE partial_pii;
DROP TABLE no_pii;
DROP SCHEMA company cascade;
"""
second_schema = """
CREATE SCHEMA company;
CREATE TABLE company.employees(name varchar, designation varchar);
CREATE TABLE company.departments(name varchar, manager varchar);
"""
@staticmethod
def execute_script(cursor, script):
for query in script.split(";"):
if len(query.strip()) > 0:
cursor.execute(query)
@pytest.fixture(scope="class")
def create_tables(self, request):
self.conn = psycopg2.connect(host="127.0.0.1",
user="******",
password="******",
database="piidb")
self.conn.autocommit = True
with self.conn.cursor() as cursor:
self.execute_script(cursor, pii_data_script)
self.execute_script(cursor, self.second_schema)
cursor.close()
def drop_tables():
with self.conn.cursor() as d_cursor:
d_cursor.execute(self.pii_db_drop)
logging.info("Executed drop script")
d_cursor.close()
self.conn.close()
request.addfinalizer(drop_tables)
def setUp(self):
self.explorer = PostgreSQLExplorer(
Namespace(
host="127.0.0.1",
user="******",
password="******",
database="piidb",
include_schema=(),
exclude_schema=(),
include_table=(),
exclude_table=(),
catalog=None,
))
def tearDown(self):
self.explorer.get_connection().close()
def test_schema(self):
names = [sch.get_name() for sch in self.explorer.get_schemas()]
self.assertCountEqual(["public", "company"], names)
def get_test_schema(self):
return "public"