🎨 [pre-commit.ci] Auto format from pre-commit.com hooks

astronomer · Dec 2, 2023 · f4f9b1c · f4f9b1c
1 parent 2ae1e25
commit f4f9b1c
Showing 1 changed file with 14 additions and 23 deletions.
diff --git a/tests/plugin/test_plugin.py b/tests/plugin/test_plugin.py
@@ -1,16 +1,19 @@
 from unittest.mock import patch
 
 import pytest
+from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException
 from airflow.utils.db import initdb, resetdb
 from airflow.www.app import cached_app
 from airflow.www.extensions.init_appbuilder import AirflowAppBuilder
-from airflow.configuration import conf
 from flask.testing import FlaskClient
 
 from cosmos.plugin import dbt_docs_view, iframe_script
 
 
+original_conf_get = conf.get
+
+
 @pytest.fixture(scope="module")
 def app() -> FlaskClient:
     initdb()
@@ -36,58 +39,46 @@ def conf_get(section, key, *args, **kwargs):
         if section == "cosmos" and key == "dbt_docs_dir":
             return "path/to/docs/dir"
         else:
-            return conf.get(section, key, *args, **kwargs)
+            return original_conf_get(section, key, *args, **kwargs)
 
-    monkeypatch.setattr(
-        "cosmos.plugin.conf",
-        "get",
-        conf_get
-    )
+    monkeypatch.setattr(conf, "get", conf_get)
 
     mock_open_file.return_value = "<head></head><body></body>"
 
     response = app.get("/cosmos/dbt_docs")
-    assert mock_open_file.assert_called_once()
+
+    mock_open_file.assert_called_once()
     assert response.status_code == 200
     assert iframe_script in response.text
 
 
 @patch("cosmos.plugin.open_file")
 @pytest.mark.parametrize("artifact", ["dbt_docs_index.html", "manifest.json", "catalog.json"])
 def test_dbt_docs_artifact(mock_open_file, monkeypatch, app, artifact):
-
     def conf_get(section, key, *args, **kwargs):
         if section == "cosmos" and key == "dbt_docs_dir":
             return "path/to/docs/dir"
         else:
-            return conf.get(section, key, *args, **kwargs)
+            return original_conf_get(section, key, *args, **kwargs)
 
-    monkeypatch.setattr(
-        "cosmos.plugin.conf",
-        "get",
-        conf_get
-    )
+    monkeypatch.setattr(conf, "get", conf_get)
 
     mock_open_file.return_value = ""
     response = app.get(f"/cosmos/{artifact}")
-    assert mock_open_file.assert_called_once()
+
+    mock_open_file.assert_called_once()
     assert response.status_code == 200
 
 
 @pytest.mark.parametrize("artifact", ["dbt_docs_index.html", "manifest.json", "catalog.json"])
 def test_dbt_docs_artifact_missing(app, artifact, monkeypatch):
-
     def conf_get(section, key, *args, **kwargs):
         if section == "cosmos":
             raise AirflowConfigException
         else:
-            return conf.get(section, key, *args, **kwargs)
+            return original_conf_get(section, key, *args, **kwargs)
 
-    monkeypatch.setattr(
-        "cosmos.plugin.conf",
-        "get",
-        conf_get
-    )
+    monkeypatch.setattr(conf, "get", conf_get)
 
     response = app.get(f"/cosmos/{artifact}")
     assert response.status_code == 404