github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/external/xlang_debeziumio_it_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  import logging
    19  import unittest
    20  
    21  from apache_beam.io.debezium import DriverClassName
    22  from apache_beam.io.debezium import ReadFromDebezium
    23  from apache_beam.options.pipeline_options import StandardOptions
    24  from apache_beam.testing.test_pipeline import TestPipeline
    25  from apache_beam.testing.util import assert_that
    26  from apache_beam.testing.util import equal_to
    27  
    28  # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
    29  try:
    30    from testcontainers.postgres import PostgresContainer
    31  except ImportError:
    32    PostgresContainer = None
    33  
    34  NUM_RECORDS = 1
    35  
    36  
    37  @unittest.skipIf(
    38      PostgresContainer is None, 'testcontainers package is not installed')
    39  @unittest.skipIf(
    40      TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is
    41      None,
    42      'Do not run this test on precommit suites.')
    43  class CrossLanguageDebeziumIOTest(unittest.TestCase):
    44    def setUp(self):
    45      self.username = 'debezium'
    46      self.password = 'dbz'
    47      self.database = 'inventory'
    48      self.start_db_container(retries=1)
    49      self.host = self.db.get_container_host_ip()
    50      self.port = self.db.get_exposed_port(5432)
    51      self.connector_class = DriverClassName.POSTGRESQL
    52      self.connection_properties = [
    53          "database.dbname=inventory",
    54          "database.server.name=dbserver1",
    55          "database.include.list=inventory",
    56          "include.schema.changes=false"
    57      ]
    58  
    59    def tearDown(self):
    60      # Sometimes stopping the container raises ReadTimeout. We can ignore it
    61      # here to avoid the test failure.
    62      try:
    63        self.db.stop()
    64      except:  # pylint: disable=bare-except
    65        logging.error('Could not stop the DB container.')
    66  
    67    def test_xlang_debezium_read(self):
    68      expected_response = [{
    69          "metadata": {
    70              "connector": "postgresql",
    71              "version": "1.3.1.Final",
    72              "name": "dbserver1",
    73              "database": "inventory",
    74              "schema": "inventory",
    75              "table": "customers"
    76          },
    77          "before": None,
    78          "after": {
    79              "fields": {
    80                  "last_name": "Thomas",
    81                  "id": 1001,
    82                  "first_name": "Sally",
    83                  "email": "sally.thomas@acme.com"
    84              }
    85          }
    86      }]
    87  
    88      with TestPipeline() as p:
    89        p.not_use_test_runner_api = True
    90        results = (
    91            p
    92            | 'Read from debezium' >> ReadFromDebezium(
    93                username=self.username,
    94                password=self.password,
    95                host=self.host,
    96                port=self.port,
    97                max_number_of_records=NUM_RECORDS,
    98                connector_class=self.connector_class,
    99                connection_properties=self.connection_properties))
   100        assert_that(results, equal_to(expected_response))
   101  
   102  
   103  # Creating a container with testcontainers sometimes raises ReadTimeout
   104  # error. In java there are 2 retries set by default.
   105  
   106    def start_db_container(self, retries):
   107      for i in range(retries):
   108        try:
   109          self.db = PostgresContainer(
   110              'debezium/example-postgres:latest',
   111              user=self.username,
   112              password=self.password,
   113              dbname=self.database)
   114          self.db.start()
   115          break
   116        except Exception as e:  # pylint: disable=bare-except
   117          if i == retries - 1:
   118            logging.error('Unable to initialize DB container.')
   119            raise e
   120  
   121  if __name__ == '__main__':
   122    logging.getLogger().setLevel(logging.INFO)
   123    unittest.main()