Speed up single entity/response service calls (#96729)

* Significantly speed up single entity/response service calls

Since the majority of service calls are single entity, we can
avoid creating tasks in this case. Since the multi-entity
service calls always check the result and raise, we can switch
the asyncio.wait to asyncio.gather

* Significantly speed up single entity/response service calls

Since the majority of service calls are single entity, we can
avoid creating tasks in this case. Since the multi-entity
service calls always check the result and raise, we can switch
the asyncio.wait to asyncio.gather

* revert

* cannot be inside pytest.raises

* one more

* Update homeassistant/helpers/service.py

homeassistant/helpers/service.py

@@ -741,6 +741,8 @@ async def entity_service_call(  # noqa: C901
     Calls all platforms simultaneously.
     """
     entity_perms: None | (Callable[[str, str], bool]) = None
+    return_response = call.return_response
+
     if call.context.user_id:
         user = await hass.auth.async_get_user(call.context.user_id)
         if user is None:
@@ -851,13 +853,27 @@ async def entity_service_call(  # noqa: C901
         entities.append(entity)
 
     if not entities:
-        if call.return_response:
+        if return_response:
             raise HomeAssistantError(
                 "Service call requested response data but did not match any entities"
             )
         return None
 
-    if call.return_response and len(entities) != 1:
+    if len(entities) == 1:
+        # Single entity case avoids creating tasks and allows returning
+        # ServiceResponse
+        entity = entities[0]
+        response_data = await _handle_entity_call(
+            hass, entity, func, data, call.context
+        )
+        if entity.should_poll:
+            # Context expires if the turn on commands took a long time.
+            # Set context again so it's there when we update
+            entity.async_set_context(call.context)
+            await entity.async_update_ha_state(True)
+        return response_data if return_response else None
+
+    if return_response:
         raise HomeAssistantError(
             "Service call requested response data but matched more than one entity"
         )
@@ -874,9 +890,8 @@ async def entity_service_call(  # noqa: C901
     )
     assert not pending
 
-    response_data: ServiceResponse | None
     for task in done:
-        response_data = task.result()  # pop exception if have
+        task.result()  # pop exception if have
 
     tasks: list[asyncio.Task[None]] = []
 
@@ -895,7 +910,7 @@ async def entity_service_call(  # noqa: C901
         for future in done:
             future.result()  # pop exception if have
 
-    return response_data if call.return_response else None
+    return None
 
 
 async def _handle_entity_call(

tests/components/devolo_home_network/test_button.py

@@ -228,7 +228,8 @@ async def test_auth_failed(hass: HomeAssistant, mock_device: MockDevice) -> None
             {ATTR_ENTITY_ID: state_key},
             blocking=True,
         )
-        await hass.async_block_till_done()
+
+    await hass.async_block_till_done()
     flows = hass.config_entries.flow.async_progress()
     assert len(flows) == 1
 

tests/components/devolo_home_network/test_switch.py

@@ -307,6 +307,9 @@ async def test_auth_failed(
         await hass.services.async_call(
             PLATFORM, SERVICE_TURN_ON, {"entity_id": state_key}, blocking=True
         )
+
+    await hass.async_block_till_done()
+
     flows = hass.config_entries.flow.async_progress()
     assert len(flows) == 1
 

tests/components/homeassistant/triggers/test_time.py

@@ -102,6 +102,7 @@ async def test_if_fires_using_at_input_datetime(
         },
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     time_that_will_not_match_right_away = trigger_dt - timedelta(minutes=1)
 
@@ -148,6 +149,7 @@ async def test_if_fires_using_at_input_datetime(
         },
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     async_fire_time_changed(hass, trigger_dt + timedelta(seconds=1))
     await hass.async_block_till_done()
@@ -556,6 +558,7 @@ async def test_datetime_in_past_on_load(hass: HomeAssistant, calls) -> None:
         },
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     assert await async_setup_component(
         hass,
@@ -587,6 +590,7 @@ async def test_datetime_in_past_on_load(hass: HomeAssistant, calls) -> None:
         },
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     async_fire_time_changed(hass, future + timedelta(seconds=1))
     await hass.async_block_till_done()

tests/components/rflink/test_light.py

@@ -285,11 +285,13 @@ async def test_signal_repetitions_cancelling(hass: HomeAssistant, monkeypatch) -
     await hass.services.async_call(
         DOMAIN, SERVICE_TURN_OFF, {ATTR_ENTITY_ID: f"{DOMAIN}.test"}
     )
+
     # Get background service time to start running
     await asyncio.sleep(0)
     await hass.services.async_call(
         DOMAIN, SERVICE_TURN_ON, {ATTR_ENTITY_ID: f"{DOMAIN}.test"}, blocking=True
     )
+    await hass.async_block_till_done()
 
     assert [call[0][1] for call in protocol.send_command_ack.call_args_list] == [
         "off",

tests/components/shelly/test_climate.py

@@ -427,6 +427,7 @@ async def test_block_set_mode_auth_error(
         {ATTR_ENTITY_ID: ENTITY_ID, ATTR_HVAC_MODE: HVACMode.HEAT},
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     assert entry.state == ConfigEntryState.LOADED
 

tests/components/shelly/test_number.py

@@ -186,6 +186,7 @@ async def test_block_set_value_auth_error(
         {ATTR_ENTITY_ID: "number.test_name_valve_position", ATTR_VALUE: 30},
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     assert entry.state == ConfigEntryState.LOADED
 

tests/components/shelly/test_switch.py

@@ -82,6 +82,7 @@ async def test_block_set_state_auth_error(
         {ATTR_ENTITY_ID: "switch.test_name_channel_1"},
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     assert entry.state == ConfigEntryState.LOADED
 
@@ -211,6 +212,7 @@ async def test_rpc_auth_error(
         {ATTR_ENTITY_ID: "switch.test_switch_0"},
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     assert entry.state == ConfigEntryState.LOADED
 

tests/components/shelly/test_update.py

@@ -203,6 +203,7 @@ async def test_block_update_auth_error(
         {ATTR_ENTITY_ID: "update.test_name_firmware_update"},
         blocking=True,
     )
+    await hass.async_block_till_done()
 
     assert entry.state == ConfigEntryState.LOADED
 
@@ -541,6 +542,7 @@ async def test_rpc_update_auth_error(
         blocking=True,
     )
 
+    await hass.async_block_till_done()
     assert entry.state == ConfigEntryState.LOADED
 
     flows = hass.config_entries.flow.async_progress()

tests/components/totalconnect/test_alarm_control_panel.py

@@ -129,7 +129,7 @@ async def test_arm_home_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 ALARM_DOMAIN, SERVICE_ALARM_ARM_HOME, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect failed to arm home test."
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_DISARMED
         assert mock_request.call_count == 2
@@ -139,7 +139,7 @@ async def test_arm_home_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 ALARM_DOMAIN, SERVICE_ALARM_ARM_HOME, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect usercode is invalid. Did not arm home"
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_DISARMED
         # should have started a re-auth flow
@@ -183,7 +183,7 @@ async def test_arm_home_instant_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 DOMAIN, SERVICE_ALARM_ARM_HOME_INSTANT, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect failed to arm home instant test."
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_DISARMED
         assert mock_request.call_count == 2
@@ -193,7 +193,7 @@ async def test_arm_home_instant_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 DOMAIN, SERVICE_ALARM_ARM_HOME_INSTANT, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert (
             f"{err.value}"
             == "TotalConnect usercode is invalid. Did not arm home instant"
@@ -240,7 +240,7 @@ async def test_arm_away_instant_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 DOMAIN, SERVICE_ALARM_ARM_AWAY_INSTANT, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect failed to arm away instant test."
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_DISARMED
         assert mock_request.call_count == 2
@@ -250,7 +250,7 @@ async def test_arm_away_instant_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 DOMAIN, SERVICE_ALARM_ARM_AWAY_INSTANT, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert (
             f"{err.value}"
             == "TotalConnect usercode is invalid. Did not arm away instant"
@@ -296,7 +296,7 @@ async def test_arm_away_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 ALARM_DOMAIN, SERVICE_ALARM_ARM_AWAY, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect failed to arm away test."
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_DISARMED
         assert mock_request.call_count == 2
@@ -306,7 +306,7 @@ async def test_arm_away_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 ALARM_DOMAIN, SERVICE_ALARM_ARM_AWAY, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect usercode is invalid. Did not arm away"
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_DISARMED
         # should have started a re-auth flow
@@ -353,7 +353,7 @@ async def test_disarm_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 ALARM_DOMAIN, SERVICE_ALARM_DISARM, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect failed to disarm test."
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_ARMED_AWAY
         assert mock_request.call_count == 2
@@ -363,7 +363,7 @@ async def test_disarm_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 ALARM_DOMAIN, SERVICE_ALARM_DISARM, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect usercode is invalid. Did not disarm"
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_ARMED_AWAY
         # should have started a re-auth flow
@@ -406,7 +406,7 @@ async def test_arm_night_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 ALARM_DOMAIN, SERVICE_ALARM_ARM_NIGHT, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect failed to arm night test."
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_DISARMED
         assert mock_request.call_count == 2
@@ -416,7 +416,7 @@ async def test_arm_night_failure(hass: HomeAssistant) -> None:
             await hass.services.async_call(
                 ALARM_DOMAIN, SERVICE_ALARM_ARM_NIGHT, DATA, blocking=True
             )
-            await hass.async_block_till_done()
+        await hass.async_block_till_done()
         assert f"{err.value}" == "TotalConnect usercode is invalid. Did not arm night"
         assert hass.states.get(ENTITY_ID).state == STATE_ALARM_DISARMED
         # should have started a re-auth flow

tests/components/vizio/test_media_player.py

@@ -456,6 +456,7 @@ async def test_options_update(
         options=new_options,
     )
     assert config_entry.options == updated_options
+    await hass.async_block_till_done()
     await _test_service(
         hass, MP_DOMAIN, "vol_up", SERVICE_VOLUME_UP, None, num=VOLUME_STEP
     )