Introduce StdoutCaptureManager to isolate stdout from logging logs

airflow-core/tests/unit/cli/commands/test_asset_command.py

@@ -18,8 +18,6 @@
 
 from __future__ import annotations
 
-import contextlib
-import io
 import json
 import os
 import typing
@@ -56,34 +54,34 @@ def parser() -> ArgumentParser:
     return cli_parser.get_parser()
 
 
-def test_cli_assets_list(parser: ArgumentParser) -> None:
+def test_cli_assets_list(parser: ArgumentParser, stdout_capture) -> None:
     args = parser.parse_args(["assets", "list", "--output=json"])
-    with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+    with stdout_capture as capture:
         asset_command.asset_list(args)
 
-    asset_list = json.loads(temp_stdout.getvalue())
+    asset_list = json.loads(capture.getvalue())
     assert len(asset_list) > 0
     assert set(asset_list[0]) == {"name", "uri", "group", "extra"}
     assert any(asset["uri"] == "s3://dag1/output_1.txt" for asset in asset_list), asset_list
 
 
-def test_cli_assets_alias_list(parser: ArgumentParser) -> None:
+def test_cli_assets_alias_list(parser: ArgumentParser, stdout_capture) -> None:
     args = parser.parse_args(["assets", "list", "--alias", "--output=json"])
-    with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+    with stdout_capture as capture:
         asset_command.asset_list(args)
 
-    alias_list = json.loads(temp_stdout.getvalue())
+    alias_list = json.loads(capture.getvalue())
     assert len(alias_list) > 0
     assert set(alias_list[0]) == {"name", "group"}
     assert any(alias["name"] == "example-alias" for alias in alias_list), alias_list
 
 
-def test_cli_assets_details(parser: ArgumentParser) -> None:
+def test_cli_assets_details(parser: ArgumentParser, stdout_capture) -> None:
     args = parser.parse_args(["assets", "details", "--name=asset1_producer", "--output=json"])
-    with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+    with stdout_capture as capture:
         asset_command.asset_details(args)
 
-    asset_detail_list = json.loads(temp_stdout.getvalue())
+    asset_detail_list = json.loads(capture.getvalue())
     assert len(asset_detail_list) == 1
 
     # No good way to statically compare these.
@@ -106,12 +104,12 @@ def test_cli_assets_details(parser: ArgumentParser) -> None:
     }
 
 
-def test_cli_assets_alias_details(parser: ArgumentParser) -> None:
+def test_cli_assets_alias_details(parser: ArgumentParser, stdout_capture) -> None:
     args = parser.parse_args(["assets", "details", "--alias", "--name=example-alias", "--output=json"])
-    with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+    with stdout_capture as capture:
         asset_command.asset_details(args)
 
-    alias_detail_list = json.loads(temp_stdout.getvalue())
+    alias_detail_list = json.loads(capture.getvalue())
     assert len(alias_detail_list) == 1
 
     # No good way to statically compare these.
@@ -121,15 +119,17 @@ def test_cli_assets_alias_details(parser: ArgumentParser) -> None:
 
 
 @mock.patch("airflow.api_fastapi.core_api.datamodels.dag_versions.hasattr")
-def test_cli_assets_materialize(mock_hasattr, parser: ArgumentParser) -> None:
+def test_cli_assets_materialize(mock_hasattr, parser: ArgumentParser, stdout_capture) -> None:
     mock_hasattr.return_value = False
     args = parser.parse_args(["assets", "materialize", "--name=asset1_producer", "--output=json"])
-    with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+    with stdout_capture as capture:
         asset_command.asset_materialize(args)
 
-    output = temp_stdout.getvalue()
-    # Skip the first line of `temp_stdout` since the current `DAGRunResponse` requires `DagBundlesManager`, which logs `INFO - DAG bundles loaded: dags-folder, example_dags`.
-    output = "\n".join(output.splitlines()[1:])
+    output = capture.getvalue()
+
+    # Check if output is empty first
+    assert output, "No output captured from asset_materialize command"
+
     run_list = json.loads(output)
     assert len(run_list) == 1
 
@@ -162,12 +162,12 @@ def test_cli_assets_materialize(mock_hasattr, parser: ArgumentParser) -> None:
     }
 
 
-def test_cli_assets_materialize_with_view_url_template(parser: ArgumentParser) -> None:
+def test_cli_assets_materialize_with_view_url_template(parser: ArgumentParser, stdout_capture) -> None:
     args = parser.parse_args(["assets", "materialize", "--name=asset1_producer", "--output=json"])
-    with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+    with stdout_capture as capture:
         asset_command.asset_materialize(args)
 
-    output = temp_stdout.getvalue()
+    output = capture.getvalue()
     run_list = json.loads(output)
     assert len(run_list) == 1
 

airflow-core/tests/unit/cli/commands/test_cheat_sheet_command.py

@@ -16,8 +16,6 @@
 # under the License.
 from __future__ import annotations
 
-import contextlib
-from io import StringIO
 from unittest import mock
 
 from airflow.cli import cli_parser
@@ -94,8 +92,8 @@ def setup_class(cls):
         cls.parser = cli_parser.get_parser()
 
     @mock.patch("airflow.cli.cli_parser.airflow_commands", MOCK_COMMANDS)
-    def test_should_display_index(self):
-        with contextlib.redirect_stdout(StringIO()) as temp_stdout:
+    def test_should_display_index(self, stdout_capture):
+        with stdout_capture as temp_stdout:
             args = self.parser.parse_args(["cheat-sheet"])
             args.func(args)
         output = temp_stdout.getvalue()