test: strengthen assertions and remove timing assumptions

tests/PostHog.AI.Tests/PostHogAIExtensionsTests.cs

@@ -38,5 +38,9 @@ public void AddPostHogOpenAIClientSucceedsWhenPostHogIsRegistered()
         // Should not throw
         var builder = services.AddPostHogOpenAIClient("sk-test-key");
         Assert.NotNull(builder);
+        using var provider = services.BuildServiceProvider();
+        Assert.NotNull(provider.GetRequiredService<OpenAI.OpenAIClient>());
+        using var handler = provider.GetRequiredService<PostHogOpenAIHandler>();
+        Assert.NotNull(handler);
     }
 }

tests/PostHog.AI.Tests/PostHogOpenAIHandlerTests.cs

@@ -466,8 +466,7 @@ public async Task SendAsyncIncludesInputAndOutputChoicesWhenPrivacyModeIsFalseSi
                 Arg.Any<string>(),
                 PostHogAIFieldNames.Generation,
                 Arg.Is<Dictionary<string, object>>(props =>
-                    props.ContainsKey(PostHogAIFieldNames.Input)
-                    && props.ContainsKey(PostHogAIFieldNames.OutputChoices)
+                    HasExpectedMessagePayloads(props)
                 ),
                 null,
                 false,
@@ -690,21 +689,23 @@ await writer.WriteAsync(
         HttpResponseMessage? response = null;
         try
         {
-            using (PostHogAIContext.BeginScope(privacyMode: true))
+            using (PostHogAIContext.BeginScope(distinctId: "request-user", traceId: "request-trace", privacyMode: true))
             {
-                response = await _client.PostAsync(
-                    new Uri("/v1/chat/completions", UriKind.Relative),
-                    requestContent
-                );
-
+                using var request = new HttpRequestMessage(HttpMethod.Post, new Uri("/v1/chat/completions", UriKind.Relative))
+                {
+                    Content = requestContent
+                };
+                response = await _client.SendAsync(request, HttpCompletionOption.ResponseHeadersRead);
                 Assert.True(response.IsSuccessStatusCode);
+                Assert.Empty(_postHogClient.ReceivedCalls());
             }
 
-            // Stream consumed outside the scope — privacy mode must still apply
-            var resultStream = await response.Content.ReadAsStreamAsync();
-            using (var reader = new StreamReader(resultStream))
+            using (PostHogAIContext.BeginScope(distinctId: "other-user", traceId: "other-trace", privacyMode: false))
             {
-                await reader.ReadToEndAsync();
+                var resultStream = await response.Content.ReadAsStreamAsync();
+                using var reader = new StreamReader(resultStream);
+                Assert.Equal(Encoding.UTF8.GetString(sseStream.ToArray()), await reader.ReadToEndAsync());
+                Assert.Empty(_postHogClient.ReceivedCalls());
             }
         }
         finally
@@ -723,11 +724,13 @@ await writer.WriteAsync(
                     && !props.ContainsKey(PostHogAIFieldNames.OutputChoices)
                     && (int)props[PostHogAIFieldNames.InputTokens] == 10
                     && (int)props[PostHogAIFieldNames.OutputTokens] == 5
+                    && (string)props[PostHogAIFieldNames.TraceId] == "request-trace"
                 ),
                 null,
                 false,
                 Arg.Any<DateTimeOffset?>()
             );
+        Assert.Equal("request-user", Assert.Single(_postHogClient.ReceivedCalls()).GetArguments()[0]);
     }
 
     [Fact]
@@ -794,15 +797,22 @@ public async Task SendAsyncIncludesInputAndOutputChoicesWhenPrivacyModeIsNullSim
                 Arg.Any<string>(),
                 PostHogAIFieldNames.Generation,
                 Arg.Is<Dictionary<string, object>>(props =>
-                    props.ContainsKey(PostHogAIFieldNames.Input)
-                    && props.ContainsKey(PostHogAIFieldNames.OutputChoices)
+                    HasExpectedMessagePayloads(props)
                 ),
                 null,
                 false,
                 Arg.Any<DateTimeOffset?>()
             );
     }
 
+    static bool HasExpectedMessagePayloads(Dictionary<string, object> properties)
+        => JsonNode.DeepEquals(
+               JsonNode.Parse("""[{"role":"user","content":"Hello"}]"""),
+               properties[PostHogAIFieldNames.Input] as JsonNode)
+           && JsonNode.DeepEquals(
+               JsonNode.Parse("""[{"index":0,"message":{"role":"assistant","content":"Hi there!"},"finish_reason":"stop"}]"""),
+               properties[PostHogAIFieldNames.OutputChoices] as JsonNode);
+
     [Fact]
     public async Task SendAsyncCapturesErrorEventOnNetworkException()
     {
@@ -933,13 +943,18 @@ public async Task SendAsyncHandlesGracefullyOnMalformedJsonResponse()
 
         Assert.True(response.IsSuccessStatusCode);
 
+        Assert.Equal("this is not json {{{{", await response.Content.ReadAsStringAsync());
+
         // Event should still be captured (with whatever properties could be extracted)
         _postHogClient
             .Received(1)
             .Capture(
                 Arg.Any<string>(),
-                Arg.Any<string>(),
-                Arg.Any<Dictionary<string, object>>(),
+                PostHogAIFieldNames.Generation,
+                Arg.Is<Dictionary<string, object>>(props =>
+                    (string)props[PostHogAIFieldNames.Model] == "gpt-4"
+                    && (string)props[PostHogAIFieldNames.Provider] == "openai"
+                    && (int)props[PostHogAIFieldNames.HttpStatus] == 200),
                 null,
                 false,
                 Arg.Any<DateTimeOffset?>()
@@ -1017,7 +1032,8 @@ public async Task SendAsyncPassesGroupsToCaptureWhenContextHasGroups()
                 Arg.Any<string>(),
                 PostHogAIFieldNames.Generation,
                 Arg.Any<Dictionary<string, object>>(),
-                Arg.Is<GroupCollection>(g => g != null && g.Count == 1),
+                Arg.Is<GroupCollection>(g => g != null && g.Count == 1
+                    && g.Single().GroupType == "company" && g.Single().GroupKey == "acme-corp"),
                 false,
                 Arg.Any<DateTimeOffset?>()
             );
@@ -1078,6 +1094,7 @@ public async Task SendAsyncContextPropertiesOverrideEventProperties()
                 properties: new Dictionary<string, object>
                 {
                     { "custom_prop", "custom_value" },
+                    { PostHogAIFieldNames.Model, "context-model" },
                 }
             )
         )
@@ -1098,6 +1115,8 @@ public async Task SendAsyncContextPropertiesOverrideEventProperties()
                 PostHogAIFieldNames.Generation,
                 Arg.Is<Dictionary<string, object>>(props =>
                     (string)props["custom_prop"] == "custom_value"
+                    && (string)props[PostHogAIFieldNames.Model] == "context-model"
+                    && (int)props[PostHogAIFieldNames.InputTokens] == 9
                 ),
                 null,
                 false,

tests/UnitTests.AspNetCore/HttpContextFeatureFlagCacheTests.cs

@@ -86,12 +86,17 @@ public async Task DoesNotCacheIfHttpContextIsNull()
             }
         };
 
-        var result = await cache.GetAndCacheFlagsAsync(
-            distinctId,
-            (_, _) => Task.FromResult(flagsResult),
-            CancellationToken.None);
+        var secondFlagsResult = new FlagsResult { RequestId = "second-fetch" };
+        var fetchCount = 0;
+        Task<FlagsResult> Fetch(string _, CancellationToken token)
+            => Task.FromResult(++fetchCount == 1 ? flagsResult : secondFlagsResult);
 
-        Assert.Equal(flagsResult, result);
+        var result = await cache.GetAndCacheFlagsAsync(distinctId, Fetch, CancellationToken.None);
+        var secondResult = await cache.GetAndCacheFlagsAsync(distinctId, Fetch, CancellationToken.None);
+
+        Assert.Equal(2, fetchCount);
+        Assert.Same(flagsResult, result);
+        Assert.Same(secondFlagsResult, secondResult);
     }
 
 
@@ -231,7 +236,8 @@ public async Task ReturnsCachedFlagsWhenPersonPropertiesMatch()
 
         var cache = new HttpContextFeatureFlagCache(httpContextAccessor);
         var distinctId = "user123";
-        var personProperties = new Dictionary<string, object?> { ["email"] = "test@example.com" };
+        var personProperties = new Dictionary<string, object?> { ["email"] = "test@example.com", ["plan"] = "pro" };
+        var equalProperties = new Dictionary<string, object?> { ["plan"] = "pro", ["email"] = "test@example.com" };
 
         var flags = new FlagsResult
         {
@@ -249,7 +255,7 @@ public async Task ReturnsCachedFlagsWhenPersonPropertiesMatch()
         };
 
         var result1 = await cache.GetAndCacheFlagsAsync(distinctId, personProperties, null, fetcher, CancellationToken.None);
-        var result2 = await cache.GetAndCacheFlagsAsync(distinctId, personProperties, null, fetcher, CancellationToken.None);
+        var result2 = await cache.GetAndCacheFlagsAsync(distinctId, equalProperties, null, fetcher, CancellationToken.None);
 
         Assert.Equal(1, fetchCount); // Should only fetch once
         Assert.Same(result1, result2);
@@ -457,11 +463,19 @@ public async Task DoesNotCacheIfHttpContextIsNull()
             { "feature1", new FeatureFlag { Key = "feature1", IsEnabled = true } }
         };
 
-        var result = await cache.GetAndCacheFeatureFlagsAsync(
-            distinctId,
-            _ => Task.FromResult((IReadOnlyDictionary<string, FeatureFlag>)featureFlags),
-            CancellationToken.None);
+        var secondFlags = new Dictionary<string, FeatureFlag>
+        {
+            ["feature1"] = new FeatureFlag { Key = "feature1", IsEnabled = false }
+        };
+        var fetchCount = 0;
+        Task<IReadOnlyDictionary<string, FeatureFlag>> Fetch(CancellationToken _)
+            => Task.FromResult<IReadOnlyDictionary<string, FeatureFlag>>(++fetchCount == 1 ? featureFlags : secondFlags);
+
+        var result = await cache.GetAndCacheFeatureFlagsAsync(distinctId, Fetch, CancellationToken.None);
+        var secondResult = await cache.GetAndCacheFeatureFlagsAsync(distinctId, Fetch, CancellationToken.None);
 
+        Assert.Equal(2, fetchCount);
         Assert.Equal(featureFlags, result);
+        Assert.Equal(secondFlags, secondResult);
     }
 }

tests/UnitTests.AspNetCore/PostHogRequestContextMiddlewareTests.cs

@@ -260,11 +260,17 @@ public async Task ConcurrentRequestsDoNotLeakRequestContext()
         var client = container.Activate<PostHogClient>();
         var results = new Dictionary<string, (string? DistinctId, string? SessionId)>();
         var gate = new object();
+        var bothEntered = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var enteredCount = 0;
 
         var middleware = CreateMiddleware(
             async context =>
             {
-                await Task.Delay(25);
+                if (Interlocked.Increment(ref enteredCount) == 2)
+                {
+                    bothEntered.SetResult(true);
+                }
+                await bothEntered.Task.WaitAsync(TimeSpan.FromSeconds(5));
                 lock (gate)
                 {
                     results[context.Request.Path.Value ?? string.Empty] = (
@@ -388,13 +394,22 @@ public async Task PreservesOriginalExceptionWhenExceptionCaptureThrows()
                 Arg.Any<DateTimeOffset?>())
             .Returns(_ => throw new InvalidOperationException("capture failed"));
 
+        var original = new NotSupportedException("original");
         var middleware = CreateMiddleware(
-            _ => throw new NotSupportedException("original"),
+            _ => throw original,
             postHog,
             options => options.CaptureExceptions = true);
 
         var exception = await Assert.ThrowsAsync<NotSupportedException>(() => middleware.InvokeAsync(CreateHttpContext()));
-        Assert.Equal("original", exception.Message);
+        Assert.Same(original, exception);
+        postHog.Received(1).CaptureException(
+            original,
+            Arg.Any<string>(),
+            Arg.Is<Dictionary<string, object>?>(properties => properties != null
+                && (int)properties["$response_status_code"] == 500),
+            Arg.Any<GroupCollection?>(),
+            Arg.Any<FeatureFlagEvaluations?>(),
+            Arg.Any<DateTimeOffset?>());
     }
 
     static PostHogRequestContextMiddleware CreateMiddleware(

tests/UnitTests/CaptureExtensionsTests.cs

@@ -74,6 +74,7 @@ public void CaptureSurveyResponsesDoesNotMutateProvidedProperties()
             "survey sent",
             Arg.Is<Dictionary<string, object>>(captured =>
                 !ReferenceEquals(captured, properties)
+                && (string)captured["source"] == "test"
                 && (string)captured["$survey_id"] == "survey-id"
                 && (string)captured["$survey_response"] == "first"
                 && (string)captured["survey_response_1"] == "second"),
@@ -88,6 +89,7 @@ static bool HasCopiedPersonProperties(
         Dictionary<string, object> personPropertiesToSet,
         Dictionary<string, object> personPropertiesToSetOnce)
         => !ReferenceEquals(captured, properties)
+           && (string)captured["source"] == "test"
            && captured["$set"] is Dictionary<string, object> set
            && !ReferenceEquals(set, personPropertiesToSet)
            && (string)set["name"] == "Max"

tests/UnitTests/Features/ETagSupportTests.cs

@@ -86,7 +86,7 @@ public async Task NotModifiedResponseReturnsCachedFlags()
             "\"etag-123\"");
 
         // Second request returns 304 Not Modified
-        container.FakeHttpMessageHandler.AddLocalEvaluationNotModifiedResponse("\"etag-123\"");
+        var notModifiedHandler = container.FakeHttpMessageHandler.AddLocalEvaluationNotModifiedResponse("\"etag-123\"");
 
         // Need batch response for $feature_flag_called event
         container.FakeHttpMessageHandler.AddBatchResponse();
@@ -99,6 +99,9 @@ public async Task NotModifiedResponseReturnsCachedFlags()
         // Second load gets 304 - should still have flags from cache
         await client.LoadFeatureFlagsAsync(CancellationToken.None);
 
+        var request = Assert.Single(notModifiedHandler.ReceivedRequests);
+        Assert.Contains(request.Headers.IfNoneMatch, etag => etag.Tag == "\"etag-123\"");
+
         // Verify flags still work from cached evaluator
         var result = await client.IsFeatureEnabledAsync("test-flag", "user-123");
         Assert.True(result);

tests/UnitTests/Features/FallbackFeatureFlagCacheTests.cs

@@ -1,3 +1,4 @@
+using Microsoft.Extensions.Time.Testing;
 using PostHog;
 using PostHog.Features;
 
@@ -8,9 +9,9 @@ public class TheGetAndCacheFeatureFlagsAsyncMethod
     [Fact]
     public async Task ReturnsItemInPrimaryCache()
     {
-        var timeProvider = TimeProvider.System;
-        var primaryCache = new MemoryFeatureFlagCache(timeProvider, 10, 0.2);
-        var secondaryCache = new MemoryFeatureFlagCache(timeProvider, 10, 0.2);
+        var timeProvider = new FakeTimeProvider();
+        using var primaryCache = new MemoryFeatureFlagCache(timeProvider, 10, 0.2);
+        using var secondaryCache = new MemoryFeatureFlagCache(timeProvider, 10, 0.2);
 
         var cache = new FallbackFeatureFlagCache(primaryCache, secondaryCache);
         var distinctId = "test-distinct-id";
@@ -30,9 +31,9 @@ public async Task ReturnsItemInPrimaryCache()
     [Fact]
     public async Task ReturnsItemInSecondaryCache()
     {
-        var timeProvider = TimeProvider.System;
-        var primaryCache = new MemoryFeatureFlagCache(timeProvider, 10, 0.2);
-        var secondaryCache = new MemoryFeatureFlagCache(timeProvider, 10, 0.2);
+        var timeProvider = new FakeTimeProvider();
+        using var primaryCache = new MemoryFeatureFlagCache(timeProvider, 10, 0.2);
+        using var secondaryCache = new MemoryFeatureFlagCache(timeProvider, 10, 0.2);
 
         var cache = new FallbackFeatureFlagCache(primaryCache, secondaryCache);
         var distinctId = "test-distinct-id";

tests/UnitTests/Features/FeatureFlagCacheKeyTests.cs

@@ -287,6 +287,8 @@ public void HandlesNullValuesInProperties()
         var key2 = FeatureFlagCacheKey.Generate("user123", properties2, null);
 
         Assert.Equal(key1, key2);
+        properties2.Remove("middle_name");
+        Assert.NotEqual(key1, FeatureFlagCacheKey.Generate("user123", properties2, null));
     }
 
     [Fact]
@@ -305,6 +307,8 @@ public void HandlesComplexNestedProperties()
         var key2 = FeatureFlagCacheKey.Generate("user123", properties, null);
 
         Assert.Equal(key1, key2);
+        ((Dictionary<string, object>)properties["metadata"]!)["version"] = 3;
+        Assert.NotEqual(key1, FeatureFlagCacheKey.Generate("user123", properties, null));
     }
 
     [Fact]
@@ -348,6 +352,8 @@ public void GeneratesStableKeyWithMultipleGroupsEachWithProperties()
         var key2 = FeatureFlagCacheKey.Generate("user123", null, groups);
 
         Assert.Equal(key1, key2);
+        groups.Single(group => group.GroupType == "team")["size"] = 51;
+        Assert.NotEqual(key1, FeatureFlagCacheKey.Generate("user123", null, groups));
     }
 
     [Fact]
@@ -367,6 +373,11 @@ public void HandlesDiversePropertyTypes()
         var key2 = FeatureFlagCacheKey.Generate("user123", properties, null);
 
         Assert.Equal(key1, key2);
+        properties["int"] = "42";
+        Assert.NotEqual(key1, FeatureFlagCacheKey.Generate("user123", properties, null));
+        properties["int"] = 42;
+        properties["array"] = new[] { "vip", "stable" };
+        Assert.NotEqual(key1, FeatureFlagCacheKey.Generate("user123", properties, null));
     }
 
     [Fact]
@@ -396,6 +407,7 @@ public void HandlesWhitespaceDistinctId()
         var key2 = FeatureFlagCacheKey.Generate("   ", null, null);
 
         Assert.Equal(key1, key2);
+        Assert.Equal("   ", key1);
     }
 
     [Fact]
@@ -436,6 +448,8 @@ public void HandlesLargeNumberOfProperties()
 
         Assert.Equal(key1, key2);
         Assert.NotEmpty(key1);
+        properties["prop99"] = "changed";
+        Assert.NotEqual(key1, FeatureFlagCacheKey.Generate("user123", properties, null));
     }
 
     [Fact]
@@ -452,6 +466,8 @@ public void HandlesLargeNumberOfGroups()
 
         Assert.Equal(key1, key2);
         Assert.NotEmpty(key1);
+        groups["group19"] = new Group("group19", "changed");
+        Assert.NotEqual(key1, FeatureFlagCacheKey.Generate("user123", null, groups));
     }
 
     [Fact]