def stop_and_clear_on_disk_postgresql_database() -> None:
"""Drops all tables in the local postgres database and stops the postgres server. Should be called in the
tearDownClass function so this only runs once per test class."""
if environment.in_gae():
raise ValueError('Running test-only code in Google App Engine.')
for declarative_base in DECLARATIVE_BASES:
use_on_disk_postgresql_database(declarative_base)
declarative_base.metadata.drop_all(SQLAlchemyEngineManager.get_engine_for_schema_base(declarative_base))
SQLAlchemyEngineManager.teardown_engine_for_schema(declarative_base)
if not environment.in_travis():
# If we are running locally, we must stop the local postgres server
os.system('pg_ctl -D /usr/local/var/postgres stop &> /dev/null')
def teardown_on_disk_postgresql_database(declarative_base: DeclarativeMeta) -> None:
"""Clears state in an on-disk postgres database for a given schema, for use once a single test has completed. As an
optimization, does not actually drop tables, just clears them. As a best practice, you should call
stop_and_clear_on_disk_postgresql_database() once all tests in a test class are complete to actually drop the
tables.
"""
if environment.in_gae():
raise ValueError('Running test-only code in Google App Engine.')
session = SessionFactory.for_schema_base(declarative_base)
try:
for table in reversed(declarative_base.metadata.sorted_tables):
session.execute(table.delete())
session.commit()
except Exception as e:
session.rollback()
raise e
finally:
session.close()
SQLAlchemyEngineManager.teardown_engine_for_schema(declarative_base)
def teardown_on_disk_postgresql_database(
declarative_base: DeclarativeMeta) -> None:
"""Clears state in an on-disk postgres database for a given schema, for use once a single test has completed. As an
optimization, does not actually drop tables, just clears them. As a best practice, you should call
stop_and_clear_on_disk_postgresql_database() once all tests in a test class are complete to actually drop the
tables.
"""
# Ensure all sessions are closed, otherwise the below may hang.
close_all_sessions()
session = SessionFactory.for_schema_base(declarative_base)
try:
for table in reversed(declarative_base.metadata.sorted_tables):
session.execute(table.delete())
session.commit()
except Exception as e:
session.rollback()
raise e
finally:
session.close()
SQLAlchemyEngineManager.teardown_engine_for_schema(declarative_base)