Merge pull request #85 from raman325/service_blocking

Add support for service blocking

Author: craigbarratt

custom_components/pyscript/function.py

@@ -190,13 +190,18 @@ def service_has_service(cls, domain, name):
     async def service_call(cls, domain, name, **kwargs):
         """Implement service.call()."""
         curr_task = asyncio.current_task()
-        if "context" in kwargs and isinstance(kwargs["context"], Context):
-            context = kwargs["context"]
-            del kwargs["context"]
-        else:
-            context = cls.task2context.get(curr_task, None)
-
-        await cls.hass.services.async_call(domain, name, kwargs, context=context)
+        hass_args = {}
+        for keyword, typ, default in [
+            ("context", [Context], cls.task2context.get(curr_task, None)),
+            ("blocking", [bool], None),
+            ("limit", [float, int], None),
+        ]:
+            if keyword in kwargs and type(kwargs[keyword]) in typ:
+                hass_args[keyword] = kwargs.pop(keyword)
+            elif default:
+                hass_args[keyword] = default
+
+        await cls.hass.services.async_call(domain, name, kwargs, **hass_args)
 
     @classmethod
     async def service_completions(cls, root):
@@ -255,15 +260,21 @@ def get(cls, name):
         def service_call_factory(domain, service):
             async def service_call(*args, **kwargs):
                 curr_task = asyncio.current_task()
-                if "context" in kwargs and isinstance(kwargs["context"], Context):
-                    context = kwargs["context"]
-                    del kwargs["context"]
-                else:
-                    context = cls.task2context.get(curr_task, None)
+                hass_args = {}
+                for keyword, typ, default in [
+                    ("context", [Context], cls.task2context.get(curr_task, None)),
+                    ("blocking", [bool], None),
+                    ("limit", [float, int], None),
+                ]:
+                    if keyword in kwargs and type(kwargs[keyword]) in typ:
+                        hass_args[keyword] = kwargs.pop(keyword)
+                    elif default:
+                        hass_args[keyword] = default
 
                 if len(args) != 0:
                     raise (TypeError, f"service {domain}.{service} takes no positional arguments")
-                await cls.hass.services.async_call(domain, service, kwargs, context=context)
+
+                await cls.hass.services.async_call(domain, service, kwargs, **hass_args)
 
             return service_call
 

custom_components/pyscript/state.py

@@ -258,11 +258,16 @@ async def get(cls, var_name):
             def service_call_factory(domain, service, entity_id, params):
                 async def service_call(*args, **kwargs):
                     curr_task = asyncio.current_task()
-                    if "context" in kwargs and isinstance(kwargs["context"], Context):
-                        context = kwargs["context"]
-                        del kwargs["context"]
-                    else:
-                        context = Function.task2context.get(curr_task, None)
+                    hass_args = {}
+                    for keyword, typ, default in [
+                        ("context", [Context], Function.task2context.get(curr_task, None)),
+                        ("blocking", [bool], None),
+                        ("limit", [float, int], None),
+                    ]:
+                        if keyword in kwargs and type(kwargs[keyword]) in typ:
+                            hass_args[keyword] = kwargs.pop(keyword)
+                        elif default:
+                            hass_args[keyword] = default
 
                     kwargs["entity_id"] = entity_id
                     if len(args) == 1 and len(params) == 1:
@@ -273,7 +278,7 @@ async def service_call(*args, **kwargs):
                         kwargs[param_name] = args[0]
                     elif len(args) != 0:
                         raise TypeError(f"service {domain}.{service} takes no positional arguments")
-                    await cls.hass.services.async_call(domain, service, kwargs, context=context)
+                    await cls.hass.services.async_call(domain, service, kwargs, **hass_args)
 
                 return service_call
 

docs/new_features.rst

@@ -54,6 +54,9 @@ The new features since 0.32 in master include:
   messages get displayed in VSCode. One benign but unresolved bug with VSCode is that when you connect
   to the pyscript kernel, VSCode starts a second pyscript Jupyter kernel, before shutting that second one
   down.
+- Service calls now accept ``blocking`` and ``limit`` parameters. The default behavior for a service call is
+  to run it in the background, but using ``blocking=True`` will force a task to wait up to ``limit`` seconds
+  for the service call to finish executing before continuing. Contributed by @raman325 (#85)
 
 The bug fixes since 0.32 in master include:
 

docs/reference.rst

@@ -210,6 +210,11 @@ called by:
 
        service.call("myservice", "flash_light", light_name="front", light_color="red")
 
+When making a service call, either using the ``service.call`` function or the service name as the
+function, you can optionally pass the keyword argument ``blocking=True`` if you would like to wait
+for the service to finish execution before continuing execution in your function. You can also
+specify a timeout for a blocking service call using the ``limit=<number_of_seconds>`` parameters.
+
 Firing events
 -------------
 
@@ -691,8 +696,11 @@ or ``float()``). Attributes keep their native type.
 Service Calls
 ^^^^^^^^^^^^^
 
-``service.call(domain, name, **kwargs)``
-  calls the service ``domain.name`` with the given keyword arguments as parameters.
+``service.call(domain, name, blocking=False, limit=10, **kwargs)``
+  calls the service ``domain.name`` with the given keyword arguments as parameters. If ``blocking``
+  is ``True``, pyscript will wait for the service to finish executing before continuing the current
+  routine, or will wait a maximum of the number of seconds specified in the `limit` keyword
+  argument.
 ``service.has_service(domain, name)``
   returns whether the service ``domain.name`` exists.
 

tests/test_function.py

@@ -12,6 +12,7 @@
 from pytest_homeassistant_custom_component.async_mock import MagicMock, Mock, mock_open, patch
 
 from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, EVENT_STATE_CHANGED
+from homeassistant.core import Context
 from homeassistant.setup import async_setup_component
 
 
@@ -804,3 +805,44 @@ def set_add(entity_id=None, val1=None, val2=None):
     assert literal_eval(await wait_until_done(notify_q)) == [4, "pyscript.var1", "32"]
     assert literal_eval(await wait_until_done(notify_q)) == [5, "pyscript.var1", "50", "HomeAssistant"]
     assert "TypeError: service pyscript.set_add takes no positional arguments" in caplog.text
+
+
+async def test_service_call_params(hass):
+    """Test that hass params get set properly on service calls."""
+    with patch.object(hass.services, "async_call") as call, patch.object(
+        Function, "service_has_service", return_value=True
+    ):
+        Function.init(hass)
+        await Function.service_call(
+            "test", "test", context=Context(id="test"), blocking=True, limit=1, other_service_data="test"
+        )
+        assert call.called
+        assert call.call_args[0] == ("test", "test", {"other_service_data": "test"})
+        assert call.call_args[1] == {"context": Context(id="test"), "blocking": True, "limit": 1}
+        call.reset_mock()
+
+        await Function.service_call(
+            "test", "test", context=Context(id="test"), blocking=False, other_service_data="test"
+        )
+        assert call.called
+        assert call.call_args[0] == ("test", "test", {"other_service_data": "test"})
+        assert call.call_args[1] == {"context": Context(id="test"), "blocking": False}
+        call.reset_mock()
+
+        await Function.get("test.test")(
+            context=Context(id="test"), blocking=True, limit=1, other_service_data="test"
+        )
+        assert call.called
+        assert call.call_args[0] == ("test", "test", {"other_service_data": "test"})
+        assert call.call_args[1] == {"context": Context(id="test"), "blocking": True, "limit": 1}
+        call.reset_mock()
+
+        await Function.get("test.test")(
+            context=Context(id="test"), blocking=False, other_service_data="test"
+        )
+        assert call.called
+        assert call.call_args[0] == ("test", "test", {"other_service_data": "test"})
+        assert call.call_args[1] == {"context": Context(id="test"), "blocking": False}
+
+    # Stop all tasks to avoid conflicts with other tests
+    await Function.reaper_stop()

tests/test_state.py

@@ -0,0 +1,43 @@
+"""Test pyscripts test module."""
+from custom_components.pyscript.state import State
+from pytest_homeassistant_custom_component.async_mock import patch
+
+from homeassistant.core import Context
+from homeassistant.helpers.state import State as HassState
+
+
+async def test_service_call(hass):
+    """Test calling a service using the entity_id as a property."""
+    with patch(
+        "custom_components.pyscript.state.async_get_all_descriptions",
+        return_value={
+            "test": {
+                "test": {"description": None, "fields": {"entity_id": "blah", "other_service_data": "blah"}}
+            }
+        },
+    ), patch.object(hass.states, "get", return_value=HassState("test.entity", "True")), patch.object(
+        hass.services, "async_call"
+    ) as call:
+        State.init(hass)
+        await State.get_service_params()
+
+        func = await State.get("test.entity.test")
+        await func(context=Context(id="test"), blocking=True, limit=1, other_service_data="test")
+        assert call.called
+        assert call.call_args[0] == (
+            "test",
+            "test",
+            {"other_service_data": "test", "entity_id": "test.entity"},
+        )
+        assert call.call_args[1] == {"context": Context(id="test"), "blocking": True, "limit": 1}
+        call.reset_mock()
+
+        func = await State.get("test.entity.test")
+        await func(context=Context(id="test"), blocking=False, other_service_data="test")
+        assert call.called
+        assert call.call_args[0] == (
+            "test",
+            "test",
+            {"other_service_data": "test", "entity_id": "test.entity"},
+        )
+        assert call.call_args[1] == {"context": Context(id="test"), "blocking": False}