From a2c58f5a13cc4e1bee172b807b7066cdce437c56 Mon Sep 17 00:00:00 2001
From: Justin Bandoro <79104794+jbandoro@users.noreply.github.com>
Date: Mon, 4 Dec 2023 03:14:49 -0800
Subject: [PATCH] Speed up integration tests (#732)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Speed up the integration tests by caching the dag bag result.

Previously, for each parametrized dag run test, it would reparse all of
the dags, which takes a non-trivial amount of time to parse all of the
cosmos example dags.

On my local machine, the total time went from 1616s to 540s, which will
save a good amount of GH minutes 😃.
---
 tests/test_example_dags.py                | 2 ++
 tests/test_example_dags_no_connections.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/tests/test_example_dags.py b/tests/test_example_dags.py
index 63abca541..41b967e92 100644
--- a/tests/test_example_dags.py
+++ b/tests/test_example_dags.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from functools import cache
 
 import airflow
 import pytest
@@ -37,6 +38,7 @@ def session():
     return get_session()
 
 
+@cache
 def get_dag_bag() -> DagBag:
     """Create a DagBag by adding the files that are not supported to .airflowignore"""
     with open(AIRFLOW_IGNORE_FILE, "w+") as file:
diff --git a/tests/test_example_dags_no_connections.py b/tests/test_example_dags_no_connections.py
index 7048b1b0e..97f39bf99 100644
--- a/tests/test_example_dags_no_connections.py
+++ b/tests/test_example_dags_no_connections.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from functools import cache
 
 import airflow
 import pytest
@@ -23,6 +24,7 @@
 }
 
 
+@cache
 def get_dag_bag() -> DagBag:
     """Create a DagBag by adding the files that are not supported to .airflowignore"""
     with open(AIRFLOW_IGNORE_FILE, "w+") as file: