refactor(databricks): make _retryable_error_async a static method and…

… add test
astronomer · Oct 30, 2023 · aa2f239 · aa2f239
1 parent 6d8ed2d
commit aa2f239
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/astronomer/providers/databricks/hooks/databricks.py b/astronomer/providers/databricks/hooks/databricks.py
@@ -152,7 +152,8 @@ async def _do_api_call_async(
                 attempt_num += 1
                 await asyncio.sleep(self.retry_delay)
 
-    def _retryable_error_async(self, exception: ClientConnectorError | ClientResponseError) -> bool:
+    @staticmethod
+    def _retryable_error_async(exception: ClientConnectorError | ClientResponseError) -> bool:
         """
         Determines whether or not an exception that was thrown might be successful
         on a subsequent attempt.

diff --git a/tests/databricks/hooks/test_databricks.py b/tests/databricks/hooks/test_databricks.py
@@ -3,6 +3,7 @@
 from unittest.mock import MagicMock
 
 import pytest
+from aiohttp import ClientConnectorError
 from airflow import __version__ as provider_version
 from airflow.exceptions import AirflowException
 from airflow.providers.databricks.hooks.databricks import (
@@ -291,3 +292,7 @@ async def test_get_run_output_response(self, mock_do_api_async):
         mock_do_api_async.return_value = MOCK_GET_OUTPUT_RESPONSE
         run_output = await hook.get_run_output_response(RUN_ID)
         assert run_output == MOCK_GET_OUTPUT_RESPONSE
+
+    def test___retryable_error_async_with_client_connector_error(self):
+        exception = ClientConnectorError(connection_key="", os_error=OSError())
+        assert DatabricksHookAsync._retryable_error_async(exception) is True