Skip to content

feat(testing): Add pytest fixtures for FlyteContext, cache, and tmp dirs#3411

Open
mvanhorn wants to merge 1 commit intoflyteorg:masterfrom
mvanhorn:osc/feat-pytest-testing-fixtures
Open

feat(testing): Add pytest fixtures for FlyteContext, cache, and tmp dirs#3411
mvanhorn wants to merge 1 commit intoflyteorg:masterfrom
mvanhorn:osc/feat-pytest-testing-fixtures

Conversation

@mvanhorn
Copy link

Summary

Adds three pytest fixtures and a context manager to flytekit.testing, eliminating common test setup boilerplate. Fixtures auto-register via a pytest11 entry point so they're available without explicit imports.

Why this matters

Flytekit's testing module currently ships two utilities (task_mock and patch). Users writing tests for Flyte workflows repeat the same boilerplate across test files:

  • FlyteContextManager.current_context() to get a context - appears 30+ times in flytekit's own test suite
  • LocalTaskCache.initialize() / LocalTaskCache.clear() before cached task tests - the manual workaround discovered in flyteorg/flyte#5657
  • tempfile.TemporaryDirectory() setup for file-based operations

Other workflow orchestrator SDKs ship dedicated testing utilities (Dagster's execute_in_process(), Prefect's prefect_test_harness, Airflow's pytest-airflow).

Changes

New files:

  • flytekit/testing/fixtures.py - Three fixtures and one context manager:
    • flyte_context - returns the current FlyteContext
    • flyte_cache - initializes and clears LocalTaskCache before/after each test
    • flyte_tmp_dir - provides a managed Path temporary directory
    • workflow_dry_run() - context manager for clean local workflow execution
  • flytekit/testing/conftest.py - pytest11 plugin entry point for auto-discovery
  • tests/flytekit/unit/testing/test_fixtures.py - 12 tests covering all fixtures

Modified files:

  • flytekit/testing/__init__.py - exports new fixtures
  • pyproject.toml - registers [project.entry-points.pytest11] so fixtures are auto-available

Usage

# No import needed - fixtures auto-register via pytest11

def test_type_transform(flyte_context):
    lt = TypeEngine.to_literal_type(int)
    lv = TypeEngine.to_literal(flyte_context, 42, int, lt)
    assert lv.scalar.primitive.integer == 42

def test_cached_task(flyte_cache):
    # Cache cleared before and after - no flaky results from stale on-disk cache
    result = my_cached_task(a=1, b=2)
    assert result == 3

def test_file_ops(flyte_tmp_dir):
    path = flyte_tmp_dir / "output.txt"
    path.write_text("hello")
    assert path.read_text() == "hello"

Testing

All 12 tests pass locally on Python 3.12:

tests/flytekit/unit/testing/test_fixtures.py - 12 passed in 0.05s

Tests verify fixture behavior, cache isolation between tests, type transformations with context, and workflow dry-run execution.

This contribution was developed with AI assistance (Claude Code).

…, and tmp dirs

Add three pytest fixtures to flytekit.testing that eliminate common test
boilerplate:

- flyte_context: provides the current FlyteContext without manually
  calling FlyteContextManager.current_context()
- flyte_cache: auto-initializes and clears LocalTaskCache before/after
  each test, preventing stale cache from causing flaky tests (see
  flyteorg/flyte#5657)
- flyte_tmp_dir: provides a managed temporary directory

Also adds workflow_dry_run() context manager for running workflows
locally with a clean cache state.

Fixtures are auto-registered via pytest11 entry point in pyproject.toml,
so they're available in any pytest session without explicit imports.

Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
@mvanhorn mvanhorn force-pushed the osc/feat-pytest-testing-fixtures branch from c37e8c6 to d285c5f Compare March 19, 2026 04:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant