From 14373c0358453a656d2e599c1a01828d6e1d2ef1 Mon Sep 17 00:00:00 2001
From: Mark Gould <mgould@spindlemedia.com>
Date: Thu, 30 Oct 2025 13:09:31 -0500
Subject: [PATCH] Fixes issue getting context when using InvokeHandler activity

---
 .../EndpointFixture.cs                                        | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/NServiceBus.Extensions.IntegrationTesting/EndpointFixture.cs b/src/NServiceBus.Extensions.IntegrationTesting/EndpointFixture.cs
index 7beef2c..94b26d9 100644
--- a/src/NServiceBus.Extensions.IntegrationTesting/EndpointFixture.cs
+++ b/src/NServiceBus.Extensions.IntegrationTesting/EndpointFixture.cs
@@ -100,8 +100,8 @@ private static async Task<ObservedMessageContexts> ExecuteAndWait<TMessageContex
 
                         break;
                     case NsbActivityNames.InvokeHandlerActivityName:
-                        var handlerContext = activitySource.Parent.GetTagItem("testing.invoke.handler.context") as IInvokeHandlerContext;
-                     
+                        var handlerContext = (activitySource.GetTagItem("testing.invoke.handler.context") ?? activitySource.Parent?.GetTagItem("testing.invoke.handler.context")) as IInvokeHandlerContext;
+
                         invokeHandlerContexts.Add(handlerContext);
 
                         if (handlerContext is TMessageContext ctx3 && predicate(ctx3))