Improve automation trace tests (#48542)

tests/components/trace/test_websocket_api.py

@@ -4,6 +4,7 @@ import pytest
 from homeassistant.bootstrap import async_setup_component
 from homeassistant.components.trace.const import STORED_TRACES
 from homeassistant.core import Context
+from homeassistant.helpers.typing import UNDEFINED
 
 from tests.common import assert_lists_same
 
@@ -26,10 +27,74 @@ def _find_traces(traces, trace_type, item_id):
     ]
 
 
+async def _setup_automation_or_script(hass, domain, configs, script_config=None):
+    """Set up automations or scripts from automation config."""
+    if domain == "script":
+        configs = {config["id"]: {"sequence": config["action"]} for config in configs}
+
+    if script_config:
+        if domain == "automation":
+            assert await async_setup_component(
+                hass, "script", {"script": script_config}
+            )
+        else:
+            configs = {**configs, **script_config}
+
+    assert await async_setup_component(hass, domain, {domain: configs})
+
+
+async def _run_automation_or_script(hass, domain, config, event, context=None):
+    if domain == "automation":
+        hass.bus.async_fire(event, context=context)
+    else:
+        await hass.services.async_call("script", config["id"], context=context)
+
+
+def _assert_raw_config(domain, config, trace):
+    if domain == "script":
+        config = {"sequence": config["action"]}
+    assert trace["config"] == config
+
+
+async def _assert_contexts(client, next_id, contexts):
+    await client.send_json({"id": next_id(), "type": "trace/contexts"})
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] == contexts
+
+
 @pytest.mark.parametrize(
-    "domain, prefix", [("automation", "action"), ("script", "sequence")]
+    "domain, prefix, extra_trace_keys, trigger, context_key, condition_results",
+    [
+        (
+            "automation",
+            "action",
+            [
+                {"trigger/0"},
+                {"trigger/0", "condition/0"},
+                {"trigger/1", "condition/0"},
+                {"trigger/0", "condition/0"},
+            ],
+            [
+                "event 'test_event'",
+                "event 'test_event2'",
+            ],
+            "parent_id",
+            [True],
+        ),
+        ("script", "sequence", [set(), set()], [UNDEFINED, UNDEFINED], "id", []),
+    ],
 )
-async def test_get_trace(hass, hass_ws_client, domain, prefix):
+async def test_get_trace(
+    hass,
+    hass_ws_client,
+    domain,
+    prefix,
+    extra_trace_keys,
+    trigger,
+    context_key,
+    condition_results,
+):
     """Test tracing an script or automation."""
     id = 1
 
@@ -55,9 +120,6 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
         },
         "action": {"event": "another_event"},
     }
-    if domain == "script":
-        sun_config = {"sequence": sun_config["action"]}
-        moon_config = {"sequence": moon_config["action"]}
 
     sun_action = {
         "limit": 10,
@@ -71,24 +133,14 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
     }
     moon_action = {"event": "another_event", "event_data": {}}
 
-    if domain == "automation":
-        assert await async_setup_component(
-            hass, domain, {domain: [sun_config, moon_config]}
-        )
-    else:
-        assert await async_setup_component(
-            hass, domain, {domain: {"sun": sun_config, "moon": moon_config}}
-        )
+    await _setup_automation_or_script(hass, domain, [sun_config, moon_config])
 
     client = await hass_ws_client()
     contexts = {}
 
     # Trigger "sun" automation / run "sun" script
     context = Context()
-    if domain == "automation":
-        hass.bus.async_fire("test_event", context=context)
-    else:
-        await hass.services.async_call("script", "sun", context=context)
+    await _run_automation_or_script(hass, domain, sun_config, "test_event", context)
     await hass.async_block_till_done()
 
     # List traces
@@ -110,25 +162,17 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
     response = await client.receive_json()
     assert response["success"]
     trace = response["result"]
-    if domain == "automation":
-        assert len(trace["trace"]) == 2
-        assert set(trace["trace"]) == {"trigger/0", f"{prefix}/0"}
-    else:
-        assert len(trace["trace"]) == 1
-        assert set(trace["trace"]) == {f"{prefix}/0"}
+    assert set(trace["trace"]) == {f"{prefix}/0"} | extra_trace_keys[0]
     assert len(trace["trace"][f"{prefix}/0"]) == 1
     assert trace["trace"][f"{prefix}/0"][0]["error"]
     assert trace["trace"][f"{prefix}/0"][0]["result"] == sun_action
-    assert trace["config"] == sun_config
+    _assert_raw_config(domain, sun_config, trace)
     assert trace["context"]
     assert trace["error"] == "Unable to find service test.automation"
     assert trace["state"] == "stopped"
     assert trace["item_id"] == "sun"
-    if domain == "automation":
-        assert trace["context"]["parent_id"] == context.id
-        assert trace["trigger"] == "event 'test_event'"
-    else:
-        assert trace["context"]["id"] == context.id
+    assert trace["context"][context_key] == context.id
+    assert trace.get("trigger", UNDEFINED) == trigger[0]
     contexts[trace["context"]["id"]] = {
         "run_id": trace["run_id"],
         "domain": domain,
@@ -136,10 +180,7 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
     }
 
     # Trigger "moon" automation, with passing condition / run "moon" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event2")
-    else:
-        await hass.services.async_call("script", "moon")
+    await _run_automation_or_script(hass, domain, moon_config, "test_event2", context)
     await hass.async_block_till_done()
 
     # List traces
@@ -161,37 +202,32 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
     response = await client.receive_json()
     assert response["success"]
     trace = response["result"]
-    if domain == "automation":
-        assert len(trace["trace"]) == 3
-        assert set(trace["trace"]) == {"trigger/0", "condition/0", f"{prefix}/0"}
-    else:
-        assert len(trace["trace"]) == 1
-        assert set(trace["trace"]) == {f"{prefix}/0"}
+    assert set(trace["trace"]) == {f"{prefix}/0"} | extra_trace_keys[1]
     assert len(trace["trace"][f"{prefix}/0"]) == 1
     assert "error" not in trace["trace"][f"{prefix}/0"][0]
     assert trace["trace"][f"{prefix}/0"][0]["result"] == moon_action
-    assert trace["config"] == moon_config
+    _assert_raw_config(domain, moon_config, trace)
     assert trace["context"]
     assert "error" not in trace
     assert trace["state"] == "stopped"
     assert trace["item_id"] == "moon"
 
-    if domain == "automation":
-        assert len(trace["trace"]["condition/0"]) == 1
-        assert trace["trace"]["condition/0"][0]["result"] == {"result": True}
-        assert trace["trigger"] == "event 'test_event2'"
+    assert trace.get("trigger", UNDEFINED) == trigger[1]
+
+    assert len(trace["trace"].get("condition/0", [])) == len(condition_results)
+    for idx, condition_result in enumerate(condition_results):
+        assert trace["trace"]["condition/0"][idx]["result"] == {
+            "result": condition_result
+        }
     contexts[trace["context"]["id"]] = {
         "run_id": trace["run_id"],
         "domain": domain,
         "item_id": trace["item_id"],
     }
 
-    if domain == "script":
+    if len(extra_trace_keys) <= 2:
         # Check contexts
-        await client.send_json({"id": next_id(), "type": "trace/contexts"})
-        response = await client.receive_json()
-        assert response["success"]
-        assert response["result"] == contexts
+        await _assert_contexts(client, next_id, contexts)
         return
 
     # Trigger "moon" automation with failing condition
@@ -217,7 +253,7 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
     response = await client.receive_json()
     assert response["success"]
     trace = response["result"]
-    assert set(trace["trace"]) == {"trigger/1", "condition/0"}
+    assert set(trace["trace"]) == extra_trace_keys[2]
     assert len(trace["trace"]["condition/0"]) == 1
     assert trace["trace"]["condition/0"][0]["result"] == {"result": False}
     assert trace["config"] == moon_config
@@ -255,8 +291,7 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
     response = await client.receive_json()
     assert response["success"]
     trace = response["result"]
-    assert len(trace["trace"]) == 3
-    assert set(trace["trace"]) == {"trigger/0", "condition/0", f"{prefix}/0"}
+    assert set(trace["trace"]) == {f"{prefix}/0"} | extra_trace_keys[3]
     assert len(trace["trace"][f"{prefix}/0"]) == 1
     assert "error" not in trace["trace"][f"{prefix}/0"][0]
     assert trace["trace"][f"{prefix}/0"][0]["result"] == moon_action
@@ -275,10 +310,7 @@ async def test_get_trace(hass, hass_ws_client, domain, prefix):
     }
 
     # Check contexts
-    await client.send_json({"id": next_id(), "type": "trace/contexts"})
-    response = await client.receive_json()
-    assert response["success"]
-    assert response["result"] == contexts
+    await _assert_contexts(client, next_id, contexts)
 
 
 @pytest.mark.parametrize("domain", ["automation", "script"])
@@ -320,18 +352,7 @@ async def test_trace_overflow(hass, hass_ws_client, domain):
         "trigger": {"platform": "event", "event_type": "test_event2"},
         "action": {"event": "another_event"},
     }
-    if domain == "script":
-        sun_config = {"sequence": sun_config["action"]}
-        moon_config = {"sequence": moon_config["action"]}
-
-    if domain == "automation":
-        assert await async_setup_component(
-            hass, domain, {domain: [sun_config, moon_config]}
-        )
-    else:
-        assert await async_setup_component(
-            hass, domain, {domain: {"sun": sun_config, "moon": moon_config}}
-        )
+    await _setup_automation_or_script(hass, domain, [sun_config, moon_config])
 
     client = await hass_ws_client()
 
@@ -341,12 +362,8 @@ async def test_trace_overflow(hass, hass_ws_client, domain):
     assert response["result"] == []
 
     # Trigger "sun" and "moon" automation / script once
-    if domain == "automation":
-        hass.bus.async_fire("test_event")
-        hass.bus.async_fire("test_event2")
-    else:
-        await hass.services.async_call("script", "sun")
-        await hass.services.async_call("script", "moon")
+    await _run_automation_or_script(hass, domain, sun_config, "test_event")
+    await _run_automation_or_script(hass, domain, moon_config, "test_event2")
     await hass.async_block_till_done()
 
     # List traces
@@ -359,10 +376,7 @@ async def test_trace_overflow(hass, hass_ws_client, domain):
 
     # Trigger "moon" enough times to overflow the max number of stored traces
     for _ in range(STORED_TRACES):
-        if domain == "automation":
-            hass.bus.async_fire("test_event2")
-        else:
-            await hass.services.async_call("script", "moon")
+        await _run_automation_or_script(hass, domain, moon_config, "test_event2")
         await hass.async_block_till_done()
 
     await client.send_json({"id": next_id(), "type": "trace/list", "domain": domain})
@@ -377,9 +391,28 @@ async def test_trace_overflow(hass, hass_ws_client, domain):
 
 
 @pytest.mark.parametrize(
-    "domain, prefix", [("automation", "action"), ("script", "sequence")]
+    "domain, prefix, trigger, last_step",
+    [
+        (
+            "automation",
+            "action",
+            [
+                "event 'test_event'",
+                "event 'test_event2'",
+                "event 'test_event3'",
+                "event 'test_event2'",
+            ],
+            ["{prefix}/0", "{prefix}/0", "condition/0", "{prefix}/0"],
+        ),
+        (
+            "script",
+            "sequence",
+            [UNDEFINED, UNDEFINED, UNDEFINED, UNDEFINED],
+            ["{prefix}/0", "{prefix}/0", "{prefix}/0", "{prefix}/0"],
+        ),
+    ],
 )
-async def test_list_traces(hass, hass_ws_client, domain, prefix):
+async def test_list_traces(hass, hass_ws_client, domain, prefix, trigger, last_step):
     """Test listing script and automation traces."""
     id = 1
 
@@ -405,18 +438,7 @@ async def test_list_traces(hass, hass_ws_client, domain, prefix):
         },
         "action": {"event": "another_event"},
     }
-    if domain == "script":
-        sun_config = {"sequence": sun_config["action"]}
-        moon_config = {"sequence": moon_config["action"]}
-
-    if domain == "automation":
-        assert await async_setup_component(
-            hass, domain, {domain: [sun_config, moon_config]}
-        )
-    else:
-        assert await async_setup_component(
-            hass, domain, {domain: {"sun": sun_config, "moon": moon_config}}
-        )
+    await _setup_automation_or_script(hass, domain, [sun_config, moon_config])
 
     client = await hass_ws_client()
 
@@ -433,10 +455,7 @@ async def test_list_traces(hass, hass_ws_client, domain, prefix):
     assert response["result"] == []
 
     # Trigger "sun" automation / run "sun" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event")
-    else:
-        await hass.services.async_call("script", "sun")
+    await _run_automation_or_script(hass, domain, sun_config, "test_event")
     await hass.async_block_till_done()
 
     # Get trace
@@ -462,24 +481,15 @@ async def test_list_traces(hass, hass_ws_client, domain, prefix):
     assert response["result"] == []
 
     # Trigger "moon" automation, with passing condition / run "moon" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event2")
-    else:
-        await hass.services.async_call("script", "moon")
+    await _run_automation_or_script(hass, domain, moon_config, "test_event2")
     await hass.async_block_till_done()
 
     # Trigger "moon" automation, with failing condition / run "moon" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event3")
-    else:
-        await hass.services.async_call("script", "moon")
+    await _run_automation_or_script(hass, domain, moon_config, "test_event3")
     await hass.async_block_till_done()
 
     # Trigger "moon" automation, with passing condition / run "moon" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event2")
-    else:
-        await hass.services.async_call("script", "moon")
+    await _run_automation_or_script(hass, domain, moon_config, "test_event2")
     await hass.async_block_till_done()
 
     # Get trace
@@ -489,48 +499,43 @@ async def test_list_traces(hass, hass_ws_client, domain, prefix):
     assert len(_find_traces(response["result"], domain, "moon")) == 3
     assert len(_find_traces(response["result"], domain, "sun")) == 1
     trace = _find_traces(response["result"], domain, "sun")[0]
-    assert trace["last_step"] == f"{prefix}/0"
+    assert trace["last_step"] == last_step[0].format(prefix=prefix)
     assert trace["error"] == "Unable to find service test.automation"
     assert trace["state"] == "stopped"
     assert trace["timestamp"]
     assert trace["item_id"] == "sun"
-    if domain == "automation":
-        assert trace["trigger"] == "event 'test_event'"
+    assert trace.get("trigger", UNDEFINED) == trigger[0]
 
     trace = _find_traces(response["result"], domain, "moon")[0]
-    assert trace["last_step"] == f"{prefix}/0"
+    assert trace["last_step"] == last_step[1].format(prefix=prefix)
     assert "error" not in trace
     assert trace["state"] == "stopped"
     assert trace["timestamp"]
     assert trace["item_id"] == "moon"
-    if domain == "automation":
-        assert trace["trigger"] == "event 'test_event2'"
+    assert trace.get("trigger", UNDEFINED) == trigger[1]
 
     trace = _find_traces(response["result"], domain, "moon")[1]
+    assert trace["last_step"] == last_step[2].format(prefix=prefix)
     assert "error" not in trace
     assert trace["state"] == "stopped"
     assert trace["timestamp"]
     assert trace["item_id"] == "moon"
-    if domain == "automation":
-        assert trace["last_step"] == "condition/0"
-        assert trace["trigger"] == "event 'test_event3'"
-    else:
-        assert trace["last_step"] == f"{prefix}/0"
+    assert trace.get("trigger", UNDEFINED) == trigger[2]
 
     trace = _find_traces(response["result"], domain, "moon")[2]
-    assert trace["last_step"] == f"{prefix}/0"
+    assert trace["last_step"] == last_step[3].format(prefix=prefix)
     assert "error" not in trace
     assert trace["state"] == "stopped"
     assert trace["timestamp"]
     assert trace["item_id"] == "moon"
-    if domain == "automation":
-        assert trace["trigger"] == "event 'test_event2'"
+    assert trace.get("trigger", UNDEFINED) == trigger[3]
 
 
 @pytest.mark.parametrize(
-    "domain, prefix", [("automation", "action"), ("script", "sequence")]
+    "domain, prefix, extra_trace_keys",
+    [("automation", "action", {"trigger/0"}), ("script", "sequence", set())],
 )
-async def test_nested_traces(hass, hass_ws_client, domain, prefix):
+async def test_nested_traces(hass, hass_ws_client, domain, prefix, extra_trace_keys):
     """Test nested automation and script traces."""
     id = 1
 
@@ -544,48 +549,24 @@ async def test_nested_traces(hass, hass_ws_client, domain, prefix):
         "trigger": {"platform": "event", "event_type": "test_event"},
         "action": {"service": "script.moon"},
     }
-    moon_config = {
-        "sequence": {"event": "another_event"},
-    }
-    if domain == "script":
-        sun_config = {"sequence": sun_config["action"]}
-
-    if domain == "automation":
-        assert await async_setup_component(hass, domain, {domain: [sun_config]})
-        assert await async_setup_component(
-            hass, "script", {"script": {"moon": moon_config}}
-        )
-    else:
-        assert await async_setup_component(
-            hass, domain, {domain: {"sun": sun_config, "moon": moon_config}}
-        )
+    moon_config = {"moon": {"sequence": {"event": "another_event"}}}
+    await _setup_automation_or_script(hass, domain, [sun_config], moon_config)
 
     client = await hass_ws_client()
 
     # Trigger "sun" automation / run "sun" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event")
-    else:
-        await hass.services.async_call("script", "sun")
+    await _run_automation_or_script(hass, domain, sun_config, "test_event")
     await hass.async_block_till_done()
 
     # List traces
     await client.send_json({"id": next_id(), "type": "trace/list", "domain": "script"})
     response = await client.receive_json()
     assert response["success"]
-    if domain == "automation":
-        assert len(response["result"]) == 1
-    else:
-        assert len(response["result"]) == 2
     assert len(_find_traces(response["result"], "script", "moon")) == 1
     moon_run_id = _find_run_id(response["result"], "script", "moon")
-    if domain == "automation":
-        await client.send_json(
-            {"id": next_id(), "type": "trace/list", "domain": "automation"}
-        )
-        response = await client.receive_json()
-        assert response["success"]
-        assert len(response["result"]) == 1
+    await client.send_json({"id": next_id(), "type": "trace/list", "domain": domain})
+    response = await client.receive_json()
+    assert response["success"]
     assert len(_find_traces(response["result"], domain, "sun")) == 1
     sun_run_id = _find_run_id(response["result"], domain, "sun")
     assert sun_run_id != moon_run_id
@@ -603,12 +584,7 @@ async def test_nested_traces(hass, hass_ws_client, domain, prefix):
     response = await client.receive_json()
     assert response["success"]
     trace = response["result"]
-    if domain == "automation":
-        assert len(trace["trace"]) == 2
-        assert set(trace["trace"]) == {"trigger/0", f"{prefix}/0"}
-    else:
-        assert len(trace["trace"]) == 1
-        assert set(trace["trace"]) == {f"{prefix}/0"}
+    assert set(trace["trace"]) == {f"{prefix}/0"} | extra_trace_keys
     assert len(trace["trace"][f"{prefix}/0"]) == 1
     child_id = trace["trace"][f"{prefix}/0"][0]["child_id"]
     assert child_id == {"domain": "script", "item_id": "moon", "run_id": moon_run_id}
@@ -652,13 +628,7 @@ async def test_breakpoints(hass, hass_ws_client, domain, prefix):
             {"event": "event8"},
         ],
     }
-    if domain == "script":
-        sun_config = {"sequence": sun_config["action"]}
-
-    if domain == "automation":
-        assert await async_setup_component(hass, domain, {domain: [sun_config]})
-    else:
-        assert await async_setup_component(hass, domain, {domain: {"sun": sun_config}})
+    await _setup_automation_or_script(hass, domain, [sun_config])
 
     client = await hass_ws_client()
 
@@ -721,10 +691,7 @@ async def test_breakpoints(hass, hass_ws_client, domain, prefix):
     )
 
     # Trigger "sun" automation / run "sun" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event")
-    else:
-        await hass.services.async_call("script", "sun")
+    await _run_automation_or_script(hass, domain, sun_config, "test_event")
 
     response = await client.receive_json()
     run_id = await assert_last_step("sun", f"{prefix}/1", "running")
@@ -830,13 +797,7 @@ async def test_breakpoints_2(hass, hass_ws_client, domain, prefix):
             {"event": "event8"},
         ],
     }
-    if domain == "script":
-        sun_config = {"sequence": sun_config["action"]}
-
-    if domain == "automation":
-        assert await async_setup_component(hass, domain, {domain: [sun_config]})
-    else:
-        assert await async_setup_component(hass, domain, {domain: {"sun": sun_config}})
+    await _setup_automation_or_script(hass, domain, [sun_config])
 
     client = await hass_ws_client()
 
@@ -860,10 +821,7 @@ async def test_breakpoints_2(hass, hass_ws_client, domain, prefix):
     assert response["success"]
 
     # Trigger "sun" automation / run "sun" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event")
-    else:
-        await hass.services.async_call("script", "sun")
+    await _run_automation_or_script(hass, domain, sun_config, "test_event")
 
     response = await client.receive_json()
     run_id = await assert_last_step("sun", f"{prefix}/1", "running")
@@ -897,10 +855,7 @@ async def test_breakpoints_2(hass, hass_ws_client, domain, prefix):
     assert not response["success"]
 
     # Trigger "sun" automation / script, should finish without stopping on breakpoints
-    if domain == "automation":
-        hass.bus.async_fire("test_event")
-    else:
-        await hass.services.async_call("script", "sun")
+    await _run_automation_or_script(hass, domain, sun_config, "test_event")
     await hass.async_block_till_done()
 
     new_run_id = await assert_last_step("sun", f"{prefix}/8", "stopped")
@@ -945,13 +900,7 @@ async def test_breakpoints_3(hass, hass_ws_client, domain, prefix):
             {"event": "event8"},
         ],
     }
-    if domain == "script":
-        sun_config = {"sequence": sun_config["action"]}
-
-    if domain == "automation":
-        assert await async_setup_component(hass, domain, {domain: [sun_config]})
-    else:
-        assert await async_setup_component(hass, domain, {domain: {"sun": sun_config}})
+    await _setup_automation_or_script(hass, domain, [sun_config])
 
     client = await hass_ws_client()
 
@@ -987,10 +936,7 @@ async def test_breakpoints_3(hass, hass_ws_client, domain, prefix):
     assert response["success"]
 
     # Trigger "sun" automation / run "sun" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event")
-    else:
-        await hass.services.async_call("script", "sun")
+    await _run_automation_or_script(hass, domain, sun_config, "test_event")
 
     response = await client.receive_json()
     run_id = await assert_last_step("sun", f"{prefix}/1", "running")
@@ -1050,10 +996,7 @@ async def test_breakpoints_3(hass, hass_ws_client, domain, prefix):
     assert response["success"]
 
     # Trigger "sun" automation / run "sun" script
-    if domain == "automation":
-        hass.bus.async_fire("test_event")
-    else:
-        await hass.services.async_call("script", "sun")
+    await _run_automation_or_script(hass, domain, sun_config, "test_event")
 
     response = await client.receive_json()
     run_id = await assert_last_step("sun", f"{prefix}/5", "running")