Back to Repositories

Testing SFTP Connection Management System in Apache Airflow

This test suite implements helper functions for SFTP connection management in Apache Airflow, focusing on secure credential handling and connection configuration for SFTP testing scenarios. The suite provides context management utilities for temporary SFTP connections during test execution.

Test Coverage Overview

The test coverage focuses on SFTP connection handling and credential management within Airflow.

Tests credential loading from JSON files
Validates connection URI generation
Covers environment variable management
Handles various SFTP configuration parameters

Implementation Analysis

The testing approach utilizes Python’s context manager pattern for managing temporary SFTP connections during test execution. The implementation leverages Airflow’s Connection model and environment patching utilities to ensure isolated testing environments.

Key patterns include JSON credential parsing, connection URI generation, and environment variable manipulation.

Technical Details

Testing components include:

contextlib for context manager implementation
Airflow Connection model for SFTP configuration
process_utils.patch_environ for environment isolation
JSON parsing for credential management
Environment variable handling for connection IDs

Best Practices Demonstrated

The test implementation showcases several testing best practices for infrastructure components.

Proper resource cleanup using context managers
Secure credential handling
Environment isolation
Clear error handling for invalid configurations
Modular test helper design

apache/airflow

tests_common/test_utils/sftp_system_helpers.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 json
import os
from contextlib import contextmanager

from airflow.exceptions import AirflowException
from airflow.models import Connection
from airflow.utils.process_utils import patch_environ

SFTP_CONNECTION_ID = os.environ.get("SFTP_CONNECTION_ID", "sftp_default")


@contextmanager
def provide_sftp_default_connection(key_file_path: str):
    """
    Context manager to provide a temporary value for sftp_default connection.

    :param key_file_path: Path to file with sftp_default credentials .json file.
    """
    if not key_file_path.endswith(".json"):
        raise AirflowException("Use a JSON key file.")
    with open(key_file_path) as credentials:
        creds = json.load(credentials)
    conn = Connection(
        conn_id=SFTP_CONNECTION_ID,
        conn_type="ssh",
        port=creds.get("port", None),
        host=creds.get("host", None),
        login=creds.get("login", None),
        password=creds.get("password", None),
        extra=json.dumps(creds.get("extra", None)),
    )
    with patch_environ({f"AIRFLOW_CONN_{conn.conn_id.upper()}": conn.get_uri()}):
        yield