astronomer · pankajastro · Jan 30, 2024 · Jan 29, 2024 · Jan 29, 2024 · Jan 30, 2024
@@ -1,4 +1,5 @@
 """Astronomer migration DAG to transform metadata from source deployment to target Astro Cloud deployment."""
+
 from datetime import datetime
 
 from airflow import DAG

@@ -1,4 +1,5 @@
 """Master Dag to run all the example dags."""
+
 import logging
 import os
 import time

@@ -1,4 +1,5 @@
 """This script fetches the latest runtime image tag from the provided Quay.io repository URL."""
+
 from __future__ import annotations
 
 import sys

@@ -44,7 +44,7 @@ repos:
         args: ["--config=mlc-config.json"]
 
   - repo: https://github.com/psf/black
-    rev: 23.12.1
+    rev: 24.1.1
     hooks:
       - id: black
 
@@ -90,7 +90,7 @@ repos:
         additional_dependencies: ['toml']
 
   - repo: https://github.com/PyCQA/bandit
-    rev: 1.7.6
+    rev: 1.7.7
     hooks:
       - id: bandit
         # Excluding Example DAGs temporarily until they can be fixed

@@ -1,4 +1,5 @@
 """DAG to nuke AWS resources."""
+
 import logging
 import os
 from datetime import datetime, timedelta

@@ -1,4 +1,5 @@
 """Example DAG for the  AWS Batch Operator Async"""
+
 import logging
 import os
 import time

@@ -1,4 +1,5 @@
 """Airflow operators example to manage AWS Redshift cluster."""
+
 import os
 from datetime import datetime, timedelta
 from typing import Any

@@ -1,4 +1,5 @@
 """This module contains async AWS Base Hook for deferrable operators and sensors."""
+
 from __future__ import annotations
 
 from typing import Any

@@ -7,6 +7,7 @@
     - `Batch <http://boto3.readthedocs.io/en/latest/reference/services/batch.html>`_
     - `Welcome <https://docs.aws.amazon.com/batch/latest/APIReference/Welcome.html>`_
 """
+
 from __future__ import annotations
 
 import warnings

@@ -1,4 +1,5 @@
 """This is an example dag for hive partition sensors."""
+
 import logging
 import os
 import time

@@ -1,4 +1,5 @@
 """This module contains the Apache HiveCli hook async."""
+
 import asyncio
 from typing import Tuple
 

@@ -3,6 +3,7 @@
 The tasks below trigger the computation of pi on the Spark instance
 using the Java and Python executables provided in the example library.
 """
+
 import logging
 import os
 import time

@@ -1,4 +1,5 @@
 """This module contains the Apache Livy hook async."""
+
 import asyncio
 import re
 import warnings

@@ -1,4 +1,5 @@
 """This module contains the Apache Livy operator async."""
+
 from __future__ import annotations
 
 import warnings

@@ -1,4 +1,5 @@
 """This module contains the Apache Livy Trigger."""
+
 import asyncio
 import warnings
 from typing import Any, AsyncIterator, Dict, Optional, Tuple, Union

@@ -1,4 +1,5 @@
 """Example DAG demonstrating the usage of the KubernetesPodOperatorAsync."""
+
 import os
 from datetime import datetime, timedelta
 

@@ -1,4 +1,5 @@
 """Example Airflow DAG which uses impersonation parameters for authenticating with Google BigQuery service."""
+
 import os
 from datetime import datetime, timedelta
 from typing import Any

@@ -2,6 +2,7 @@
 Example Airflow DAG for Google BigQuery service.
 Uses Async version of BigQueryInsertJobOperator and BigQueryCheckOperator.
 """
+
 import os
 from datetime import datetime, timedelta
 from typing import Any

@@ -1,4 +1,5 @@
 """Example Airflow DAG for Google BigQuery Sensors."""
+
 import os
 from datetime import datetime, timedelta
 from typing import Any

@@ -1,4 +1,5 @@
 """DAG to delete stale GCP resources."""
+
 import os
 from datetime import timedelta
 

@@ -1,4 +1,5 @@
 """Example Airflow DAG for Google Cloud Storage operators."""
+
 import os
 from datetime import datetime, timedelta
 

@@ -9,7 +9,7 @@
 from airflow.exceptions import AirflowException
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from airflow.utils.process_utils import execute_in_subprocess, patch_environ
-from google.auth import impersonated_credentials  # type: ignore[attr-defined]
+from google.auth import impersonated_credentials
 from google.cloud.container_v1 import ClusterManagerClient
 from google.oauth2.service_account import Credentials
 from kubernetes_asyncio.config.kube_config import KubeConfigLoader
@@ -76,9 +76,9 @@ def _get_gke_config_file(
 
     # Write config to a temp file and set the environment variable to point to it.
     # This is to avoid race conditions of reading/writing a single file
-    with tempfile.NamedTemporaryFile() as conf_file, patch_environ(
-        {KUBE_CONFIG_ENV_VAR: conf_file.name}
-    ), hook.provide_authorized_gcloud():
+    with tempfile.NamedTemporaryFile() as conf_file:
+        with patch_environ({KUBE_CONFIG_ENV_VAR: conf_file.name}):
+            hook.provide_authorized_gcloud()
         # Attempt to get/update credentials
         # We call gcloud directly instead of using google-cloud-python api
         # because there is no way to write kubernetes config to a file, which is

@@ -1,4 +1,5 @@
 """This module contains a Google Cloud Storage hook."""
+
 import warnings
 from typing import Any, cast
 

@@ -1,4 +1,5 @@
 """This module contains Google BigQueryAsync providers."""
+
 from __future__ import annotations
 
 import warnings

@@ -1,4 +1,5 @@
 """This module contains Google Dataproc operators."""
+
 from __future__ import annotations
 
 import warnings

@@ -1,4 +1,5 @@
 """This module contains Google GKE operators."""
+
 from __future__ import annotations
 
 from typing import Any, Sequence

@@ -1,4 +1,5 @@
 """This module contains Google Big Query sensors."""
+
 from __future__ import annotations
 
 import warnings

@@ -1,4 +1,5 @@
 """This module contains Google Cloud Storage sensors."""
+
 from __future__ import annotations
 
 import warnings

@@ -1,4 +1,5 @@
 """This module contains the Azure Data Factory hook's asynchronous implementation."""
+
 from __future__ import annotations
 
 import inspect

@@ -1,4 +1,5 @@
 """This module contains the Azure WASB hook's asynchronous implementation."""
+
 from __future__ import annotations
 
 import warnings

@@ -65,9 +65,11 @@ def parse_providers_release_testing_gh_issue(issue_url: str) -> list[str]:
     h2_titles = first_comment.find_all("h2")
     package_urls = [_parse_pypi_url_from_h2_title(h2_title) for h2_title in h2_titles]
     pinned_packages = [
-        _parse_pinned_package_from_pypi_url(url)
-        if "apache-airflow-providers-cncf-kubernetes" not in url
-        else _parse_pinned_pacakge_from_cncf(url)
+        (
+            _parse_pinned_package_from_pypi_url(url)
+            if "apache-airflow-providers-cncf-kubernetes" not in url
+            else _parse_pinned_pacakge_from_cncf(url)
+        )
         for url in package_urls
     ]
     return pinned_packages

@@ -11,28 +11,24 @@
 
 class TestHttpHookAsync:
     @pytest.mark.asyncio
-    async def test_do_api_call_async_non_retryable_error(self, aioresponse):
+    async def test_do_api_call_async_non_retryable_error(self, aioresponse, monkeypatch):
         hook = HttpHookAsync(method="GET")
         aioresponse.get("http://httpbin.org/non_existent_endpoint", status=400)
 
-        with pytest.raises(AirflowException) as exc, mock.patch.dict(
-            "os.environ",
-            AIRFLOW_CONN_HTTP_DEFAULT="http://httpbin.org/",
-        ):
+        monkeypatch.setenv("AIRFLOW_CONN_HTTP_DEFAULT", "http://httpbin.org/")
+        with pytest.raises(AirflowException) as exc:
             await hook.run(endpoint="non_existent_endpoint")
 
         assert str(exc.value) == "400:Bad Request"
 
     @pytest.mark.asyncio
-    async def test_do_api_call_async_retryable_error(self, caplog, aioresponse):
+    async def test_do_api_call_async_retryable_error(self, caplog, aioresponse, monkeypatch):
         caplog.set_level(logging.WARNING, logger="astronomer.providers.http.hooks.http")
         hook = HttpHookAsync(method="GET")
         aioresponse.get("http://httpbin.org/non_existent_endpoint", status=500, repeat=True)
 
-        with pytest.raises(AirflowException) as exc, mock.patch.dict(
-            "os.environ",
-            AIRFLOW_CONN_HTTP_DEFAULT="http://httpbin.org/",
-        ):
+        monkeypatch.setenv("AIRFLOW_CONN_HTTP_DEFAULT", "http://httpbin.org/")
+        with pytest.raises(AirflowException) as exc:
             await hook.run(endpoint="non_existent_endpoint")
 
         assert str(exc.value) == "500:Internal Server Error"
@@ -72,7 +68,8 @@ async def test_post_request(self, aioresponse):
         )
 
         with mock.patch(
-            "airflow.hooks.base.BaseHook.get_connection", side_effect=self.get_airflow_connection
+            "airflow.hooks.base.BaseHook.get_connection",
+            side_effect=self.get_airflow_connection,
         ):
             resp = await hook.run("v1/test")
             assert resp.status == 200
@@ -90,7 +87,8 @@ async def test_post_request_and_get_json_without_keep_response(self, aioresponse
         )
 
         with mock.patch(
-            "airflow.hooks.base.BaseHook.get_connection", side_effect=self.get_airflow_connection
+            "airflow.hooks.base.BaseHook.get_connection",
+            side_effect=self.get_airflow_connection,
         ):
             resp = await hook.run("v1/test")
             with pytest.raises(ClientConnectionError, match="Connection closed"):
@@ -109,7 +107,8 @@ async def test_post_request_and_get_json_with_keep_response(self, aioresponse):
         )
 
         with mock.patch(
-            "airflow.hooks.base.BaseHook.get_connection", side_effect=self.get_airflow_connection
+            "airflow.hooks.base.BaseHook.get_connection",
+            side_effect=self.get_airflow_connection,
         ):
             resp = await hook.run("v1/test")
             resp_payload = await resp.json()
@@ -128,7 +127,8 @@ async def test_post_request_with_error_code(self, aioresponse):
         )
 
         with mock.patch(
-            "airflow.hooks.base.BaseHook.get_connection", side_effect=self.get_airflow_connection
+            "airflow.hooks.base.BaseHook.get_connection",
+            side_effect=self.get_airflow_connection,
         ):
             with pytest.raises(AirflowException):
                 await hook.run("v1/test")