Create decorator to check service permissions (#22667)

* Create decorator to check service permissions * Typing * Linting * Member comments * Linting * Member comments * Updated import * Owner comments * Linting * Linting * More work * Fixed tests * Removed service helper tests in RainMachine * Linting * Owner comments * Linting * Owner comments Co-Authored-By: bachya <bachya1208@gmail.com>
2019-04-13 13:54:29 -06:00 · 2019-04-13 13:54:29 -06:00 · fc481133e7
commit fc481133e7
parent 7a6950fd72
5 changed files with 177 additions and 118 deletions
--- a/tests/helpers/test_service.py
+++ b/tests/helpers/test_service.py
@ -38,12 +38,14 @@ def mock_entities():
        available=True,
        should_poll=False,
        supported_features=1,
+        platform='test_domain',
    )
    living_room = Mock(
        entity_id='light.living_room',
        available=True,
        should_poll=False,
        supported_features=0,
+        platform='test_domain',
    )
    entities = OrderedDict()
    entities[kitchen.entity_id] = kitchen
@ -461,3 +463,116 @@ async def test_register_admin_service(hass, hass_read_only_user,
        ))
    assert len(calls) == 1
    assert calls[0].context.user_id == hass_admin_user.id
+
+
+async def test_domain_control_not_async(hass, mock_entities):
+    """Test domain verification in a service call with an unknown user."""
+    calls = []
+
+    def mock_service_log(call):
+        """Define a protected service."""
+        calls.append(call)
+
+    with pytest.raises(exceptions.HomeAssistantError):
+        hass.helpers.service.verify_domain_control(
+            'test_domain')(mock_service_log)
+
+
+async def test_domain_control_unknown(hass, mock_entities):
+    """Test domain verification in a service call with an unknown user."""
+    calls = []
+
+    async def mock_service_log(call):
+        """Define a protected service."""
+        calls.append(call)
+
+    with patch('homeassistant.helpers.entity_registry.async_get_registry',
+               return_value=mock_coro(Mock(entities=mock_entities))):
+        protected_mock_service = hass.helpers.service.verify_domain_control(
+            'test_domain')(mock_service_log)
+
+        hass.services.async_register(
+            'test_domain', 'test_service', protected_mock_service, schema=None)
+
+        with pytest.raises(exceptions.UnknownUser):
+            await hass.services.async_call(
+                'test_domain',
+                'test_service', {},
+                blocking=True,
+                context=ha.Context(user_id='fake_user_id'))
+        assert len(calls) == 0
+
+
+async def test_domain_control_unauthorized(
+        hass, hass_read_only_user, mock_entities):
+    """Test domain verification in a service call with an unauthorized user."""
+    calls = []
+
+    async def mock_service_log(call):
+        """Define a protected service."""
+        calls.append(call)
+
+    with patch('homeassistant.helpers.entity_registry.async_get_registry',
+               return_value=mock_coro(Mock(entities=mock_entities))):
+        protected_mock_service = hass.helpers.service.verify_domain_control(
+            'test_domain')(mock_service_log)
+
+        hass.services.async_register(
+            'test_domain', 'test_service', protected_mock_service, schema=None)
+
+        with pytest.raises(exceptions.Unauthorized):
+            await hass.services.async_call(
+                'test_domain',
+                'test_service', {},
+                blocking=True,
+                context=ha.Context(user_id=hass_read_only_user.id))
+
+
+async def test_domain_control_admin(hass, hass_admin_user, mock_entities):
+    """Test domain verification in a service call with an admin user."""
+    calls = []
+
+    async def mock_service_log(call):
+        """Define a protected service."""
+        calls.append(call)
+
+    with patch('homeassistant.helpers.entity_registry.async_get_registry',
+               return_value=mock_coro(Mock(entities=mock_entities))):
+        protected_mock_service = hass.helpers.service.verify_domain_control(
+            'test_domain')(mock_service_log)
+
+        hass.services.async_register(
+            'test_domain', 'test_service', protected_mock_service, schema=None)
+
+        await hass.services.async_call(
+            'test_domain',
+            'test_service', {},
+            blocking=True,
+            context=ha.Context(user_id=hass_admin_user.id))
+
+        assert len(calls) == 1
+
+
+async def test_domain_control_no_user(hass, mock_entities):
+    """Test domain verification in a service call with no user."""
+    calls = []
+
+    async def mock_service_log(call):
+        """Define a protected service."""
+        calls.append(call)
+
+    with patch('homeassistant.helpers.entity_registry.async_get_registry',
+               return_value=mock_coro(Mock(entities=mock_entities))):
+        protected_mock_service = hass.helpers.service.verify_domain_control(
+            'test_domain')(mock_service_log)
+
+        hass.services.async_register(
+            'test_domain', 'test_service', protected_mock_service, schema=None)
+
+        await hass.services.async_call(
+            'test_domain',
+            'test_service', {},
+            blocking=True,
+            context=ha.Context(user_id=None))
+
+        assert len(calls) == 1