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"