Add context to event trigger (#40932)

homeassistant/components/homeassistant/triggers/event.py

@@ -11,6 +11,7 @@ from homeassistant.helpers import config_validation as cv
 
 CONF_EVENT_TYPE = "event_type"
 CONF_EVENT_DATA = "event_data"
+CONF_EVENT_CONTEXT = "context"
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -19,35 +20,41 @@ TRIGGER_SCHEMA = vol.Schema(
         vol.Required(CONF_PLATFORM): "event",
         vol.Required(CONF_EVENT_TYPE): cv.string,
         vol.Optional(CONF_EVENT_DATA): dict,
+        vol.Optional(CONF_EVENT_CONTEXT): dict,
     }
 )
 
 
+def _populate_schema(config, config_parameter):
+    if config_parameter not in config:
+        return None
+
+    return vol.Schema(
+        {vol.Required(key): value for key, value in config[config_parameter].items()},
+        extra=vol.ALLOW_EXTRA,
+    )
+
+
 async def async_attach_trigger(
     hass, config, action, automation_info, *, platform_type="event"
 ):
     """Listen for events based on configuration."""
     event_type = config.get(CONF_EVENT_TYPE)
-    event_data_schema = None
+    event_data_schema = _populate_schema(config, CONF_EVENT_DATA)
-    if config.get(CONF_EVENT_DATA):
+    event_context_schema = _populate_schema(config, CONF_EVENT_CONTEXT)
-        event_data_schema = vol.Schema(
-            {
-                vol.Required(key): value
-                for key, value in config.get(CONF_EVENT_DATA).items()
-            },
-            extra=vol.ALLOW_EXTRA,
-        )
 
     @callback
     def handle_event(event):
         """Listen for events and calls the action when data matches."""
-        if event_data_schema:
-            # Check that the event data matches the configured
-            # schema if one was provided
         try:
+            # Check that the event data and context match the configured
+            # schema if one was provided
+            if event_data_schema:
                 event_data_schema(event.data)
+            if event_context_schema:
+                event_context_schema(event.context.as_dict())
         except vol.Invalid:
-                # If event data doesn't match requested schema, skip event
+            # If event doesn't match, skip event
             return
 
         hass.async_run_job(

tests/components/homeassistant/triggers/test_event.py

@@ -15,6 +15,12 @@ def calls(hass):
     return async_mock_service(hass, "test", "automation")
 
 
+@pytest.fixture
+def context_with_user():
+    """Track calls to a mock service."""
+    return Context(user_id="test_user_id")
+
+
 @pytest.fixture(autouse=True)
 def setup_comp(hass):
     """Initialize components."""
@@ -53,8 +59,8 @@ async def test_if_fires_on_event(hass, calls):
     assert len(calls) == 1
 
 
-async def test_if_fires_on_event_extra_data(hass, calls):
+async def test_if_fires_on_event_extra_data(hass, calls, context_with_user):
-    """Test the firing of events still matches with event data."""
+    """Test the firing of events still matches with event data and context."""
     assert await async_setup_component(
         hass,
         automation.DOMAIN,
@@ -65,8 +71,9 @@ async def test_if_fires_on_event_extra_data(hass, calls):
             }
         },
     )
-
+    hass.bus.async_fire(
-    hass.bus.async_fire("test_event", {"extra_key": "extra_data"})
+        "test_event", {"extra_key": "extra_data"}, context=context_with_user
+    )
     await hass.async_block_till_done()
     assert len(calls) == 1
 
@@ -82,8 +89,8 @@ async def test_if_fires_on_event_extra_data(hass, calls):
     assert len(calls) == 1
 
 
-async def test_if_fires_on_event_with_data(hass, calls):
+async def test_if_fires_on_event_with_data_and_context(hass, calls, context_with_user):
-    """Test the firing of events with data."""
+    """Test the firing of events with data and context."""
     assert await async_setup_component(
         hass,
         automation.DOMAIN,
@@ -96,6 +103,7 @@ async def test_if_fires_on_event_with_data(hass, calls):
                         "some_attr": "some_value",
                         "second_attr": "second_value",
                     },
+                    "context": {"user_id": context_with_user.user_id},
                 },
                 "action": {"service": "test.automation"},
             }
@@ -105,17 +113,31 @@ async def test_if_fires_on_event_with_data(hass, calls):
     hass.bus.async_fire(
         "test_event",
         {"some_attr": "some_value", "another": "value", "second_attr": "second_value"},
+        context=context_with_user,
     )
     await hass.async_block_till_done()
     assert len(calls) == 1
 
-    hass.bus.async_fire("test_event", {"some_attr": "some_value", "another": "value"})
+    hass.bus.async_fire(
+        "test_event",
+        {"some_attr": "some_value", "another": "value"},
+        context=context_with_user,
+    )
     await hass.async_block_till_done()
     assert len(calls) == 1  # No new call
 
+    hass.bus.async_fire(
+        "test_event",
+        {"some_attr": "some_value", "another": "value", "second_attr": "second_value"},
+    )
+    await hass.async_block_till_done()
+    assert len(calls) == 1
 
-async def test_if_fires_on_event_with_empty_data_config(hass, calls):
+
-    """Test the firing of events with empty data config.
+async def test_if_fires_on_event_with_empty_data_and_context_config(
+    hass, calls, context_with_user
+):
+    """Test the firing of events with empty data and context config.
 
     The frontend automation editor can produce configurations with an
     empty dict for event_data instead of no key.
@@ -129,13 +151,18 @@ async def test_if_fires_on_event_with_empty_data_config(hass, calls):
                     "platform": "event",
                     "event_type": "test_event",
                     "event_data": {},
+                    "context": {},
                 },
                 "action": {"service": "test.automation"},
             }
         },
     )
 
-    hass.bus.async_fire("test_event", {"some_attr": "some_value", "another": "value"})
+    hass.bus.async_fire(
+        "test_event",
+        {"some_attr": "some_value", "another": "value"},
+        context=context_with_user,
+    )
     await hass.async_block_till_done()
     assert len(calls) == 1
 
@@ -165,7 +192,7 @@ async def test_if_fires_on_event_with_nested_data(hass, calls):
 
 
 async def test_if_not_fires_if_event_data_not_matches(hass, calls):
-    """Test firing of event if no match."""
+    """Test firing of event if no data match."""
     assert await async_setup_component(
         hass,
         automation.DOMAIN,
@@ -184,3 +211,27 @@ async def test_if_not_fires_if_event_data_not_matches(hass, calls):
     hass.bus.async_fire("test_event", {"some_attr": "some_other_value"})
     await hass.async_block_till_done()
     assert len(calls) == 0
+
+
+async def test_if_not_fires_if_event_context_not_matches(
+    hass, calls, context_with_user
+):
+    """Test firing of event if no context match."""
+    assert await async_setup_component(
+        hass,
+        automation.DOMAIN,
+        {
+            automation.DOMAIN: {
+                "trigger": {
+                    "platform": "event",
+                    "event_type": "test_event",
+                    "context": {"user_id": "some_user"},
+                },
+                "action": {"service": "test.automation"},
+            }
+        },
+    )
+
+    hass.bus.async_fire("test_event", {}, context=context_with_user)
+    await hass.async_block_till_done()
+    assert len(calls) == 0