#
# 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 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
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.internal.clients import bigquery
from apache_beam.io.gcp.spanner_wrapper import SpannerWrapper
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.utils import python_callable
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_bigquery_table(project, prefix='yaml_bq_it_'):
  """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.

  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_'.

  Yields:
    str: The full path for a temporary BigQuery table within the created
    dataset.
         Example: 'my-project.yaml_bq_it_a1b2c3d4e5f6...tmp_table'
  """
  bigquery_client = BigQueryWrapper()
  dataset_id = '%s_%s' % (prefix, uuid.uuid4().hex)
  bigquery_client.get_or_create_dataset(project, dataset_id)
  logging.info("Created dataset %s in project %s", dataset_id, project)
  yield f'{project}.{dataset_id}.tmp_table'
  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()

    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 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_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)


@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


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()
