This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7e56dac75c DatabricksHook: fix status property to work with
ClientResponse used in async mode (#43333)
7e56dac75c is described below
commit 7e56dac75cb2665e0c69bd331d7cc2f3d5056bab
Author: Luca Furrer <[email protected]>
AuthorDate: Fri Oct 25 10:04:53 2024 +0200
DatabricksHook: fix status property to work with ClientResponse used in
async mode (#43333)
* fix status property to work with ClientResponse used in async mode in
Databricks Hook
* add unittest for Databricks async api call
---
.../airflow/providers/databricks/hooks/databricks_base.py | 2 +-
providers/tests/databricks/hooks/test_databricks.py | 14 ++++++++++++++
2 files changed, 15 insertions(+), 1 deletion(-)
diff --git
a/providers/src/airflow/providers/databricks/hooks/databricks_base.py
b/providers/src/airflow/providers/databricks/hooks/databricks_base.py
index cf80b58d42..8a4a7335a4 100644
--- a/providers/src/airflow/providers/databricks/hooks/databricks_base.py
+++ b/providers/src/airflow/providers/databricks/hooks/databricks_base.py
@@ -645,7 +645,7 @@ class BaseDatabricksHook(BaseHook):
headers={**headers, **self.user_agent_header},
timeout=self.timeout_seconds,
) as response:
- self.log.debug("Response Status Code: %s",
response.status_code)
+ self.log.debug("Response Status Code: %s",
response.status)
self.log.debug("Response text: %s", response.text)
response.raise_for_status()
return await response.json()
diff --git a/providers/tests/databricks/hooks/test_databricks.py
b/providers/tests/databricks/hooks/test_databricks.py
index e2323f548d..4eaeddf972 100644
--- a/providers/tests/databricks/hooks/test_databricks.py
+++ b/providers/tests/databricks/hooks/test_databricks.py
@@ -1309,6 +1309,20 @@ class
TestDatabricksHookConnSettings(TestDatabricksHookToken):
mock_get.assert_called_once()
assert mock_get.call_args.args ==
(f"http://{HOST}:7908/api/2.1/foo/bar";,)
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
+ async def
test_async_do_api_call_only_existing_response_properties_are_read(self,
mock_get):
+ self.hook.log.setLevel("DEBUG")
+ response = mock_get.return_value.__aenter__.return_value
+ response.mock_add_spec(aiohttp.ClientResponse, spec_set=True)
+ response.json = AsyncMock(return_value={"bar": "baz"})
+ async with self.hook:
+ run_page_url = await self.hook._a_do_api_call(("GET",
"api/2.1/foo/bar"))
+
+ assert run_page_url == {"bar": "baz"}
+ mock_get.assert_called_once()
+ assert mock_get.call_args.args ==
(f"http://{HOST}:7908/api/2.1/foo/bar";,)
+
class TestRunState:
def test_is_terminal_true(self):
