Fix test_flow_run_instrumentation_captures_labels

src/prefect/types/__init__.py

@@ -1,6 +1,6 @@
 from functools import partial
-from typing import Annotated, Any, Dict, List, Optional, Set, TypeAlias, TypeVar, Union
-from typing_extensions import Literal
+from typing import Annotated, Any, Dict, List, Optional, Set, TypeVar, Union
+from typing_extensions import Literal, TypeAlias
 import orjson
 import pydantic
 

tests/telemetry/instrumentation_tester.py

@@ -95,8 +95,8 @@ def get_finished_spans(self):
     def assert_has_attributes(obj: HasAttributes, attributes: Dict[str, Any]):
         assert obj.attributes is not None
         for key, val in attributes.items():
-            assert key in obj.attributes
-            assert obj.attributes[key] == val
+            assert key in obj.attributes, f"Key {key!r} not found in attributes"
+            assert obj.attributes[key] == val, f"Value for key {key!r} does not match"
 
     @staticmethod
     def assert_span_instrumented_for(span: Union[Span, ReadableSpan], module):

tests/test_flow_engine.py

@@ -16,7 +16,7 @@
 from prefect import Flow, __development_base_path__, flow, task
 from prefect.client.orchestration import PrefectClient, SyncPrefectClient
 from prefect.client.schemas.filters import FlowFilter, FlowRunFilter
-from prefect.client.schemas.objects import FlowRun, StateType
+from prefect.client.schemas.objects import StateType
 from prefect.client.schemas.sorting import FlowRunSort
 from prefect.concurrency.asyncio import concurrency as aconcurrency
 from prefect.concurrency.sync import concurrency
@@ -1891,27 +1891,18 @@ def instrumented_flow():
         assert span.status.status_code == trace.StatusCode.OK
 
     def test_flow_run_instrumentation_captures_labels(
-        self, instrumentation: InstrumentationTester, monkeypatch
+        self,
+        instrumentation: InstrumentationTester,
+        sync_prefect_client: SyncPrefectClient,
     ):
-        # simulate server responding with labels on flow run
-        class FlowRunWithLabels(FlowRun):
-            labels: KeyValueLabels = pydantic.Field(
-                default_factory=lambda: {
-                    "prefect.deployment.id": "some-id",
-                    "my-label": "my-value",
-                }
-            )
-
-        monkeypatch.setattr(
-            "prefect.client.orchestration.FlowRun",
-            FlowRunWithLabels,
-        )
-
         @flow
         def instrumented_flow():
             pass
 
-        instrumented_flow()
+        state = instrumented_flow(return_state=True)
+
+        assert state.state_details.flow_run_id is not None
+        flow_run = sync_prefect_client.read_flow_run(state.state_details.flow_run_id)
 
         spans = instrumentation.get_finished_spans()
         assert len(spans) == 1
@@ -1921,11 +1912,10 @@ def instrumented_flow():
         instrumentation.assert_has_attributes(
             span,
             {
+                **flow_run.labels,
                 "prefect.run.type": "flow",
                 "prefect.flow.name": "instrumented-flow",
                 "prefect.run.id": mock.ANY,
-                "prefect.deployment.id": "some-id",
-                "my-label": "my-value",
             },
         )