sdks/python/conftest.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.
#

"""Pytest configuration and custom hooks."""

import gc
import math
import os
import sys
import threading
import time

import pytest

from apache_beam.options import pipeline_options
from apache_beam.testing.test_pipeline import TestPipeline

MAX_SUPPORTED_PYTHON_VERSION = (3, 14)


def pytest_addoption(parser):
  parser.addoption(
      '--test-pipeline-options',
      help='Options to use in test pipelines. NOTE: Tests may '
      'ignore some or all of these options.')


# See pytest.ini for main collection rules.
collect_ignore_glob = [
    '*_py3%d.py' % minor for minor in range(
        sys.version_info.minor + 1, MAX_SUPPORTED_PYTHON_VERSION[1] + 1)
]


@pytest.fixture(scope="session", autouse=True)
def configure_beam_rpc_timeouts():
  """
  Configure gRPC and RPC timeouts for Beam tests
  to prevent DEADLINE_EXCEEDED errors.
  """
  print("\n--- Applying Beam RPC timeout configuration ---")

  timeout_env_vars = {
      'GRPC_ARG_KEEPALIVE_TIME_MS': '30000',
      'GRPC_ARG_KEEPALIVE_TIMEOUT_MS': '10000',
      'GRPC_ARG_HTTP2_MAX_PINGS_WITHOUT_DATA': '0',
      'GRPC_ARG_KEEPALIVE_PERMIT_WITHOUT_CALLS': '1',
      'GRPC_ARG_HTTP2_MIN_RECV_PING_INTERVAL_WITHOUT_DATA_MS': '300000',
      'GRPC_ARG_HTTP2_MIN_SENT_PING_INTERVAL_WITHOUT_DATA_MS': '10000',

      # Additional stability settings for DinD environment
      'GRPC_ARG_MAX_RECONNECT_BACKOFF_MS': '120000',
      'GRPC_ARG_INITIAL_RECONNECT_BACKOFF_MS': '1000',
      'GRPC_ARG_MAX_CONNECTION_IDLE_MS': '600000',
      'GRPC_ARG_MAX_CONNECTION_AGE_MS': '1800000',

      # Beam-specific retry and timeout settings
      'BEAM_RETRY_MAX_ATTEMPTS': '5',
      'BEAM_RETRY_INITIAL_DELAY_MS': '1000',
      'BEAM_RETRY_MAX_DELAY_MS': '60000',
      'BEAM_RUNNER_BUNDLE_TIMEOUT_MS': '300000',

      # Force deterministic execution in DinD environment
      'BEAM_TESTING_FORCE_SINGLE_BUNDLE': 'true',
      'BEAM_TESTING_DETERMINISTIC_ORDER': 'true',
      'BEAM_SDK_WORKER_PARALLELISM': '1',
      'BEAM_WORKER_POOL_SIZE': '1',
      'BEAM_FN_API_CONTROL_PORT': '0',
      'BEAM_FN_API_DATA_PORT': '0',

  # Container-specific stability settings
      'PYTHONHASHSEED': '0',
      'OMP_NUM_THREADS': '1',
      'OPENBLAS_NUM_THREADS': '1',

      # Force sequential pytest execution (CRITICAL for DinD stability)
      'PYTEST_XDIST_WORKER_COUNT': '1',
      'PYTEST_CURRENT_TEST_TIMEOUT': '300',

      # Mock and test isolation improvements
      'PYTEST_MOCK_TIMEOUT': '60',
      'BEAM_TEST_ISOLATION_MODE': 'strict',
  }

  for key, value in timeout_env_vars.items():
    os.environ[key] = value
    print(f"Set {key}={value}")

  print("Successfully configured Beam RPC timeouts")


@pytest.fixture(scope="class", autouse=True)
def ensure_clean_state():
  """
  Ensure clean state before each test class
  to prevent cross-test contamination.
  Runs once per test class instead of per test to reduce overhead.
  """
  # Force garbage collection to clean up any lingering resources
  gc.collect()

  # Log active thread count for debugging
  thread_count = threading.active_count()
  if thread_count > 50:
    print(f"Warning: {thread_count} active threads detected before test class")
    # Force a brief pause to let threads settle
    time.sleep(0.5)
    gc.collect()

  yield

  # Enhanced cleanup after test class
  try:
    # Force more aggressive cleanup
    gc.collect()
    # Brief pause to let any async operations complete
    time.sleep(0.1)
    # Additional garbage collection
    gc.collect()
  except Exception as e:
    print(f"Warning: Cleanup error: {e}")


@pytest.fixture(scope="class", autouse=True)
def enhance_mock_stability():
  """
  Enhance mock stability in DinD environment.
  Runs once per test class instead of per test to reduce overhead.
  """
  # Brief pause before test class to ensure clean mock state
  time.sleep(0.05)

  yield

  # Brief pause after test class to let mocks clean up
  time.sleep(0.05)


def _configure_testcontainers_timeouts():
  """Applies TC_* environment settings to testcontainers' global config.

  testcontainers 4.x computes its wait timeout as ``max_tries * sleep_time``
  from ``testcontainers_config``. It reads ``TC_MAX_TRIES`` and
  ``TC_POOLING_INTERVAL`` itself, but not ``TC_TIMEOUT`` or ``TC_SLEEP_TIME``,
  which Beam's CI workflows set. Translate them here so that the effective
  timeout is at least ``TC_TIMEOUT`` seconds.
  """
  from testcontainers.core.config import testcontainers_config

  sleep_time = float(
      os.getenv("TC_SLEEP_TIME") or os.getenv("TC_POOLING_INTERVAL") or "1")
  timeout = float(os.getenv("TC_TIMEOUT", "120"))
  max_tries = int(os.getenv("TC_MAX_TRIES", "0"))

  testcontainers_config.sleep_time = sleep_time
  testcontainers_config.max_tries = max(
      max_tries, math.ceil(timeout / sleep_time))
  print(
      "Set testcontainers max_tries=%d, sleep_time=%s (timeout=%ss)" % (
          testcontainers_config.max_tries,
          testcontainers_config.sleep_time,
          testcontainers_config.timeout))


def pytest_configure(config):
  """Saves options added in pytest_addoption for later use.
  This is necessary since pytest-xdist workers do not have the
  same sys.argv as the main pytest invocation.
  xdist does seem to pickle TestPipeline
  """
  # for the entire test session.
  print("\n--- Applying global testcontainers timeout configuration ---")
  try:
    _configure_testcontainers_timeouts()
  except ImportError:
    print("The testcontainers library is not installed.")

  TestPipeline.pytest_test_pipeline_options = config.getoption(
      'test_pipeline_options', default='')
  # Enable optional type checks on all tests.
  pipeline_options.enable_all_additional_type_checks()