def test_nested_sql_query(self):
# launching first container
with PostgresContainer("postgres:14.1") as database_first:
env_vars = get_env_vars(database_first, "first")
# launching second container
with PostgresContainer("postgres:14.1") as database_second:
env_vars = {
**get_env_vars(database_second, "second"),
**env_vars
}
# patching ENV vars with connection to both containers
with mock.patch.dict(os.environ, env_vars):
pgmagic.Postgres().load_params()
# starting parent session
with pgmagic.session("first") as s:
ret_first = s.execute("SELECT 1").fetchone()
# starting nested session
with pgmagic.session("second") as s:
ret_second = s.execute("SELECT 1").fetchone()
self.assertEqual(ret_first[0], ret_second[0])
def start_postgres_container(self, retries):
for i in range(retries):
try:
self.postgres = PostgresContainer('postgres:12.3')
self.postgres.start()
break
except Exception as e: # pylint: disable=bare-except
if i == retries - 1:
logging.error('Unable to initialize postgreSQL container.')
raise e
def setUp(self):
self.postgres = PostgresContainer('postgres:latest')
self.postgres.start()
self.engine = sqlalchemy.create_engine(
self.postgres.get_connection_url())
self.username = '******'
self.password = '******'
self.host = self.postgres.get_container_host_ip()
self.port = self.postgres.get_exposed_port(5432)
self.database_name = 'test'
self.driver_class_name = 'org.postgresql.Driver'
self.jdbc_url = 'jdbc:postgresql://{}:{}/{}'.format(
self.host, self.port, self.database_name)
def start_db_container(self, retries):
for i in range(retries):
try:
self.db = PostgresContainer('debezium/example-postgres:latest',
user=self.username,
password=self.password,
dbname=self.database)
self.db.start()
break
except Exception as e: # pylint: disable=bare-except
if i == retries - 1:
logging.error('Unable to initialize DB container.')
raise e
def test_create_session(self):
with PostgresContainer("postgres:14.1") as c:
env_vars = get_env_vars(c)
with mock.patch.dict(os.environ, env_vars):
custom = pgmagic.pg_cfg(database="default")
self.assertEqual(custom, c.get_connection_url())
def test_docker_run_postgress():
postgres_container = PostgresContainer("postgres:9.5")
with postgres_container as postgres:
e = sqlalchemy.create_engine(postgres.get_connection_url())
result = e.execute("select version()")
for row in result:
print("server version:", row[0])
def test_docker_run_postgres():
postgres_container = PostgresContainer("postgres:9.5")
with postgres_container as postgres:
e = sqlalchemy.create_engine(postgres.get_connection_url())
result = e.execute("select version()")
for row in result:
assert row[0].lower().startswith("postgresql 9.5")
def setUp(self) -> None:
self.sample_table_name = 'crif_sample'
self.file_type = 'customers'
self.sample_df = pd.DataFrame(sample_dict)
self.postgres_container = PostgresContainer("postgres:9.5").start()
with create_engine(self.postgres_container.get_connection_url()
).connect() as conn, conn.begin():
conn.execute("create schema if not exists mart_compliance ")
self.sample_df.to_sql(self.sample_table_name,
conn,
schema='mart_compliance')
self.crif_table = CrifTable(
self.sample_table_name, 'customers', datetime.now(),
self.postgres_container.get_connection_url())
def test_docker_run_greenplum():
postgres_container = PostgresContainer("datagrip/greenplum:6.8",
user="******", password="******", dbname="guest")
with postgres_container as postgres:
e = sqlalchemy.create_engine(postgres.get_connection_url())
result = e.execute("select version()")
for row in result:
print("server version:", row[0])
def db():
environ["POSTGRES_DB"] = "user-manager"
postgres_container = PostgresContainer("postgres:12.2-alpine")
with postgres_container as postgres:
environ["DATABASE_URL"] = postgres.get_connection_url()
database = Database(get_test_settings())
system("alembic upgrade head")
yield database
def session():
with PostgresContainer('postgres:9.6') as psql:
engine = create_engine(psql.get_connection_url())
ModelBase.metadata.create_all(engine)
Session = sessionmaker(bind=engine, expire_on_commit=False)
session = Session()
yield session
session.close()
def run_tests(self, test_labels, extra_tests=None, **kwargs):
with PostgresContainer('postgres:12.3') as container:
conf.settings.DATABASES['default'][
'HOST'] = container.get_container_host_ip()
conf.settings.DATABASES['default'][
'PORT'] = container.get_exposed_port(self.POSTGRESQL_PORT)
conf.settings.DATABASES['default']['USER'] = "******"
conf.settings.DATABASES['default']['PASSWORD'] = "******"
super().run_tests(test_labels, extra_tests=None, **kwargs)
def db(scope="session"):
with PostgresContainer("postgres:11.4") as postgres:
engine = create_engine(postgres.get_connection_url())
connection = engine.raw_connection()
yield {
"connection": connection,
"engine": engine,
}
connection.close()
def setup():
environ["POSTGRES_DB"] = "cash-flow"
postgres_container = PostgresContainer("postgres:12.2-alpine")
with postgres_container as postgres:
environ["DATABASE_URL"] = postgres.get_connection_url()
engine = create_engine(get_url())
system("alembic upgrade head")
Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
yield Session
def db_session(request):
postgres = PostgresContainer()
postgres.start()
db_url = postgres.get_connection_url()
engine = create_engine(db_url)
session = scoped_session(
sessionmaker(autocommit=False, autoflush=False, bind=engine))
Base.query = session.query_property()
Base.metadata.create_all(engine)
def stop_db():
postgres.stop()
request.addfinalizer(stop_db)
return session
def pytest_configure(config: Config) -> None:
app_config = load_configuration("testing")
if app_config.tests.testcontainers:
config.database_container = PostgresContainer(
image="postgres:12",
user="******",
password="******",
dbname="backend",
)
config.redis_container = RedisContainer(image="redis:alpine")
config.database_container.start()
config.redis_container.start()
class TestPullCrifTable(unittest.TestCase):
def setUp(self) -> None:
self.sample_table_name = 'crif_sample'
self.file_type = 'customers'
self.sample_df = pd.DataFrame(sample_dict)
self.postgres_container = PostgresContainer("postgres:9.5").start()
with create_engine(self.postgres_container.get_connection_url()
).connect() as conn, conn.begin():
conn.execute("create schema if not exists mart_compliance ")
self.sample_df.to_sql(self.sample_table_name,
conn,
schema='mart_compliance')
self.crif_table = CrifTable(
self.sample_table_name, 'customers', datetime.now(),
self.postgres_container.get_connection_url())
def tearDown(self) -> None:
self.postgres_container.stop()
def test_pull_table(self):
month_a = datetime.strptime('2020-02-01', '%Y-%m-%d')
month_b = datetime.strptime('2020-04-01', '%Y-%m-%d')
with CrifTable(
self.sample_table_name, 'customers', month_a,
self.postgres_container.get_connection_url()) as crif_table:
crif_table.pull_table_as_dicts()
assert len(crif_table.table) == 6
with CrifTable(
self.sample_table_name, 'customers', month_b,
self.postgres_container.get_connection_url()) as crif_table:
crif_table.pull_table_as_dicts()
assert len(crif_table.table) == 8
def test_docker_run_postgress():
postgres_container = PostgresContainer("user", "secret", version="9.5")
with postgres_container as postgres:
conn = psycopg2.connect(host=postgres.host_ip,
user=postgres.username,
password=postgres.password,
database=postgres.database)
cur = conn.cursor()
cur.execute("SELECT VERSION()")
row = cur.fetchone()
print("server version:", row[0])
cur.close()
assert len(row) > 0
def psql():
"""
We need to modify the Flask Application to accept a seperate
sessionmaker for testing. !TODO
"""
psql_container = PostgresContainer("postgres:9.6")
psql_container.POSTGRES_USER = "******"
psql_container.POSTGRES_PASSWORD = "******"
psql_container.POSTGRES_DB = "test"
psql_container.start()
engine = sqlalchemy.create_engine(psql_container.get_connection_url())
session = sessionmaker(bind=engine)()
# alert webservice to port for this database
port = psql_container.get_connection_url().split(":")[-1].split("/")[0]
header = copy.copy(TEST_HEADER)
header["port"] = port
return {"container": psql_container, "session": session, "header": header}
def main(sqs):
pg = None
try:
print("Starting postgres db...")
pg = PostgresContainer("postgres:11.6-alpine")
pg.start()
os.environ['DB_URL'] = pg.get_connection_url()
sqs.create_queue(QueueName="test-queue.fifo",
Attributes={"FifoQueue": "true"})
from app import main
yield main
finally:
if pg is not None:
pg.stop()
def test(debug, sbt_cmd, user, user_home):
print("=========================================")
print("CI environment detected [" + OS_TYPE + "]")
print("Running container user as " + user)
print("Test commencing")
print("=========================================")
try:
postgres = PostgresContainer("postgres:9.6")
postgres.POSTGRES_DB = environ.get("POSTGRES_DB")
postgres.POSTGRES_USER = environ.get("POSTGRES_USER")
postgres.POSTGRES_PASSWORD = environ.get("POSTGRES_PASSWORD")
with postgres.with_name("random_menu_selector_db") as db:
system("docker network create random_menu_selector_network")
system(
"docker network connect random_menu_selector_network random_menu_selector_db"
)
system("mkdir -p ${HOME}/.sbt && mkdir -p ${HOME}/.ivy2")
test = " && ".join([
"sbt compile", # Liquibase requires db jar dependency
"db/evolve.sh",
"sbt " + ["", "-jvm-debug 5005"][debug] + " \\\"" + sbt_cmd +
"\\\"",
])
system(" ".join([
"docker run",
"--name",
"random_menu_selector_build",
"--network",
"random_menu_selector_network",
"--env-file",
"${PWD}/test.env",
["", "-p 5005:5005"][debug],
"-u",
user,
"-w",
user_home + "/random_menu_selector",
"-v",
"/var/run/docker.sock:/var/run/docker.sock",
"-v",
"${HOME}/.ivy2:" + user_home + "/.ivy2",
"-v",
"${HOME}/.sbt:" + user_home + "/.sbt",
"-v",
"${PWD}/..:" + user_home + "/random_menu_selector",
"fsw0422/random_menu_selector/cicd:latest",
"bash -c \"{cmd}\"".format(cmd=test),
]))
finally:
system("docker rm -f random_menu_selector_build")
system("docker rm -f random_menu_selector_db")
system("printf 'y\n' | docker network prune")
system("printf 'y\n' | docker volume prune")
import sys
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate
from testcontainers.postgres import PostgresContainer
SQLALCHEMY_DATABASE_URI = None
# somewhat hacky setup to run a testcontainer when running tests
testing = "pytest" in sys.modules
postgres_container = None
def stop_postgres_container():
if postgres_container != None:
postgres_container.stop()
if testing:
postgres_container = PostgresContainer("ratestask-db:latest")
postgres_container.start()
SQLALCHEMY_DATABASE_URI = postgres_container.get_connection_url()
else:
SQLALCHEMY_DATABASE_URI = "postgresql+psycopg2://postgres:ratestask@localhost/postgres"
database = SQLAlchemy()
migrate = Migrate(db=database)
class CrossLanguageJdbcIOTest(unittest.TestCase):
def setUp(self):
self.postgres = PostgresContainer('postgres:latest')
self.postgres.start()
self.engine = sqlalchemy.create_engine(
self.postgres.get_connection_url())
self.username = '******'
self.password = '******'
self.host = self.postgres.get_container_host_ip()
self.port = self.postgres.get_exposed_port(5432)
self.database_name = 'test'
self.driver_class_name = 'org.postgresql.Driver'
self.jdbc_url = 'jdbc:postgresql://{}:{}/{}'.format(
self.host, self.port, self.database_name)
def tearDown(self):
self.postgres.stop()
def test_xlang_jdbc_write(self):
table_name = 'jdbc_external_test_write'
self.engine.execute(
"CREATE TABLE {}(f_id INTEGER, f_real REAL, f_string VARCHAR)".
format(table_name))
inserted_rows = [
JdbcWriteTestRow(i, i + 0.1, 'Test{}'.format(i))
for i in range(ROW_COUNT)
]
with TestPipeline() as p:
p.not_use_test_runner_api = True
_ = (p
|
beam.Create(inserted_rows).with_output_types(JdbcWriteTestRow)
| 'Write to jdbc' >> WriteToJdbc(
driver_class_name=self.driver_class_name,
jdbc_url=self.jdbc_url,
username=self.username,
password=self.password,
statement='INSERT INTO {} VALUES(?, ?, ?)'.format(
table_name),
))
fetched_data = self.engine.execute(
"SELECT * FROM {}".format(table_name))
fetched_rows = [
JdbcWriteTestRow(int(row[0]), float(row[1]), str(row[2]))
for row in fetched_data
]
self.assertEqual(
set(fetched_rows),
set(inserted_rows),
'Inserted data does not fit data fetched from table',
)
def test_xlang_jdbc_read(self):
table_name = 'jdbc_external_test_read'
self.engine.execute(
"CREATE TABLE {}(f_int INTEGER)".format(table_name))
for i in range(ROW_COUNT):
self.engine.execute("INSERT INTO {} VALUES({})".format(
table_name, i))
with TestPipeline() as p:
p.not_use_test_runner_api = True
result = (p
| 'Read from jdbc' >> ReadFromJdbc(
driver_class_name=self.driver_class_name,
jdbc_url=self.jdbc_url,
username=self.username,
password=self.password,
query='SELECT f_int FROM {}'.format(table_name),
))
assert_that(
result,
equal_to([JdbcReadTestRow(i) for i in range(ROW_COUNT)]))
class CrossLanguageJdbcIOTest(unittest.TestCase):
def setUp(self):
self.start_postgres_container(retries=3)
self.engine = sqlalchemy.create_engine(
self.postgres.get_connection_url())
self.username = '******'
self.password = '******'
self.host = self.postgres.get_container_host_ip()
self.port = self.postgres.get_exposed_port(5432)
self.database_name = 'test'
self.driver_class_name = 'org.postgresql.Driver'
self.jdbc_url = 'jdbc:postgresql://{}:{}/{}'.format(
self.host, self.port, self.database_name)
def tearDown(self):
# Sometimes stopping the container raises ReadTimeout. We can ignore it
# here to avoid the test failure.
try:
self.postgres.stop()
except: # pylint: disable=bare-except
logging.error('Could not stop the postgreSQL container.')
def test_xlang_jdbc_write(self):
table_name = 'jdbc_external_test_write'
self.engine.execute(
"CREATE TABLE {}(f_id INTEGER, f_real REAL, f_string VARCHAR)".
format(table_name))
inserted_rows = [
JdbcWriteTestRow(i, i + 0.1, 'Test{}'.format(i))
for i in range(ROW_COUNT)
]
with TestPipeline() as p:
p.not_use_test_runner_api = True
_ = (
p
|
beam.Create(inserted_rows).with_output_types(JdbcWriteTestRow)
# TODO(BEAM-10750) Add test with overridden write_statement
| 'Write to jdbc' >> WriteToJdbc(
table_name=table_name,
driver_class_name=self.driver_class_name,
jdbc_url=self.jdbc_url,
username=self.username,
password=self.password,
))
fetched_data = self.engine.execute(
"SELECT * FROM {}".format(table_name))
fetched_rows = [
JdbcWriteTestRow(int(row[0]), float(row[1]), str(row[2]))
for row in fetched_data
]
self.assertEqual(
set(fetched_rows),
set(inserted_rows),
'Inserted data does not fit data fetched from table',
)
def test_xlang_jdbc_read(self):
table_name = 'jdbc_external_test_read'
self.engine.execute(
"CREATE TABLE {}(f_int INTEGER)".format(table_name))
for i in range(ROW_COUNT):
self.engine.execute("INSERT INTO {} VALUES({})".format(
table_name, i))
with TestPipeline() as p:
p.not_use_test_runner_api = True
result = (
p
# TODO(BEAM-10750) Add test with overridden read_query
| 'Read from jdbc' >> ReadFromJdbc(
table_name=table_name,
driver_class_name=self.driver_class_name,
jdbc_url=self.jdbc_url,
username=self.username,
password=self.password,
))
assert_that(
result,
equal_to([JdbcReadTestRow(i) for i in range(ROW_COUNT)]))
# Creating a container with testcontainers sometimes raises ReadTimeout
# error. In java there are 2 retries set by default.
def start_postgres_container(self, retries):
for i in range(retries):
try:
self.postgres = PostgresContainer('postgres:12.3')
self.postgres.start()
break
except Exception as e: # pylint: disable=bare-except
if i == retries - 1:
logging.error('Unable to initialize postgreSQL container.')
raise e
class CrossLanguageJdbcIOTest(unittest.TestCase):
DbData = typing.NamedTuple('DbData', [('container_fn', typing.Any),
('classpath', typing.List[str]),
('db_string', str),
('connector', str)])
DB_CONTAINER_CLASSPATH_STRING = {
'postgres':
DbData(lambda: PostgresContainer('postgres:12.3'), None, 'postgresql',
'org.postgresql.Driver'),
'mysql':
DbData(lambda: MySqlContainer(), ['mysql:mysql-connector-java:8.0.28'],
'mysql', 'com.mysql.cj.jdbc.Driver')
}
def _setUpTestCase(self, container_init: Callable[[],
Union[PostgresContainer,
MySqlContainer]],
db_string: str, driver: str):
# This method is not the normal setUp from unittest, because the test has
# beem parameterized. The setup then needs extra parameters to initialize.
self.start_db_container(retries=3, container_init=container_init)
self.engine = sqlalchemy.create_engine(self.db.get_connection_url())
self.username = '******'
self.password = '******'
self.host = self.db.get_container_host_ip()
self.port = self.db.get_exposed_port(self.db.port_to_expose)
self.database_name = 'test'
self.driver_class_name = driver
self.jdbc_url = 'jdbc:{}://{}:{}/{}'.format(db_string, self.host,
self.port,
self.database_name)
def tearDown(self):
# Sometimes stopping the container raises ReadTimeout. We can ignore it
# here to avoid the test failure.
try:
self.db.stop()
except: # pylint: disable=bare-except
logging.error('Could not stop the postgreSQL container.')
@parameterized.expand(['postgres', 'mysql'])
def test_xlang_jdbc_write(self, database):
container_init, classpath, db_string, driver = (
CrossLanguageJdbcIOTest.DB_CONTAINER_CLASSPATH_STRING[database])
self._setUpTestCase(container_init, db_string, driver)
table_name = 'jdbc_external_test_write'
self.engine.execute(
"CREATE TABLE {}(f_id INTEGER, f_real FLOAT, f_string VARCHAR(100))"
.format(table_name))
inserted_rows = [
JdbcWriteTestRow(i, i + 0.1, 'Test{}'.format(i))
for i in range(ROW_COUNT)
]
with TestPipeline() as p:
p.not_use_test_runner_api = True
_ = (
p
|
beam.Create(inserted_rows).with_output_types(JdbcWriteTestRow)
# TODO(https://github.com/apache/beam/issues/20446) Add test with
# overridden write_statement
| 'Write to jdbc' >> WriteToJdbc(
table_name=table_name,
driver_class_name=self.driver_class_name,
jdbc_url=self.jdbc_url,
username=self.username,
password=self.password,
classpath=classpath,
))
fetched_data = self.engine.execute(
"SELECT * FROM {}".format(table_name))
fetched_rows = [
JdbcWriteTestRow(int(row[0]), float(row[1]), str(row[2]))
for row in fetched_data
]
self.assertEqual(
set(fetched_rows),
set(inserted_rows),
'Inserted data does not fit data fetched from table',
)
@parameterized.expand(['postgres', 'mysql'])
def test_xlang_jdbc_read(self, database):
container_init, classpath, db_string, driver = (
CrossLanguageJdbcIOTest.DB_CONTAINER_CLASSPATH_STRING[database])
self._setUpTestCase(container_init, db_string, driver)
table_name = 'jdbc_external_test_read'
self.engine.execute(
"CREATE TABLE {}(f_int INTEGER)".format(table_name))
for i in range(ROW_COUNT):
self.engine.execute("INSERT INTO {} VALUES({})".format(
table_name, i))
with TestPipeline() as p:
p.not_use_test_runner_api = True
result = (
p
# TODO(https://github.com/apache/beam/issues/20446) Add test with
# overridden read_query
| 'Read from jdbc' >> ReadFromJdbc(
table_name=table_name,
driver_class_name=self.driver_class_name,
jdbc_url=self.jdbc_url,
username=self.username,
password=self.password,
classpath=classpath))
assert_that(
result,
equal_to([JdbcReadTestRow(i) for i in range(ROW_COUNT)]))
# Creating a container with testcontainers sometimes raises ReadTimeout
# error. In java there are 2 retries set by default.
def start_db_container(self, retries, container_init):
for i in range(retries):
try:
self.db = container_init()
self.db.start()
break
except Exception as e: # pylint: disable=bare-except
if i == retries - 1:
logging.error('Unable to initialize database container.')
raise e
def test_docker_run_postgress(cls) -> None:
postgres_container = PostgresContainer("postgres:9.5")
with postgres_container as postgres:
e = sqlalchemy.create_engine(postgres.get_connection_url())
result = e.execute("SELECT version()")
def test_docker_run_postgres_with_driver_pg8000():
postgres_container = PostgresContainer("postgres:9.5", driver="pg8000")
with postgres_container as postgres:
e = sqlalchemy.create_engine(postgres.get_connection_url())
e.execute("select 1=1")
class CrossLanguageDebeziumIOTest(unittest.TestCase):
def setUp(self):
self.username = '******'
self.password = '******'
self.database = 'inventory'
self.start_db_container(retries=1)
self.host = self.db.get_container_host_ip()
self.port = self.db.get_exposed_port(5432)
self.connector_class = DriverClassName.POSTGRESQL
self.connection_properties = [
"database.dbname=inventory", "database.server.name=dbserver1",
"database.include.list=inventory", "include.schema.changes=false"
]
def tearDown(self):
# Sometimes stopping the container raises ReadTimeout. We can ignore it
# here to avoid the test failure.
try:
self.db.stop()
except: # pylint: disable=bare-except
logging.error('Could not stop the DB container.')
def test_xlang_debezium_read(self):
expected_response = [{
"metadata": {
"connector": "postgresql",
"version": "1.3.1.Final",
"name": "dbserver1",
"database": "inventory",
"schema": "inventory",
"table": "customers"
},
"before": None,
"after": {
"fields": {
"last_name": "Thomas",
"id": 1001,
"first_name": "Sally",
"email": "*****@*****.**"
}
}
}]
with TestPipeline() as p:
p.not_use_test_runner_api = True
results = (p
| 'Read from debezium' >> ReadFromDebezium(
username=self.username,
password=self.password,
host=self.host,
port=self.port,
max_number_of_records=NUM_RECORDS,
connector_class=self.connector_class,
connection_properties=self.connection_properties))
assert_that(results, equal_to(expected_response))
# Creating a container with testcontainers sometimes raises ReadTimeout
# error. In java there are 2 retries set by default.
def start_db_container(self, retries):
for i in range(retries):
try:
self.db = PostgresContainer('debezium/example-postgres:latest',
user=self.username,
password=self.password,
dbname=self.database)
self.db.start()
break
except Exception as e: # pylint: disable=bare-except
if i == retries - 1:
logging.error('Unable to initialize DB container.')
raise e
