Back to Repositories

Testing Configuration Management Utilities in Apache Airflow

This test suite implements configuration management utilities for Apache Airflow testing, focusing on temporary configuration overrides and environment variable manipulation. It provides essential context managers for modifying Airflow’s configuration and environment variables during test execution.

Test Coverage Overview

The test suite covers configuration management functionality in Apache Airflow with comprehensive coverage of both configuration file settings and environment variables.

Key areas tested include:

Temporary configuration overrides using conf_vars
Environment variable manipulation using env_vars
Section and key management in configuration
Proper restoration of original settings

Implementation Analysis

The testing approach utilizes Python’s contextlib.contextmanager to create safe, temporary modifications to configuration settings. The implementation employs context managers that ensure proper cleanup after test execution, preventing configuration state leakage between tests.

Technical patterns include:

Context manager implementation for configuration handling
Environment variable manipulation
Section-key pair management
Automatic cleanup and restoration

Technical Details

Testing tools and components:

contextlib for context manager implementation
os module for environment variable handling
Airflow settings and configuration modules
Python dictionary-based configuration overrides

Best Practices Demonstrated

The test implementation showcases excellent testing practices for configuration management. It demonstrates proper resource cleanup, isolation of test environments, and prevention of test interference.

Notable practices include:

Proper state restoration after tests
Isolation of configuration changes
Clean separation of concerns
Robust error handling

apache/airflow

tests_common/test_utils/config.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.
from __future__ import annotations

import contextlib
import os

from airflow import settings
from airflow.configuration import conf


@contextlib.contextmanager
def conf_vars(overrides):
    original = {}
    original_env_vars = {}
    for (section, key), value in overrides.items():
        env = conf._env_var_name(section, key)
        if env in os.environ:
            original_env_vars[env] = os.environ.pop(env)

        if conf.has_option(section, key):
            original[(section, key)] = conf.get(section, key)
        else:
            original[(section, key)] = None
        if value is not None:
            if not conf.has_section(section):
                conf.add_section(section)
            conf.set(section, key, value)
        else:
            if conf.has_section(section):
                conf.remove_option(section, key)
    settings.configure_vars()
    try:
        yield
    finally:
        for (section, key), value in original.items():
            if value is not None:
                if not conf.has_section(section):
                    conf.add_section(section)
                conf.set(section, key, value)
            else:
                if conf.has_section(section):
                    conf.remove_option(section, key)
        for env, value in original_env_vars.items():
            os.environ[env] = value
        settings.configure_vars()


@contextlib.contextmanager
def env_vars(overrides):
    """
    Temporarily patches env vars, restoring env as it was after context exit.

    Example:
        with env_vars({'AIRFLOW_CONN_AWS_DEFAULT': 's3://@'}):
            # now we have an aws default connection available
    """
    orig_vars = {}
    new_vars = []
    for env, value in overrides.items():
        if env in os.environ:
            orig_vars[env] = os.environ.pop(env, "")
        else:
            new_vars.append(env)
        os.environ[env] = value
    try:
        yield
    finally:
        for env, value in orig_vars.items():
            os.environ[env] = value
        for env in new_vars:
            os.environ.pop(env)