amoghrajesh commented on code in PR #72329:
URL: https://github.com/apache/airflow/pull/72329#discussion_r4014235929
##########
task-sdk/tests/task_sdk/definitions/test_variables.py:
##########
@@ -204,6 +204,134 @@ def test_keys_raises_on_unexpected_response_type(self,
mock_supervisor_comms):
list(results)
+class TestAsyncVariables:
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("deserialize_json", "value", "expected_value"),
+ [
+ pytest.param(False, "my_value", "my_value", id="simple-value"),
+ pytest.param(
+ True,
+ '{"key": "value", "number": 42, "flag": true}',
+ {"key": "value", "number": 42, "flag": True},
+ id="deser-object-value",
+ ),
+ ],
+ )
+ async def test_avar_get(self, deserialize_json, value, expected_value,
mock_supervisor_comms):
+ mock_supervisor_comms.send.return_value = VariableResult(key="my_key",
value=value)
+
+ var = await Variable.aget(key="my_key",
deserialize_json=deserialize_json)
+ assert var == expected_value
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("key", "value", "description", "serialize_json"),
+ [
+ pytest.param("key", "value", "description", False,
id="simple-value"),
+ pytest.param(
+ "key2",
+ {"hi": "there", "hello": 42, "flag": True},
+ "description2",
+ True,
+ id="serialize-json-value",
+ ),
+ ],
+ )
+ async def test_avar_set(self, key, value, description, serialize_json,
mock_supervisor_comms):
+ from unittest.mock import AsyncMock
Review Comment:
Just write `mock_supervisor_comms.asend.return_value = xyz` with no
reassignment and no import instead of importing and assigning to `AsyncMock`
for every test
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -373,6 +391,48 @@ def _get_variable(key: str, deserialize_json: bool) -> Any:
)
+async def _async_get_variable(key: str, deserialize_json: bool) -> Any:
Review Comment:
`ExecutionAPISecretsBackend` already has a `aget_variable` implemented.
Instead of this, we should do: `getattr(secrets_backend, "aget_connection",
None)` like how `_async_get_connection` did.
##########
task-sdk/tests/task_sdk/definitions/test_variables.py:
##########
@@ -204,6 +204,134 @@ def test_keys_raises_on_unexpected_response_type(self,
mock_supervisor_comms):
list(results)
+class TestAsyncVariables:
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("deserialize_json", "value", "expected_value"),
+ [
+ pytest.param(False, "my_value", "my_value", id="simple-value"),
+ pytest.param(
+ True,
+ '{"key": "value", "number": 42, "flag": true}',
+ {"key": "value", "number": 42, "flag": True},
+ id="deser-object-value",
+ ),
+ ],
+ )
+ async def test_avar_get(self, deserialize_json, value, expected_value,
mock_supervisor_comms):
+ mock_supervisor_comms.send.return_value = VariableResult(key="my_key",
value=value)
+
+ var = await Variable.aget(key="my_key",
deserialize_json=deserialize_json)
+ assert var == expected_value
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("key", "value", "description", "serialize_json"),
+ [
+ pytest.param("key", "value", "description", False,
id="simple-value"),
+ pytest.param(
+ "key2",
+ {"hi": "there", "hello": 42, "flag": True},
+ "description2",
+ True,
+ id="serialize-json-value",
+ ),
+ ],
+ )
+ async def test_avar_set(self, key, value, description, serialize_json,
mock_supervisor_comms):
+ from unittest.mock import AsyncMock
Review Comment:
Here and everywhere, use top level import pls. In fact import isnt even
needed, see my above comment.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
