diff --git a/tests/Aspire.Hosting.LocalStack.Unit.Tests/Extensions/LocalStackProjectExtensionsTests.cs b/tests/Aspire.Hosting.LocalStack.Unit.Tests/Extensions/LocalStackProjectExtensionsTests.cs
index fbea36c..310220d 100644
--- a/tests/Aspire.Hosting.LocalStack.Unit.Tests/Extensions/LocalStackProjectExtensionsTests.cs
+++ b/tests/Aspire.Hosting.LocalStack.Unit.Tests/Extensions/LocalStackProjectExtensionsTests.cs
@@ -23,19 +23,21 @@ public async Task WithReference_Should_Add_LocalStack_Reference_To_Project()
     public async Task WithReference_Should_Return_Builder_When_LocalStack_Is_Null()
     {
         const string testProjectResourceName = "test-project";
+        IResourceBuilder<ProjectResource>? capturedResult = null;
+        IResourceBuilder<ProjectResource>? capturedBuilder = null;
 
         var (app, projectResource) = TestApplicationBuilder.CreateWithResource<ProjectResource>(testProjectResourceName, builder =>
         {
             var projectBuilder = builder.AddProject(testProjectResourceName, TestDataBuilders.GetTestProjectPath());
             var result = projectBuilder.WithReference(localStackBuilder: null);
 
-            // Should return the same builder
-            if (!ReferenceEquals(result, projectBuilder))
-            {
-                throw new InvalidOperationException("Builder should be the same reference");
-            }
+            capturedBuilder = projectBuilder;
+            capturedResult = result;
         });
 
+        // Should return the same builder
+        await Assert.That(capturedResult).IsSameReferenceAs(capturedBuilder);
+
         await Assert.That(app.HasResource<ILocalStackResource>("localstack")).IsFalse();
         await projectResource.ShouldNotHaveLocalStackEnabledAnnotation();
     }