Correct trace for choose and repeat script actions (#47973)

* Correct trace for choose and repeat script actions

* only choose-wrap the choices

* Update tests

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/helpers/script.py

@@ -640,6 +640,7 @@ class _ScriptRun:
         result = traced_test_conditions(self._hass, self._variables)
         return result
 
+    @trace_path("repeat")
     async def _async_repeat_step(self):
         """Repeat a sequence."""
         description = self._action.get(CONF_ALIAS, "sequence")
@@ -658,7 +659,7 @@ class _ScriptRun:
 
         async def async_run_sequence(iteration, extra_msg=""):
             self._log("Repeating %s: Iteration %i%s", description, iteration, extra_msg)
-            with trace_path(str(self._step)):
+            with trace_path("sequence"):
                 await self._async_run_script(script)
 
         if CONF_COUNT in repeat:
@@ -724,19 +725,21 @@ class _ScriptRun:
         # pylint: disable=protected-access
         choose_data = await self._script._async_get_choose_data(self._step)
 
-        for idx, (conditions, script) in enumerate(choose_data["choices"]):
+        with trace_path("choose"):
-            with trace_path(str(idx)):
+            for idx, (conditions, script) in enumerate(choose_data["choices"]):
-                try:
+                with trace_path(str(idx)):
-                    if self._test_conditions(conditions, "choose"):
+                    try:
-                        trace_set_result(choice=idx)
+                        if self._test_conditions(conditions, "choose"):
-                        await self._async_run_script(script)
+                            trace_set_result(choice=idx)
-                        return
+                            with trace_path("sequence"):
-                except exceptions.ConditionError as ex:
+                                await self._async_run_script(script)
-                    _LOGGER.warning("Error in 'choose' evaluation:\n%s", ex)
+                                return
+                    except exceptions.ConditionError as ex:
+                        _LOGGER.warning("Error in 'choose' evaluation:\n%s", ex)
 
         if choose_data["default"]:
             trace_set_result(choice="default")
-            with trace_path("default"):
+            with trace_path(["default", "sequence"]):
                 await self._async_run_script(choose_data["default"])
 
     async def _async_wait_for_trigger_step(self):

tests/helpers/test_script.py

@@ -1236,7 +1236,7 @@ async def test_repeat_count(hass, caplog, count):
     assert_action_trace(
         {
             "0": [{}],
-            "0/0/0": [{}] * min(count, script.ACTION_TRACE_NODE_MAX_LEN),
+            "0/sequence/0": [{}] * min(count, script.ACTION_TRACE_NODE_MAX_LEN),
         }
     )
 
@@ -1578,6 +1578,10 @@ async def test_choose(hass, caplog, var, result):
             },
         }
     )
+
+    # Prepare tracing
+    trace.trace_get()
+
     script_obj = script.Script(hass, sequence, "Test Name", "test_domain")
 
     await script_obj.async_run(MappingProxyType({"var": var}), Context())
@@ -1590,6 +1594,21 @@ async def test_choose(hass, caplog, var, result):
         expected_choice = "default"
     assert f"{alias}: {expected_choice}: Executing step {aliases[var]}" in caplog.text
 
+    expected_trace = {"0": [{}]}
+    if var >= 1:
+        expected_trace["0/choose/0"] = [{}]
+        expected_trace["0/choose/0/conditions/0"] = [{}]
+    if var >= 2:
+        expected_trace["0/choose/1"] = [{}]
+        expected_trace["0/choose/1/conditions/0"] = [{}]
+    if var == 1:
+        expected_trace["0/choose/0/sequence/0"] = [{}]
+    if var == 2:
+        expected_trace["0/choose/1/sequence/0"] = [{}]
+    if var == 3:
+        expected_trace["0/default/sequence/0"] = [{}]
+    assert_action_trace(expected_trace)
+
 
 @pytest.mark.parametrize(
     "action",