Stacks

Write unit tests for untested code in Apache Airflow with Atlas in 2026

Updated 8 min read

In 2026, Apache Airflow developers can efficiently write unit tests for untested code using Atlas, which integrates directly with `pytest (DagBag)` for execution, `uv` for package management, and `ruff format` for code style. Atlas reads your existing codebase, identifies untested modules, and then generates new test files that mirror your project's established conventions, ensuring direct integration and maintainability.

How to find untested Apache Airflow code and existing test patterns?

In 2026, Atlas helps Apache Airflow developers identify untested code and existing test conventions by first using the `lsp` tool's `documentSymbol` operation to enumerate all public functions within a target module. It then employs `grep` to search for existing test files, ensuring new tests align with the repository's established framework and naming styles.

To begin writing unit tests for an untested Apache Airflow module, Atlas first needs to understand the module's public interface and the project's existing testing patterns. Atlas uses its `read` tool to ingest the target module's source code. Following this, the `lsp` tool's `documentSymbol` operation is invoked to precisely enumerate all exported symbols and public functions, ensuring no part of the module's API is overlooked. This detailed symbol list forms the basis for generating comprehensive test cases. Concurrently, Atlas utilizes its `grep` tool to scan the Apache Airflow repository for existing test files, typically found in a `tests/` directory or alongside `dags/` folders. This step is crucial for identifying the project's established `pytest (DagBag)` framework, preferred import styles, and naming conventions. By copying these existing patterns, Atlas ensures that any newly generated tests direct integrate into the codebase, maintaining consistency and reducing review friction for the developer.

How Atlas writes new `pytest (DagBag)` tests for Apache Airflow?

Atlas drafts new unit tests for Apache Airflow modules in 2026 by leveraging its `write` tool, which generates a new test file based on the identified module symbols and existing test conventions. This process ensures that every public function receives a corresponding test case, adhering to the project's `pytest (DagBag)` structure and import styles.

Once Atlas has a clear understanding of the target Apache Airflow module's public functions and the repository's testing conventions, it proceeds to draft the new unit test file. The `write` tool is central to this process. Atlas synthesizes the information gathered from the `lsp` and `grep` operations to construct a new test file, for example, `tests/dags/new_module_test.py`. This file will include boilerplate for `pytest (DagBag)` tests, appropriate import statements matching the project's style, and initial test stubs for each public function identified. For Apache Airflow specific testing, Atlas can generate `DagBag` import tests, which are critical for validating DAG parsing and structure. Before any changes are committed to disk, Atlas presents a unified diff of the proposed test file, allowing the developer to review and approve the generated code, ensuring full transparency and control over the agent's output.

How to run and debug `pytest (DagBag)` tests in Apache Airflow using Atlas?

After drafting, Atlas executes the new `pytest (DagBag)` test suite for Apache Airflow using its `bash` tool, providing immediate feedback on test failures. If the output exceeds 2000 lines or 50 KB, Atlas truncates it in the TUI but saves the full log to a file, allowing developers to use the `edit` tool to iterate on fixes.

The true value of a test lies in its execution. Atlas facilitates this by using its `bash` tool to run the newly created `pytest (DagBag)` test suite for your Apache Airflow module. The command `pytest tests/dags/new_module_test.py` is executed directly within your terminal, providing real-time feedback on test outcomes. If tests fail, Atlas displays the output in its TUI. For extensive test runs where output might exceed 2000 lines or 50 KB, Atlas intelligently truncates the display in the terminal but saves the complete log to a file, ensuring no diagnostic information is lost. Developers can then use the `edit` tool to modify the test file or the module under test, iterating on fixes. For larger modules or complex failures, the `todowrite` tool can be used to maintain a list of pending tasks, streamlining the debugging process until all tests pass. After edits, Atlas can re-run `pytest (DagBag)` and then apply `ruff format` via `atlas bash ruff format tests/dags/new_module_test.py` to ensure code style consistency.

How Atlas ensures safe and reviewable Apache Airflow test changes?

Atlas ensures safety and reviewability for Apache Airflow test changes in 2026 by computing a unified diff for every file edit and surfacing it for explicit approval before writing to disk. Every Atlas tool call, including `write` and `bash`, is permission-gated against allow, ask, and deny rules, providing developers with granular control over agent actions.

Safety and developer control are paramount when Atlas makes changes to an Apache Airflow codebase. Before any modifications are written to disk, Atlas computes a unified diff for every proposed file edit and presents it to the developer for explicit approval. This ensures that developers always have the final say on what code lands in their repository. Furthermore, every Atlas tool call, whether it's `write` for generating test files or `bash` for running `pytest (DagBag)`, is permission-gated. Developers can configure allow, ask, or deny rules for each tool, providing fine-grained control over Atlas's actions. Atlas operates with a read-only plan agent before switching to a build agent for execution, and it snapshots file changes as git patches, allowing edits to be easily diffed and rolled back if necessary. Atlas also reads git branches, status, and diffs, and can stage and create commits on your behalf, integrating direct into your existing Apache Airflow development workflow.

What are the setup steps for Atlas in an Apache Airflow repository?

Integrating Atlas into an Apache Airflow development workflow in 2026 begins by running Atlas within a repository containing a `dags/` folder and a `pyproject.toml` file pinning `apache-airflow`. Atlas then reads your DAG objects, task dependencies, and custom operators, enabling it to understand the codebase context for effective test generation and refactoring.

Setting up Atlas for an Apache Airflow project is straightforward. The primary requirement is to run Atlas within a repository that contains a `dags/` folder, which is standard for Airflow projects, and a `pyproject.toml` file that explicitly pins the `apache-airflow` package. This setup allows Atlas to correctly identify the project context. Once initialized, Atlas can read your Apache Airflow DAG objects, understand task dependencies, parse connection and variable lookups, and recognize any custom operators you have defined. This deep contextual understanding is crucial for Atlas to generate relevant and accurate unit tests, especially `pytest (DagBag)` tests for DAG parsing. Atlas also integrates with your chosen package manager, such as `uv`, and your formatter, `ruff format`, ensuring that all generated or modified code adheres to your project's standards from the outset. This comprehensive integration makes Atlas a powerful ally for maintaining high code quality in your Apache Airflow projects.

Step by step

  1. 01Run `atlas read <path/to/untested_module.py>` to ingest the target Apache Airflow module's code.
  2. 02Execute `atlas lsp documentSymbol <path/to/untested_module.py>` to enumerate all public functions for testing.
  3. 03Use `atlas grep 'import pytest'` in your `tests/` directory to find an existing `pytest (DagBag)` test file (e.g., `tests/dags/example_dag_test.py`) and copy its conventions.
  4. 04Draft the new test file (e.g., `tests/dags/new_module_test.py`) with `atlas write tests/dags/new_module_test.py`, reviewing the generated diff for approval.
  5. 05Run the new `pytest (DagBag)` test suite using `atlas bash pytest tests/dags/new_module_test.py` to check for failures.
  6. 06If tests fail, use `atlas edit <path/to/test_file.py>` or `atlas edit <path/to/module.py>` to make corrections, leveraging `todowrite` for larger tasks.
  7. 07Re-run `atlas bash pytest tests/dags/new_module_test.py` until all tests pass and the suite is green.
  8. 08Apply code formatting to the new test file with `atlas bash ruff format tests/dags/new_module_test.py`.

Frequently asked questions

How does Atlas ensure new Apache Airflow tests match my repo's style?
Atlas uses its `grep` tool to find existing `pytest (DagBag)` test files in your Apache Airflow repository, then copies their framework, import style, and naming conventions when generating new tests.
Can Atlas run `pytest (DagBag)` directly for Apache Airflow tests?
Yes, Atlas uses its `bash` tool to execute `pytest (DagBag)` commands directly within your Apache Airflow environment, providing real-time feedback on test results and allowing for iterative debugging.
What Apache Airflow specific information does Atlas understand?
Atlas reads Apache Airflow DAG objects, task dependencies, connection and variable lookups, and any custom operators, building a comprehensive understanding of your codebase for accurate test generation.
How does Atlas handle large test output when running Apache Airflow tests?
If `pytest (DagBag)` output exceeds 2000 lines or 50 KB, Atlas truncates it in the TUI but saves the full log to a file, which you can then inspect using your preferred editor.
Is it safe to let Atlas modify my Apache Airflow codebase?
Yes, Atlas is permission-gated. Every file edit generates a unified diff for your approval, and all tool calls are subject to allow, ask, or deny rules before execution, ensuring full developer control.
Does Atlas support `ruff format` for Apache Airflow code?
Absolutely. Atlas can run `ruff format` via its `bash` tool to ensure all generated or modified Apache Airflow test code adheres to your project's formatting standards, maintaining code consistency.
How does Atlas identify all functions in an Apache Airflow module for testing?
Atlas uses the `lsp` tool's `documentSymbol` operation to enumerate all exported symbols and public functions within a target Apache Airflow module, ensuring comprehensive test coverage for every public API.

Try Atlas in your terminal

The terminal-native AI coding agent. Free core, single binary.

Install Atlas

Related guides

Write Unit Tests for Untested Code with Atlas in 2026

How to write unit tests for untested code with Atlas in 2026: the lsp tool enumerates exported symbols, grep copies repo conventions, and bash actually runs the suite.

Review a pull request in Apache Airflow with Atlas in 2026

Streamline Apache Airflow pull request reviews in 2026 with Atlas. Catch subtle bugs in DAG definitions, task dependencies, and top-level code using AI-powered context and real Airflow tools.

Migrate a Deprecated API Across Every Callsite in Apache Airflow with Atlas in 2026

Efficiently migrate deprecated APIs across your entire Apache Airflow codebase with Atlas. Enumerate all callsites, apply context-anchored patches, and validate changes with `pytest (DagBag)` and `ruff format` for a

Diagnose a Hanging or Long-Running Command in Apache Airflow with Atlas in 2026

In 2026, use Atlas to diagnose and resolve hanging or slow Apache Airflow commands. Identify interactive input blocks or genuinely slow scripts within your DAGs and unblock them efficiently.

Automate GitHub Issue and Pull Request Triage in Apache Airflow with Atlas in 2026

Automate GitHub issue and pull request triage in Apache Airflow with Atlas. Safely respond to events, enforce trusted users, and integrate with your dags/ and pyproject.toml setup. Atlas leverages pytest (DagBag) and

Run the Test Suite and Triage Failures in Apache Airflow with Atlas in 2026

Streamline Apache Airflow test triage in 2026. Atlas helps developers turn walls of `pytest (DagBag)` output into prioritized root causes, using `uv` and `ruff format`.

Extract a Shared Helper from Duplicated Code in Apache Airflow with Atlas in 2026

Streamline your Apache Airflow DAGs in 2026 by extracting duplicated logic into a shared, tested helper using Atlas. Find semantic duplicates, create new modules, and apply changes with confidence.

Refactor a Legacy Module in Apache Airflow with Atlas in 2026

Streamline Apache Airflow DAGs and tasks by refactoring legacy Python modules with Atlas. Safely restructure code, verify behavior with pytest (DagBag), and ensure no breaking changes in 2026.

Browse this resource hub