Exemple #1
0
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
Exemple #2
0
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()
Exemple #3
0
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)]))