Add missing conext preservation to bayesian and universal (#40178)

We already do this for template sensors, but it was
missing for bayesian and universal

homeassistant/components/bayesian/binary_sensor.py

@@ -182,6 +182,7 @@ class BayesianBinarySensor(BinarySensorEntity):
             entity = event.data.get("entity_id")
 
             self.current_observations.update(self._record_entity_observations(entity))
+            self.async_set_context(event.context)
             self._recalculate_and_write_state()
 
         self.async_on_remove(
@@ -220,6 +221,8 @@ class BayesianBinarySensor(BinarySensorEntity):
                     obs_entry = None
                 self.current_observations[obs["id"]] = obs_entry
 
+            if event:
+                self.async_set_context(event.context)
             self._recalculate_and_write_state()
 
         for template in self.observations_by_template:

homeassistant/components/universal/media_player.py

@@ -145,8 +145,9 @@ class UniversalMediaPlayer(MediaPlayerEntity):
         """Subscribe to children and template state changes."""
 
         @callback
-        def _async_on_dependency_update(*_):
+        def _async_on_dependency_update(event):
             """Update ha state when dependencies update."""
+            self.async_set_context(event.context)
             self.async_schedule_update_ha_state(True)
 
         @callback
@@ -158,6 +159,10 @@ class UniversalMediaPlayer(MediaPlayerEntity):
                 self._state_template_result = None
             else:
                 self._state_template_result = result
+
+            if event:
+                self.async_set_context(event.context)
+
             self.async_schedule_update_ha_state(True)
 
         if self._state_template is not None:

tests/components/bayesian/test_binary_sensor.py

@@ -9,7 +9,14 @@ from homeassistant.components.homeassistant import (
     DOMAIN as HA_DOMAIN,
     SERVICE_UPDATE_ENTITY,
 )
-from homeassistant.const import ATTR_ENTITY_ID, SERVICE_RELOAD, STATE_UNKNOWN
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    SERVICE_RELOAD,
+    STATE_OFF,
+    STATE_ON,
+    STATE_UNKNOWN,
+)
+from homeassistant.core import Context, callback
 from homeassistant.setup import async_setup_component, setup_component
 
 from tests.async_mock import patch
@@ -686,3 +693,79 @@ async def test_reload(hass):
 
 def _get_fixtures_base_path():
     return path.dirname(path.dirname(path.dirname(__file__)))
+
+
+async def test_template_triggers(hass):
+    """Test sensor with template triggers."""
+    hass.states.async_set("input_boolean.test", STATE_OFF)
+    config = {
+        "binary_sensor": {
+            "name": "Test_Binary",
+            "platform": "bayesian",
+            "observations": [
+                {
+                    "platform": "template",
+                    "value_template": "{{ states.input_boolean.test.state }}",
+                    "prob_given_true": 1999.9,
+                },
+            ],
+            "prior": 0.2,
+            "probability_threshold": 0.32,
+        }
+    }
+
+    await async_setup_component(hass, "binary_sensor", config)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("binary_sensor.test_binary").state == STATE_OFF
+
+    events = []
+    hass.helpers.event.async_track_state_change_event(
+        "binary_sensor.test_binary", callback(lambda event: events.append(event))
+    )
+
+    context = Context()
+    hass.states.async_set("input_boolean.test", STATE_ON, context=context)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    assert events[0].context == context
+
+
+async def test_state_triggers(hass):
+    """Test sensor with state triggers."""
+    hass.states.async_set("sensor.test_monitored", STATE_OFF)
+
+    config = {
+        "binary_sensor": {
+            "name": "Test_Binary",
+            "platform": "bayesian",
+            "observations": [
+                {
+                    "platform": "state",
+                    "entity_id": "sensor.test_monitored",
+                    "to_state": "off",
+                    "prob_given_true": 999.9,
+                    "prob_given_false": 999.4,
+                },
+            ],
+            "prior": 0.2,
+            "probability_threshold": 0.32,
+        }
+    }
+    await async_setup_component(hass, "binary_sensor", config)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("binary_sensor.test_binary").state == STATE_OFF
+
+    events = []
+    hass.helpers.event.async_track_state_change_event(
+        "binary_sensor.test_binary", callback(lambda event: events.append(event))
+    )
+
+    context = Context()
+    hass.states.async_set("sensor.test_monitored", STATE_ON, context=context)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    assert events[0].context == context

tests/components/universal/test_media_player.py

@@ -20,6 +20,7 @@ from homeassistant.const import (
     STATE_PLAYING,
     STATE_UNKNOWN,
 )
+from homeassistant.core import Context, callback
 from homeassistant.setup import async_setup_component
 
 from tests.async_mock import patch
@@ -812,10 +813,18 @@ async def test_master_state_with_template(hass):
     await hass.async_block_till_done()
     hass.states.get("media_player.tv").state == STATE_ON
 
-    hass.states.async_set("input_boolean.test", STATE_ON)
+    events = []
+
+    hass.helpers.event.async_track_state_change_event(
+        "media_player.tv", callback(lambda event: events.append(event))
+    )
+
+    context = Context()
+    hass.states.async_set("input_boolean.test", STATE_ON, context=context)
     await hass.async_block_till_done()
 
     hass.states.get("media_player.tv").state == STATE_OFF
+    assert events[0].context == context
 
 
 async def test_reload(hass):