sdks/python/apache_beam/yaml/integration_tests.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Runs integration tests in the tests directory."""

import contextlib
import copy
import glob
import itertools
import json
import logging
import os
import random
import secrets
import shutil
import sqlite3
import string
import struct
import tempfile
import time
import unittest
import uuid
from datetime import datetime
from datetime import timezone

import mock
import requests

from apache_beam.coders import Coder
from apache_beam.coders.coder_impl import CoderImpl


class BigEndianIntegerCoderImpl(CoderImpl):
  """Coder implementation for big-endian integers used in cross-language tests.

  This is needed because Java's BigEndianIntegerCoder falls back to the generic
  'beam:coders:javasdk:0.1' URN when used in cross-language pipelines, and
  Python's FnApiRunner needs to know how to decode it.
  """
  def encode_to_stream(self, value, stream, nested):
    stream.write(struct.pack('>i', value))

  def decode_from_stream(self, stream, nested):
    return struct.unpack('>i', stream.read(4))[0]


class BigEndianIntegerCoder(Coder):
  def get_impl(self):
    return BigEndianIntegerCoderImpl()


# Register the coder with the fallback URN used by the Java SDK for this coder.
# This allows the Python FnApiRunner to handle data sharded by Java transforms
# using BigEndianIntegerCoder in integration tests.
Coder.register_urn(
    'beam:coders:javasdk:0.1',
    None, lambda payload, components, context: BigEndianIntegerCoder())

import psycopg2
import pyarrow as pa
import pyarrow.parquet as pq
import pytds
import sqlalchemy
import yaml
from apitools.base.py.exceptions import HttpError
from google.cloud import pubsub_v1
from google.cloud.bigtable import client
from google.cloud.bigtable_admin_v2.types import instance

try:
  from google.cloud import firestore
except ImportError:
  firestore = None
try:
  from google.cloud import secretmanager
except ImportError:
  secretmanager = None
from testcontainers.core.container import DockerContainer
from testcontainers.core.waiting_utils import wait_for_logs
from testcontainers.google import PubSubContainer
from testcontainers.kafka import KafkaContainer
from testcontainers.mongodb import MongoDbContainer
from testcontainers.mssql import SqlServerContainer
from testcontainers.mysql import MySqlContainer
from testcontainers.postgres import PostgresContainer

import apache_beam as beam
from apache_beam.io import filesystems
from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
from apache_beam.io.gcp.bigquery_tools import get_table_schema_from_string
from apache_beam.io.gcp.internal.clients import bigquery
from apache_beam.io.gcp.spanner_wrapper import SpannerWrapper
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.utils import python_callable
from apache_beam.utils import retry
from apache_beam.yaml import yaml_provider
from apache_beam.yaml import yaml_transform
from apache_beam.yaml.conftest import yaml_test_files_dir

_LOGGER = logging.getLogger(__name__)

_MONGO_CONTAINER_IMAGE = 'mongo:7.0.7'


@contextlib.contextmanager
def gcs_temp_dir(bucket):
  """Context manager to create and clean up a temporary GCS directory.

  Creates a unique temporary directory within the specified GCS bucket
  and yields the path. Upon exiting the context, the directory and its
  contents are deleted.

  Args:
    bucket (str): The GCS bucket name (e.g., 'gs://my-bucket').

  Yields:
    str: The full path to the created temporary GCS directory.
         Example: 'gs://my-bucket/yaml-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx'
  """
  gcs_tempdir = bucket + '/yaml-' + str(uuid.uuid4())
  yield gcs_tempdir
  filesystems.FileSystems.delete([gcs_tempdir])


@contextlib.contextmanager
def temp_spanner_table(project, prefix='temp_spanner_db_'):
  """Context manager to create and clean up a temporary Spanner database and
  table.

  Creates a unique temporary Spanner database within the specified project
  and a predefined table named 'tmp_table' with columns ['UserId', 'Key'].
  It yields connection details for the created resources. Upon exiting the
  context, the temporary database (and its table) is deleted.

  Args:
    project (str): The Google Cloud project ID.
    prefix (str): A prefix to use for the temporary database name.
      Defaults to 'temp_spanner_db_'.

  Yields:
    list[str]: A list containing connection details:
      [project_id, instance_id, database_id, table_name, list_of_columns].
      Example: ['my-project', 'beam-test', 'temp_spanner_db_...', 'tmp_table',
        ['UserId', 'Key']]
  """
  spanner_client = SpannerWrapper(project)
  spanner_client._create_database()
  instance = "beam-test"
  database = spanner_client._test_database
  table = 'tmp_table'
  columns = ['UserId', 'Key']
  logging.info("Created Spanner database: %s", database)
  try:
    yield [f'{project}', f'{instance}', f'{database}', f'{table}', columns]
  finally:
    logging.info("Deleting Spanner database: %s", database)
    spanner_client._delete_database()


@contextlib.contextmanager
def temp_firestore_collection(
    project='apache-beam-testing',
    database='firestoredb',
    prefix='yaml_firestore_it_'):
  """Context manager for an isolated Firestore collection used in YAML ITs.

  Uses the shared Beam test project and the ``firestoredb`` database, matching
  the Java Firestore integration tests.

  Args:
    project (str): GCP project id.
    database (str): Firestore database id.
    prefix (str): Prefix for the temporary collection name.

  Yields:
    dict: Keys ``PROJECT``, ``DATABASE``, and ``COLLECTION``.
  """
  if firestore is None:
    raise RuntimeError("google-cloud-firestore is not installed.")

  client = firestore.Client(project=project, database=database)
  collection_id = f'{prefix}{uuid.uuid4().hex}'
  logging.info(
      'Using Firestore collection %s in project %s database %s',
      collection_id,
      project,
      database)
  try:
    yield {
        'PROJECT': project,
        'DATABASE': database,
        'COLLECTION': collection_id,
    }
  finally:
    logging.info('Deleting documents in Firestore collection %s', collection_id)
    collection_ref = client.collection(collection_id)
    batch = client.batch()
    pending = 0
    for doc in collection_ref.stream():
      batch.delete(doc.reference)
      pending += 1
      if pending >= 400:
        batch.commit()
        batch = client.batch()
        pending = 0
    if pending:
      batch.commit()


class AssertBigQueryQueryEqual(beam.PTransform):
  """Asserts that a BigQuery SQL query returns exactly the expected rows.

  Used by CDC YAML ITs: Storage Read (ReadFromBigQuery table:) does not see
  unapplied CDC UPSERTs, while a query job applies them (same approach as
  xlang_bigqueryio_it_test.BigqueryFullResultMatcher).
  """
  def __init__(self, project, query, expected):
    self._project = project
    self._query = query
    self._expected = expected

  def expand(self, pcoll):
    project = self._project
    query = self._query
    expected = [dict(row) for row in self._expected]

    def _verify(_):
      # Local import: google-cloud-bigquery is a gcp_test extra.
      from google.cloud import bigquery as gcbq
      client = gcbq.Client(project=project)
      actual = [dict(row.items()) for row in client.query(query).result()]

      # Compare as sorted JSON for stable, type-tolerant equality.
      def _key(row):
        return json.dumps(row, sort_keys=True, default=str)

      if sorted(actual, key=_key) != sorted(expected, key=_key):
        raise AssertionError(
            'BigQuery query %r returned %r, expected %r' %
            (query, actual, expected))
      return _

    return pcoll | beam.Map(_verify)


@contextlib.contextmanager
def temp_bigquery_table(
    project,
    prefix='yaml_bq_it_',
    use_colon_table_spec=False,
    table_schema=None,
    primary_key=None):
  """Context manager to create and clean up a temporary BigQuery dataset.

  Creates a unique temporary BigQuery dataset within the specified project.
  It yields a placeholder table name string within that dataset (e.g.,
  'project.dataset_id.tmp_table'). The actual table is expected to be
  created by the test using this context, unless ``table_schema`` is set.

  Upon exiting the context, the temporary dataset and all its contents
  (including any tables created within it) are deleted.

  Args:
    project (str): The Google Cloud project ID.
    prefix (str): A prefix to use for the temporary dataset name.
      Defaults to 'yaml_bq_it_'.
    use_colon_table_spec (bool): If True, yield ``project:dataset.tmp_table``.
      Defaults to False (``project.dataset.tmp_table``).
    table_schema (str|None): Optional BigQuery schema string
      (``name:TYPE,...``). When set, creates ``tmp_table`` before yield.
    primary_key (list|None): Optional primary-key column names. Requires
      ``table_schema``. Primary-key columns are created as REQUIRED.

  Yields:
    str: The full path for a temporary BigQuery table within the created
    dataset.
         Example: 'my-project.yaml_bq_it_a1b2c3d4e5f6...tmp_table'
  """
  if primary_key and not table_schema:
    raise ValueError('primary_key requires table_schema')
  options = PipelineOptions()
  options.view_as(GoogleCloudOptions).project = project
  bigquery_client = BigQueryWrapper.from_pipeline_options(options)
  dataset_id = '%s_%s' % (prefix, uuid.uuid4().hex)
  table_id = 'tmp_table'
  bigquery_client.get_or_create_dataset(project, dataset_id)
  logging.info("Created dataset %s in project %s", dataset_id, project)
  if table_schema:
    schema = get_table_schema_from_string(table_schema)
    additional_create_parameters = None
    if primary_key:
      pk_columns = set(primary_key)
      for field in schema.fields:
        if field.name in pk_columns:
          field.mode = 'REQUIRED'
      additional_create_parameters = {
          'tableConstraints': bigquery.TableConstraints(
              primaryKey=bigquery.TableConstraints.PrimaryKeyValue(
                  columns=list(primary_key)))
      }
    bigquery_client.get_or_create_table(
        project_id=project,
        dataset_id=dataset_id,
        table_id=table_id,
        schema=schema,
        create_disposition='CREATE_IF_NEEDED',
        write_disposition='WRITE_APPEND',
        additional_create_parameters=additional_create_parameters)
    logging.info(
        "Created table %s.%s.%s (primary_key=%s)",
        project,
        dataset_id,
        table_id,
        primary_key)
  if use_colon_table_spec:
    yield f'{project}:{dataset_id}.{table_id}'
  else:
    yield f'{project}.{dataset_id}.{table_id}'
  request = bigquery.BigqueryDatasetsDeleteRequest(
      projectId=project, datasetId=dataset_id, deleteContents=True)
  logging.info("Deleting dataset %s in project %s", dataset_id, project)
  bigquery_client.client.datasets.Delete(request)


def instance_prefix(instance):
  datestr = "".join(filter(str.isdigit, str(datetime.now(timezone.utc).date())))
  instance_id = '%s-%s-%s' % (instance, datestr, secrets.token_hex(4))
  assert len(instance_id) < 34, "instance id length needs to be within [6, 33]"
  return instance_id


@contextlib.contextmanager
def temp_bigtable_table(project, prefix='yaml_bt_it_'):
  INSTANCE = "bt-write-tests"
  TABLE_ID = "test-table"

  instance_id = instance_prefix(INSTANCE)

  clientT = client.Client(admin=True, project=project)
  # create cluster and instance
  instanceT = clientT.instance(
      instance_id,
      display_name=INSTANCE,
      instance_type=instance.Instance.Type.DEVELOPMENT)
  cluster = instanceT.cluster("test-cluster", "us-central1-a")
  operation = instanceT.create(clusters=[cluster])
  operation.result(timeout=500)
  _LOGGER.info("Created instance [%s] in project [%s]", instance_id, project)

  # create table inside instance
  table = instanceT.table(TABLE_ID)
  table.create()
  _LOGGER.info("Created table [%s]", table.table_id)
  # in the table that is created, make a new family called cf1
  col_fam = table.column_family('cf1')
  col_fam.create()

  # another family called cf2
  col_fam = table.column_family('cf2')
  col_fam.create()

  #yielding the tmp table for all the bigTable tests
  yield instance_id

  #try catch for deleting table and instance after all tests are ran
  try:
    _LOGGER.info("Deleting table [%s]", table.table_id)
    table.delete()
    instanceT.delete()
  except HttpError:
    _LOGGER.warning("Failed to clean up instance")


@contextlib.contextmanager
def temp_mongodb_table():
  """
  provides a temporary MongoDB instance.

  starts a MongoDB container, creates a unique database
  and collection name for test isolation, and yields them as a dictionary.

  This allows YAML test files to get connection details without hardcoding them.
  Example usage in a YAML test file's fixture section:

  fixtures:
    - name: mongo_vars
      type: path.to.this.file.mongodb_fixture

  Then, in the pipeline definition, you can use placeholders like:
  - uri: ${mongo_vars.URI}
  - database: ${mongo_vars.DATABASE}
  - collection: ${mongo_vars.COLLECTION}
  """
  _LOGGER.info("Setting up MongoDB fixture...")
  mongo_container = MongoDbContainer(_MONGO_CONTAINER_IMAGE)
  try:
    mongo_container.start()
    mongo_uri = mongo_container.get_connection_url()

    # MongoDbContainer's entrypoint restarts mongod after init scripts; wait
    # until the server is stably accepting connections and responds to ping.
    @retry.with_exponential_backoff(num_retries=5, initial_delay_secs=1)
    def _wait_for_mongo():
      mongo_client = mongo_container.get_connection_client()
      mongo_client.admin.command('ping')

    _wait_for_mongo()

    db_name = f'db_{uuid.uuid4().hex}'
    collection_name = f'collection_{uuid.uuid4().hex}'

    _LOGGER.info(
        "MongoDB container started. URI: [%s], DB: [%s], Collection: [%s]",
        mongo_uri,
        db_name,
        collection_name)

    yield {
        'URI': mongo_uri,
        'DATABASE': db_name,
        'COLLECTION': collection_name,
    }

  finally:
    _LOGGER.info("Tearing down MongoDB fixture...")
    mongo_container.stop()
    _LOGGER.info("MongoDB container stopped.")


@contextlib.contextmanager
def temp_sqlite_database(prefix='yaml_jdbc_it_'):
  """Context manager to provide a temporary SQLite database via JDBC for
  testing.

  This function creates a temporary SQLite database file on the local
  filesystem. It establishes a connection using 'sqlite3', creates a predefined
  'tmp_table', and then yields a JDBC connection string suitable for use in
  tests that require a generic JDBC connection (specifically configured for
  SQLite in this case).

  The SQLite database file is automatically cleaned up (closed and deleted)
  when the context manager exits.

  Args:
      prefix (str): A prefix to use for the temporary database file name.

  Yields:
      str: A JDBC connection string for the temporary SQLite database.
           Example format: "jdbc:sqlite:<path_to_db_file>"

  Raises:
      sqlite3.Error: If there's an error connecting to or interacting with
                     the SQLite database during setup.
      Exception: Any other exception encountered during the setup or cleanup
                 process.
  """
  conn = cursor = None
  try:
    # Establish connection to the temp file
    db_name = f'{prefix}{uuid.uuid4().hex}.db'
    conn = sqlite3.connect(db_name)
    cursor = conn.cursor()

    # Create a temp table for tests
    cursor.execute(
        '''
      CREATE TABLE tmp_table (
        value INTEGER PRIMARY KEY,
        rank INTEGER
      )
    ''')
    conn.commit()
    yield f'jdbc:sqlite:{db_name}'
  except (sqlite3.Error, Exception) as err:
    logging.error("Error interacting with temporary SQLite DB: %s", err)
    raise err
  finally:
    # Close connections
    if cursor:
      cursor.close()
    if conn:
      conn.close()
    try:
      if os.path.exists(db_name):
        os.remove(db_name)
    except Exception as err:
      logging.error("Error deleting temporary SQLite DB: %s", err)
      raise err


@contextlib.contextmanager
def temp_mysql_database():
  """Context manager to provide a temporary MySQL database for testing.

  This function utilizes the 'testcontainers' library to spin up a
  MySQL instance within a Docker container. It then connects
  to this temporary database using 'mysql.connector', creates a predefined
  'tmp_table', and yields a JDBC connection string suitable for use in tests.

  The Docker container and the database instance are automatically managed
  and torn down when the context manager exits.

  Yields:
      str: A JDBC connection string for the temporary MySQL database.
           Example format:
           "jdbc:mysql://<host>:<port>/<db_name>?
              user=<user>&password=<password>"

  Raises:
      mysql.connector.Error: If there's an error connecting to or interacting
                             with the MySQL database during setup.
      Exception: Any other exception encountered during the setup process.
  """
  with MySqlContainer(init=True, dialect='pymysql') as mysql_container:
    # Make connection to temp database and create tmp table
    engine = sqlalchemy.create_engine(mysql_container.get_connection_url())
    with engine.begin() as connection:
      connection.execute(
          sqlalchemy.text(
              "CREATE TABLE tmp_table (value INTEGER, `rank` INTEGER);"))

    # Construct the JDBC url for connections later on by tests
    jdbc_url = (
        f"jdbc:mysql://{mysql_container.get_container_host_ip()}:"
        f"{mysql_container.get_exposed_port(mysql_container.port)}/"
        f"{mysql_container.dbname}?"
        f"user={mysql_container.username}&"
        f"password={mysql_container.password}")

    yield jdbc_url


@contextlib.contextmanager
def snowflake_fixture():
  options = PipelineOptions().get_all_options()

  required = [
      'server_name',
      'username',
      'password',
      'staging_bucket_name',
      'storage_integration_name',
      'database',
      'schema',
      'table',
  ]

  missing = [name for name in required if not options.get(name)]
  if missing:
    raise unittest.SkipTest(
        'Snowflake YAML integration test requires external configuration: ' +
        ', '.join(missing))

  yield {
      'SERVER_NAME': options['server_name'],
      'USERNAME': options['username'],
      'PASSWORD': options['password'],
      'STAGING_BUCKET_NAME': options['staging_bucket_name'],
      'STORAGE_INTEGRATION_NAME': options['storage_integration_name'],
      'DATABASE': options['database'],
      'SCHEMA': options['schema'],
      'TABLE': options['table'],
      'WAREHOUSE': options.get('warehouse'),
      'ROLE': options.get('role'),
  }


@contextlib.contextmanager
def temp_debezium_postgres_database():
  """Provides a temporary PostgreSQL database configured for Debezium CDC."""

  container = (
      DockerContainer('quay.io/debezium/example-postgres:latest').with_env(
          'POSTGRES_USER',
          'debezium').with_env('POSTGRES_PASSWORD', 'dbz').with_env(
              'POSTGRES_DB', 'inventory').with_exposed_ports(5432))

  try:
    container.start()
    wait_for_logs(container, 'database system is ready to accept connections')

    host = container.get_container_host_ip()
    port = int(container.get_exposed_port(5432))

    connection = None
    for _ in range(30):
      try:
        connection = psycopg2.connect(
            host=host,
            port=port,
            user='debezium',
            password='dbz',
            dbname='inventory')
        break
      except psycopg2.OperationalError:
        time.sleep(1)

    if connection is None:
      raise RuntimeError(
          'Debezium PostgreSQL container failed to become ready.')

    try:
      with connection.cursor() as cursor:
        cursor.execute(
            """
          CREATE TABLE IF NOT EXISTS customers (
            id INTEGER PRIMARY KEY,
            name VARCHAR(255)
          )
          """)
        cursor.execute(
            """
          INSERT INTO customers (id, name)
          VALUES
            (1, 'Alice'),
            (2, 'Bob')
          ON CONFLICT (id) DO NOTHING
          """)
      connection.commit()
    finally:
      connection.close()

    yield {
        'HOST': host,
        'PORT': port,
        'USERNAME': 'debezium',
        'PASSWORD': 'dbz',
        'DATABASE': 'inventory',
        'TABLE': 'public.customers',
    }

  finally:
    container.stop()


@contextlib.contextmanager
def temp_postgres_database():
  """Context manager to provide a temporary PostgreSQL database for testing.

  This function utilizes the 'testcontainers' library to spin up a
  PostgreSQL instance within a Docker container. It then connects
  to this temporary database using 'psycopg2', creates a predefined 'tmp_table',
  and yields a JDBC connection string suitable for use in tests.

  The Docker container and the database instance are automatically managed
  and torn down when the context manager exits.

  Yields:
      str: A JDBC connection string for the temporary PostgreSQL database.
           Example format:
           "jdbc:postgresql://<host>:<port>/<db_name>?
              user=<user>&password=<password>"

  Raises:
      psycopg2.Error: If there's an error connecting to or interacting with
                      the PostgreSQL database during setup.
      Exception: Any other exception encountered during the setup process.
  """
  default_port = 5432

  # Start the postgress container using testcontainers
  with PostgresContainer(port=default_port) as postgres_container:
    try:
      # Make connection to temp database and create tmp table
      engine = sqlalchemy.create_engine(postgres_container.get_connection_url())
      with engine.begin() as connection:
        connection.execute(
            sqlalchemy.text(
                "CREATE TABLE tmp_table (value INTEGER, rank INTEGER);"))

      # Construct the JDBC url for connections later on by tests
      jdbc_url = (
          f"jdbc:postgresql://{postgres_container.get_container_host_ip()}:"
          f"{postgres_container.get_exposed_port(default_port)}/"
          f"{postgres_container.dbname}?"
          f"user={postgres_container.username}&"
          f"password={postgres_container.password}")

      yield jdbc_url
    except (psycopg2.Error, Exception) as err:
      logging.error("Error interacting with temporary Postgres DB: %s", err)
      raise err


@contextlib.contextmanager
def temp_postgres_database_with_secret_manager(
    project=None, prefix='yaml_jdbc_sm_it_'):
  """Context manager to provide a temporary PostgreSQL database authenticated
  via GCP Secret Manager for testing.

  This function utilizes the 'testcontainers' library to spin up a
  PostgreSQL instance within a Docker container, creates a predefined 'tmp_table',
  registers the database container password in GCP Secret Manager, and yields
  a dictionary containing the JDBC connection URL, username, secret specification,
  and secret manager identifier.

  The Docker container, database instance, and GCP secret are automatically
  managed and torn down when the context manager exits.

  Args:
      project (str): Google Cloud Project ID. If not provided, reads from
        the GOOGLE_CLOUD_PROJECT environment variable or defaults to
        'apache-beam-testing'.
      prefix (str): Prefix to use for the temporary GCP secret name.

  Yields:
      dict: A dictionary containing connection and secret details:
            {
                'URL': 'jdbc:postgresql://<host>:<port>/<dbname>',
                'USERNAME': '<username>',
                'PASSWORD_SPEC': '{"name": "<secret_id>", "project": "<project_id>"}',
                'SECRET_MANAGER': 'googlecloudsecretmanager',
            }
  """
  if secretmanager is None:
    raise RuntimeError("google-cloud-secret-manager is not installed.")

  project_id = project or os.environ.get(
      'GOOGLE_CLOUD_PROJECT', 'apache-beam-testing')
  secret_client = secretmanager.SecretManagerServiceClient()

  default_port = 5432
  with PostgresContainer(port=default_port) as postgres_container:
    secret_postfix = (
        datetime.now(timezone.utc).strftime('%m%d_%H%M%S') + '_' +
        uuid.uuid4().hex[:6])
    secret_id = f'{prefix}{secret_postfix}'
    project_path = f'projects/{project_id}'
    secret_path = f'{project_path}/secrets/{secret_id}'

    _LOGGER.info("Creating GCP secret %s in project %s", secret_id, project_id)
    try:
      secret_client.get_secret(request={'name': secret_path})
    except Exception:
      secret_client.create_secret(
          request={
              'parent': project_path,
              'secret_id': secret_id,
              'secret': {
                  'replication': {
                      'automatic': {}
                  }
              }
          })

    secret_client.add_secret_version(
        request={
            'parent': secret_path,
            'payload': {
                'data': postgres_container.password.encode('utf-8')
            }
        })

    try:
      # Make connection to temp database and create tmp table
      engine = sqlalchemy.create_engine(postgres_container.get_connection_url())
      with engine.begin() as connection:
        connection.execute(
            sqlalchemy.text(
                "CREATE TABLE tmp_table (value INTEGER, rank INTEGER);"))

      # Construct the JDBC url for connections
      jdbc_url = (
          f"jdbc:postgresql://{postgres_container.get_container_host_ip()}:"
          f"{postgres_container.get_exposed_port(default_port)}/"
          f"{postgres_container.dbname}")

      secret_password_spec = json.dumps({
          'name': secret_id,
          'project': project_id,
      })

      yield {
          'URL': jdbc_url,
          'USERNAME': postgres_container.username,
          'PASSWORD_SPEC': secret_password_spec,
          'SECRET_MANAGER': 'googlecloudsecretmanager',
      }
    except (psycopg2.Error, Exception) as err:
      logging.error(
          "Error interacting with temporary Postgres DB with secret manager: %s",
          err)
      raise err
    finally:
      try:
        _LOGGER.info("Deleting GCP secret: %s", secret_path)
        secret_client.delete_secret(request={'name': secret_path})
      except Exception as err:
        _LOGGER.warning("Could not delete GCP secret %s: %s", secret_path, err)


@contextlib.contextmanager
def cloudsql_postgres_fixture():
  """Context manager to provide a PostgreSQL testcontainer database for testing."""
  default_port = 5432
  with PostgresContainer(port=default_port) as postgres_container:
    try:
      engine = sqlalchemy.create_engine(postgres_container.get_connection_url())
      with engine.begin() as connection:
        connection.execute(
            sqlalchemy.text(
                "CREATE TABLE tmp_table (id INTEGER, name VARCHAR(255), score FLOAT);"
            ))

      jdbc_url = (
          f"jdbc:postgresql://{postgres_container.get_container_host_ip()}:"
          f"{postgres_container.get_exposed_port(default_port)}/"
          f"{postgres_container.dbname}?"
          f"user={postgres_container.username}&"
          f"password={postgres_container.password}")

      yield {
          'JDBC_URL': jdbc_url,
          'DRIVER_CLASS_NAME': 'org.postgresql.Driver',
          'USERNAME': postgres_container.username,
          'PASSWORD': postgres_container.password,
          'DATABASE': postgres_container.dbname,
          'TABLE': 'tmp_table',
      }
    except (psycopg2.Error, Exception) as err:
      logging.error(
          "Error interacting with temporary Postgres DB in cloudsql_postgres_fixture: %s",
          err)
      raise err


@contextlib.contextmanager
def temp_sqlserver_database():
  """Context manager to provide a temporary SQL Server database for testing.

  This function utilizes the 'testcontainers' library to spin up a
  Microsoft SQL Server instance within a Docker container. It then connects
  to this temporary database using 'pytds', creates a predefined 'tmp_table',
  and yields a JDBC connection string suitable for use in tests.

  The Docker container and the database instance are automatically managed
  and torn down when the context manager exits.

  Yields:
      str: A JDBC connection string for the temporary SQL Server database.
           Example format:
           "jdbc:sqlserver://<host>:<port>;
              databaseName=<db_name>;
              user=<user>;
              password=<password>;
              encrypt=false;
              trustServerCertificate=true"

  Raises:
      pytds.Error: If there's an error connecting to or interacting with
                     the SQL Server database during setup.
      Exception: Any other exception encountered during the setup process.
  """
  default_port = 1433

  # Start the sql server using testcontainers
  with SqlServerContainer(port=default_port,
                          dialect='mssql+pytds') as sqlserver_container:
    try:
      # Make connection to temp database and create tmp table
      engine = sqlalchemy.create_engine(
          sqlserver_container.get_connection_url())
      with engine.begin() as connection:
        connection.execute(
            sqlalchemy.text(
                "CREATE TABLE tmp_table (value INTEGER, rank INTEGER);"))

      # Construct the JDBC url for connections later on by tests
      # NOTE: encrypt=false and trustServerCertificate=true is generally
      # needed for test container connections without proper certificates setup
      jdbc_url = (
          f"jdbc:sqlserver://{sqlserver_container.get_container_host_ip()}:"
          f"{int(sqlserver_container.get_exposed_port(default_port))};"
          f"databaseName={sqlserver_container.dbname};"
          f"user={sqlserver_container.username};"
          f"password={sqlserver_container.password};"
          f"encrypt=true;"
          f"trustServerCertificate=true")

      yield jdbc_url
    except (pytds.Error, Exception) as err:
      logging.error("Error interacting with temporary SQL Server DB: %s", err)
      raise err


class OracleTestContainer(DockerContainer):
  """
  OracleTestContainer is an updated version of OracleDBContainer that goes
  ahead and sets the oracle password, waits for logs to establish that the
  container is ready before calling get_exposed_port, and uses a more modern
  oracle driver.
  """
  def __init__(self):
    super().__init__("gvenzl/oracle-xe:21-slim")
    self.with_env("ORACLE_PASSWORD", "oracle")
    self.with_exposed_ports(1521)

  def start(self):
    super().start()
    wait_for_logs(self, "DATABASE IS READY TO USE!", timeout=300)
    return self

  def get_connection_url(self):
    port = self.get_exposed_port(1521)
    return \
      f"oracle+oracledb://system:oracle@localhost:{port}/?service_name=XEPDB1"


@contextlib.contextmanager
def temp_oracle_database():
  """Context manager to provide a temporary Oracle database for testing.

  This function utilizes the 'testcontainers' library to spin up an
  Oracle Database instance within a Docker container. It then connects
  to this temporary database using 'oracledb', creates a predefined

  NOTE: A custom OracleTestContainer class was created due to the current
  version (OracleDBContainer) that calls get_exposed_port too soon causing the
  service to hang until timeout.

  Yields:
      str: A JDBC connection string for the temporary Oracle database.
           Example format:
           "jdbc:oracle:thin:system/oracle@localhost:{port}/XEPDB1"

  Raises:
      oracledb.Error: If there's an error connecting to or interacting with
                      the Oracle database during setup.
      Exception: Any other exception encountered during the setup process.
  """
  with OracleTestContainer() as oracle:
    engine = sqlalchemy.create_engine(oracle.get_connection_url())
    with engine.connect() as connection:
      connection.execute(
          sqlalchemy.text(
              """
                CREATE TABLE tmp_table (
                    value NUMBER PRIMARY KEY,
                    rank NUMBER
                )
            """))
      connection.commit()
    port = oracle.get_exposed_port(1521)
    yield f"jdbc:oracle:thin:system/oracle@localhost:{port}/XEPDB1"


@contextlib.contextmanager
def temp_iceberg_table_with_pk(table_data):

  # Create a temp dir that will be shared between host and container.
  # We use the exact same path on both to avoid path mapping issues.
  # We create it in the current working directory (workspace) because
  # Docker in GitHub Actions often cannot mount directories from /tmp.
  temp_dir = tempfile.mkdtemp(dir=os.getcwd())
  os.chmod(temp_dir, 0o777)

  # Start the Iceberg REST catalog container
  container = DockerContainer("tabulario/iceberg-rest:0.6.0")
  container.with_exposed_ports(8181)
  container.with_volume_mapping(temp_dir, temp_dir, mode='rw')
  container.with_env("HADOOP_USER_NAME", "iceberg")
  container.with_env("CATALOG_WAREHOUSE", temp_dir)
  container.with_env(
      "CATALOG_IO__IMPL", "org.apache.iceberg.hadoop.HadoopFileIO")

  try:
    container.start()

    ip = container.get_container_host_ip()
    port = container.get_exposed_port(8181)
    api_url = f"http://{ip}:{port}"

    # Poll the REST API until it is ready
    for _ in range(30):
      try:
        response = requests.get(f"{api_url}/v1/config", timeout=5)
        if response.status_code == 200:
          break
      except (requests.exceptions.ConnectionError, requests.exceptions.Timeout):
        pass
      time.sleep(1)
    else:
      raise RuntimeError("Iceberg REST catalog failed to start in time.")

    # Create namespace 'db'
    requests.post(
        f"{api_url}/v1/namespaces",
        json={"namespace": ["db"]},
        headers={"Content-Type": "application/json"},
        timeout=10)

    # Create table with primary key
    response = requests.post(
        f"{api_url}/v1/namespaces/db/tables",
        json=table_data,
        headers={"Content-Type": "application/json"},
        timeout=10)
    if response.status_code != 200:
      raise RuntimeError(f"Failed to create Iceberg table: {response.text}")

    # Change permissions of the created directories inside the container
    # so the host user can write to them.
    container.get_wrapped_container().exec_run(f"chmod -R 777 {temp_dir}")

    yield {
        "api_url": api_url,
        "temp_dir": temp_dir,
        "table": f"db.{table_data['name']}"
    }

  finally:
    container.stop()
    shutil.rmtree(temp_dir, ignore_errors=True)


_LOCALSTACK_KINESIS_IMAGE = 'localstack/localstack:3.8.1'
_LOCALSTACK_KINESIS_ACCESS_KEY = 'accesskey'
_LOCALSTACK_KINESIS_SECRET_KEY = 'secretkey'
_LOCALSTACK_KINESIS_REGION = 'us-east-1'


def _create_kinesis_stream(kinesis_client, stream_name):
  """Creates a Kinesis stream and waits until it is active."""
  retries = 10
  for attempt in range(retries):
    try:
      kinesis_client.create_stream(StreamName=stream_name, ShardCount=1)
      time.sleep(2)
      break
    except Exception as exn:
      if attempt == retries - 1:
        _LOGGER.error('Could not create kinesis stream')
        raise exn

  for attempt in range(retries):
    stream = kinesis_client.describe_stream(StreamName=stream_name)
    if stream['StreamDescription']['StreamStatus'] == 'ACTIVE':
      return
    time.sleep(2)
    if attempt == retries - 1:
      raise RuntimeError(
          'Unable to initialize Kinesis stream %s. Status: %s' % (
              stream['StreamDescription']['StreamName'],
              stream['StreamDescription']['StreamStatus']))


@contextlib.contextmanager
def temp_kinesis_localstack():
  """Context manager to provide a temporary LocalStack Kinesis stream.

  Starts a LocalStack container, creates a unique Kinesis stream, and yields
  connection details for YAML Kinesis read/write transforms.

  Yields:
      dict: Keys include STREAM_NAME, AWS_ACCESS_KEY, AWS_SECRET_KEY, REGION,
          SERVICE_ENDPOINT (https), and VERIFY_CERTIFICATE (False).
  """
  try:
    import boto3
  except ImportError as exn:
    raise unittest.SkipTest('boto3 is not installed') from exn

  _LOGGER.info('Setting up LocalStack Kinesis fixture...')
  localstack = DockerContainer(_LOCALSTACK_KINESIS_IMAGE).with_bind_ports(
      4566, 4566)
  for port in range(4510, 4560):
    localstack = localstack.with_bind_ports(port, port)

  for attempt in range(4):
    try:
      localstack.start()
      break
    except Exception as exn:
      if attempt == 3:
        _LOGGER.error('Could not initialize localstack container')
        raise exn

  stream_name = f'yaml_kinesis_{uuid.uuid4().hex}'
  kinesis_client = None
  try:
    host = localstack.get_container_host_ip()
    port = localstack.get_exposed_port('4566')
    service_endpoint = f'https://localhost.localstack.cloud:{port}'
    http_endpoint = f'http://{host}:{port}'

    kinesis_client = boto3.client(
        service_name='kinesis',
        region_name=_LOCALSTACK_KINESIS_REGION,
        endpoint_url=http_endpoint,
        aws_access_key_id=_LOCALSTACK_KINESIS_ACCESS_KEY,
        aws_secret_access_key=_LOCALSTACK_KINESIS_SECRET_KEY,
    )
    _create_kinesis_stream(kinesis_client, stream_name)

    _LOGGER.info(
        'LocalStack Kinesis ready. Stream: [%s], endpoint: [%s]',
        stream_name,
        service_endpoint)

    yield {
        'STREAM_NAME': stream_name,
        'AWS_ACCESS_KEY': _LOCALSTACK_KINESIS_ACCESS_KEY,
        'AWS_SECRET_KEY': _LOCALSTACK_KINESIS_SECRET_KEY,
        'REGION': _LOCALSTACK_KINESIS_REGION,
        'SERVICE_ENDPOINT': service_endpoint,
        'VERIFY_CERTIFICATE': False,
    }
  finally:
    _LOGGER.info('Tearing down LocalStack Kinesis fixture...')
    if kinesis_client is not None:
      try:
        kinesis_client.delete_stream(
            StreamName=stream_name, EnforceConsumerDeletion=True)
      except Exception:
        _LOGGER.warning('Failed to delete kinesis stream [%s]', stream_name)
    try:
      localstack.stop()
    except Exception:
      _LOGGER.warning('Failed to stop localstack container')
    _LOGGER.info('LocalStack Kinesis fixture stopped.')


@contextlib.contextmanager
def temp_jms_activemq_server():
  """Context manager to provide a temporary ActiveMQ broker for JMS tests."""

  broker = DockerContainer('apache/activemq-classic:5.18.3').with_exposed_ports(
      61616)

  try:
    broker.start()
    wait_for_logs(broker, '.*ActiveMQ .* started.*', timeout=30)

    host = broker.get_container_host_ip()
    port = broker.get_exposed_port(61616)

    yield {
        'SERVER_URI': f'tcp://{host}:{port}',
        'CONNECTION_FACTORY_CLASS_NAME': 'org.apache.activemq.ActiveMQConnectionFactory',
    }
  finally:
    broker.stop()


@contextlib.contextmanager
def temp_ibm_mq_server():
  container = (
      DockerContainer('icr.io/ibm-messaging/mq:9.3.0.25-r1').with_env(
          'LICENSE', 'accept').with_env('MQ_QMGR_NAME', 'QM1').with_env(
              'MQ_APP_PASSWORD', 'admin123').with_exposed_ports(1414))

  try:
    container.start()
    wait_for_logs(container, '.*(MQQMNAME|Started queue manager).*', timeout=45)

    host = container.get_container_host_ip()
    port = container.get_exposed_port(1414)

    yield {
        'SERVER_URI': f'tcp://{host}:{port}?channel=DEV.APP.SVRCONN&queueManager=QM1',
        'CONNECTION_FACTORY_CLASS_NAME': 'com.ibm.mq.jakarta.jms.MQConnectionFactory',
        'USERNAME': 'app',
        'PASSWORD': 'admin123',
        'SOURCE_QUEUE': 'DEV.QUEUE.1',
        'SINK_QUEUE': 'DEV.QUEUE.2',
    }

  finally:
    container.stop()


@contextlib.contextmanager
def temp_kafka_server():
  """Context manager to provide a temporary Kafka server for testing.

  This function utilizes the 'testcontainers' library to spin up a Kafka
  instance within a Docker container. It then yields the bootstrap server
  string, which can be used by Kafka clients to connect to this temporary
  server.

  The Docker container and the Kafka instance are automatically managed
  and torn down when the context manager exits.

  Yields:
      str: The bootstrap server string for the temporary Kafka instance.
           Example format: "localhost:XXXXX" or "PLAINTEXT://localhost:XXXXX"

  Raises:
      Exception: If there's an error starting the Kafka container or
                 interacting with the temporary Kafka server.
  """
  with KafkaContainer() as kafka_container:
    try:
      yield kafka_container.get_bootstrap_server()
    except Exception as err:
      logging.error("Error interacting with temporary Kakfa Server: %s", err)
      raise err


@contextlib.contextmanager
def temp_pubsub_emulator(project_id="apache-beam-testing"):
  """
  Context manager to provide a temporary Pub/Sub emulator for testing.

  This function uses 'testcontainers' to spin up a Google Cloud SDK
  container running the Pub/Sub emulator. It yields the emulator host
  string (e.g., "localhost:xxxxx") that can be used to configure Pub/Sub
  clients.

  The Docker container is automatically managed and torn down when the
  context manager exits.

  Args:
      project_id (str): The GCP project ID to be used by the emulator.
                        This doesn't need to be a real project for the emulator.

  Yields:
      str: The host and port for the Pub/Sub emulator (e.g., "localhost:xxxx").
            This will be the address to point your Pub/Sub client to.

  Raises:
      Exception: If the container fails to start or the emulator isn't ready.
  """
  with PubSubContainer(project=project_id) as pubsub_container:
    publisher = pubsub_v1.PublisherClient()
    random_front_charactor = random.choice(string.ascii_lowercase)
    topic_id = f"{random_front_charactor}{uuid.uuid4().hex[:8]}"
    topic_name_to_create = \
      f"projects/{pubsub_container.project}/topics/{topic_id}"
    created_topic_object = publisher.create_topic(name=topic_name_to_create)
    yield created_topic_object.name


@contextlib.contextmanager
def temp_delta_table():
  with tempfile.TemporaryDirectory() as temp_dir:
    log_dir = os.path.join(temp_dir, "_delta_log")
    os.makedirs(log_dir, exist_ok=True)
    table_data = pa.table({"name": ["a", "b", "c"]})
    parquet_path = os.path.join(temp_dir, "part-00000.parquet")
    pq.write_table(table_data, parquet_path)
    file_size = os.path.getsize(parquet_path)
    commit_content = (
        '{"protocol":{"minReaderVersion":1,"minWriterVersion":2}}\n'
        '{"metaData":{"id":"test-id","format":{"provider":"parquet","options":{}},"schemaString":"{\\"type\\":\\"struct\\",\\"fields\\":[{\\"name\\":\\"name\\",\\"type\\":\\"string\\",\\"nullable\\":true,\\"metadata\\":{}}]}","partitionColumns":[],"configuration":{},"createdAt":123456789}}\n'
        f'{{"add":{{"path":"part-00000.parquet","partitionValues":{{}},"size":{file_size},"modificationTime":123456789,"dataChange":true}}}}\n'
    )
    commit_file = os.path.join(log_dir, "00000000000000000000.json")
    with open(commit_file, "w") as f:
      f.write(commit_content)
    yield temp_dir


@contextlib.contextmanager
def temp_delta_cdc_table():
  try:
    from deltalake import write_deltalake
  except ImportError as exn:
    raise unittest.SkipTest('deltalake is not installed') from exn

  with tempfile.TemporaryDirectory() as temp_dir:
    # Version 0 commit
    table_data = pa.table({"name": ["a", "b"]})
    write_deltalake(
        temp_dir,
        table_data,
        mode="overwrite",
        configuration={"delta.enableChangeDataFeed": "true"})

    # Version 1 commit
    table_data_1 = pa.table({"name": ["c"]})
    write_deltalake(temp_dir, table_data_1, mode="append")

    yield temp_dir


@contextlib.contextmanager
def temp_delta_cdc_table_with_updates():
  try:
    from deltalake import DeltaTable
    from deltalake import write_deltalake
  except ImportError as exn:
    raise unittest.SkipTest('deltalake is not installed') from exn

  # The CDC sink needs a non-null key.
  schema = pa.schema([
      pa.field("id", pa.int64(), nullable=False),
      pa.field("item", pa.string()),
      pa.field("qty", pa.int64()),
  ])

  def orders(*rows):
    return pa.Table.from_pylist([dict(zip(schema.names, row)) for row in rows],
                                schema=schema)

  with tempfile.TemporaryDirectory() as temp_dir:
    # Version 0: three new orders.
    write_deltalake(
        temp_dir,
        orders((1, "apple", 3), (2, "pear", 5), (3, "plum", 2)),
        configuration={"delta.enableChangeDataFeed": "true"})
    # Versions 1 to 3: update order 1, delete order 2, add order 4.
    table = DeltaTable(temp_dir)
    table.update(updates={"qty": "4"}, predicate="id = 1")
    table.delete("id = 2")
    write_deltalake(temp_dir, orders((4, "kiwi", 1)), mode="append")

    yield temp_dir


def replace_recursive(spec, vars):
  """Recursively replaces string placeholders in a spec with values from vars.

  Traverses a nested structure (dicts, lists, or other types). If a string
  is encountered and contains placeholders in the format '{key}', it attempts
  to replace them using the `vars` dictionary.

  Args:
    spec: The (potentially nested) structure to process.
    vars: A dictionary of variable names to their replacement values.

  Returns:
    The spec with placeholders replaced.

  Raises:
    ValueError: If a string formatting error occurs.
  """
  if isinstance(spec, dict):
    return {
        key: replace_recursive(value, vars)
        for (key, value) in spec.items()
    }
  elif isinstance(spec, list):
    return [replace_recursive(value, vars) for value in spec]
  # TODO(https://github.com/apache/beam/issues/35067): Consider checking for
  # callable in the if branch above instead of checking lambda here.
  elif isinstance(
      spec, str) and '{' in spec and '{\n' not in spec and 'lambda' not in spec:
    try:
      return spec.format(**vars)
    except Exception as exn:
      raise ValueError(f"Error evaluating {spec}: {exn}") from exn
  else:
    return spec


def transform_types(spec):
  """Recursively extracts all transform types from a pipeline specification.

  This generator function traverses a nested pipeline specification (likely
  parsed from YAML). It identifies and yields the 'type' string for each
  transform defined within the specification, including those within
  'composite' or 'chain' structures.

  Args:
    spec (dict): A dictionary representing a pipeline or transform
      specification.

  Yields:
    str: The 'type' of each transform found in the specification.
  """
  if spec.get('type', None) in (None, 'composite', 'chain'):
    if 'source' in spec:
      yield from transform_types(spec['source'])
    for t in spec.get('transforms', []):
      yield from transform_types(t)
    if 'sink' in spec:
      yield from transform_types(spec['sink'])
  else:
    yield spec['type']


def provider_sets(spec, require_available=False):
  """
  Generates all relevant combinations of providers for a given pipeline spec.

  This function analyzes a pipeline specification to identify transforms that
  can be implemented by multiple underlying providers (e.g., a generic
  transform vs. a SQL-backed one). It then yields different "provider sets,"
  each representing a unique combination of choices for these multi-provider
  transforms.

  If no transforms have multiple available providers, it yields a single
  provider set using the standard defaults.

  Args:
    spec (dict): The pipeline specification, typically loaded from YAML.
    require_available (bool): If True, raises an error if a provider
      needed for a transform is not available. If False (default),
      unavailable providers are skipped, potentially reducing the number
      of yielded combinations.

  Yields:
    tuple: A tuple where the first element is a string suffix uniquely
      identifying the provider combination (e.g., "MyTransform_SqlProvider_0"),
      and the second element is a dictionary mapping transform types to a list
      containing the selected provider(s) for that combination.
  """
  try:
    for p in spec['pipelines']:
      _ = yaml_transform.preprocess(copy.deepcopy(p['pipeline']))
  except Exception as exn:
    print(exn)
    all_transform_types = []
  else:
    all_transform_types = set.union(
        *(
            set(
                transform_types(
                    yaml_transform.preprocess(copy.deepcopy(p['pipeline']))))
            for p in spec['pipelines']))

  def filter_to_available(t, providers):
    if t == 'LogForTesting':
      # Don't fan out to all the (many) possibilities for this one...
      return [
          p for p in providers if isinstance(p, yaml_provider.InlineProvider)
      ][:1]
    elif require_available:
      for p in providers:
        if not p.available():
          raise ValueError("Provider {p} required for {t} is not available.")
      return providers
    else:
      return [p for p in providers if p.available()]

  standard_providers = yaml_provider.standard_providers()
  multiple_providers = {
      t: filter_to_available(t, standard_providers[t])
      for t in all_transform_types
      if len(filter_to_available(t, standard_providers[t])) > 1
  }
  if not multiple_providers:
    yield 'only', standard_providers
  else:
    names, provider_lists = zip(*sorted(multiple_providers.items()))
    for ix, c in enumerate(itertools.product(*provider_lists)):
      yield (
          '_'.join(
              n + '_' + type(p.underlying_provider()).__name__
              for (n, p) in zip(names, c)) + f'_{ix}',
          dict(standard_providers, **{n: [p]
                                      for (n, p) in zip(names, c)}))


def create_test_methods(spec):
  """Dynamically creates test methods based on a YAML specification.

  This function takes a YAML specification (`spec`) which defines pipelines,
  fixtures, and potentially options. It iterates through different
  combinations of "providers" (which determine how YAML transforms are
  implemented, e.g., using Python or SQL).

  For each combination of providers:
    1. It constructs a unique test method name (e.g., `test_only`).
    2. It defines a test method that:
        a. Sets up any specified fixtures, making their values available as
           variables.
        b. Mocks the standard YAML providers to use the current combination
           of providers for this test run.
        c. For each pipeline defined in the `spec`:
            i. Creates a `beam.Pipeline` instance with specified options.
            ii. Expands the YAML pipeline definition using
               `yaml_transform.expand_pipeline`, substituting any fixture
               variables.
            iii. Runs the Beam pipeline.

  The function yields tuples of (test_method_name, test_method_function),
  which can then be used to populate a `unittest.TestCase` class.

  Args:
    spec (dict): A dictionary parsed from a YAML test specification file.
      It's expected to have keys like 'fixtures' (optional) and 'pipelines'.

  Yields:
    tuple: A tuple containing:
      - str: The generated name for the test method (e.g., "test_only").
      - function: The dynamically generated test method.
  """
  for suffix, providers in provider_sets(spec):

    def test(self, providers=providers):  # default arg to capture loop value
      vars = {}
      with contextlib.ExitStack() as stack:
        stack.enter_context(
            mock.patch(
                'apache_beam.yaml.yaml_provider.standard_providers',
                lambda: providers))
        for fixture in spec.get('fixtures', []):
          vars[fixture['name']] = stack.enter_context(
              python_callable.PythonCallableWithSource.
              load_from_fully_qualified_name(fixture['type'])(
                  **yaml_transform.SafeLineLoader.strip_metadata(
                      fixture.get('config', {}))))
        for pipeline_spec in spec['pipelines']:
          try:
            with beam.Pipeline(options=PipelineOptions(
                pickle_library='cloudpickle',
                **replace_recursive(
                    yaml_transform.SafeLineLoader.strip_metadata(
                        pipeline_spec.get('options', {})),
                    vars))) as p:
              yaml_transform.expand_pipeline(
                  p, replace_recursive(pipeline_spec, vars))
          except ValueError as exn:
            # FnApiRunner currently does not support this requirement in
            # some xlang scenarios (e.g. Iceberg YAML pipelines).
            if 'beam:requirement:pardo:on_window_expiration:v1' in str(exn):
              self.skipTest(
                  'Runner does not support '
                  'beam:requirement:pardo:on_window_expiration:v1')
            raise

    yield f'test_{suffix}', test


_SICKBAY_TESTS = {
    'ml_transform.yaml': 'Requires broken TFT dependency types (e.g. ScaleTo01)',
}


def parse_test_files(filepattern):
  """Parses YAML test files and dynamically creates test cases.

  This function iterates through all files matching the given glob pattern.
  For each YAML file found, it:
    1. Reads the file content.
    2. Determines a test suite name based on the file name.
    3. Calls `create_test_methods` to generate test methods from the
       YAML specification.
    4. Dynamically creates a new TestCase class (inheriting from
       `unittest.TestCase`) and populates it with the generated test methods.
    5. Adds this newly created TestCase class to the global scope, making it
       discoverable by the unittest framework.

  Args:
    filepattern (str): A glob pattern specifying the YAML test files to parse.
      For example, 'path/to/tests/*.yaml'.
  """
  for path in glob.glob(filepattern):
    with open(path) as fin:
      filename = os.path.basename(path)
      suite_name = os.path.splitext(filename)[0].title().replace(
          '-', '') + 'Test'
      print(path, suite_name)
      methods = dict(
          create_test_methods(
              yaml.load(fin, Loader=yaml_transform.SafeLineLoader)))
      suite_class = type(suite_name, (unittest.TestCase, ), methods)
      if filename in _SICKBAY_TESTS:
        suite_class = unittest.skip(f"Sickbayed: {_SICKBAY_TESTS[filename]}")(
            suite_class)
      globals()[suite_name] = suite_class


# Logging setups
logging.getLogger().setLevel(logging.INFO)

# Dynamically create test methods from the tests directory.
# yaml_test_files_dir comes from conftest.py and set by pytest_configure.
_test_files_dir = yaml_test_files_dir
_file_pattern = os.path.join(
    os.path.dirname(__file__), _test_files_dir, '*.yaml')
parse_test_files(_file_pattern)

if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()