fix: preserve StateBag when creating per-test TestBuilderContext

[thomhurst]

Summary

Fixes #5267

• Data added to TestBuilderContext.StateBag during data source generation (e.g., in MethodDataSource methods) was lost because the per-test TestBuilderContext created at each data combination in TestBuilder.cs:282 did not copy the StateBag from the discovery context
• The fix copies the discovery-time StateBag into each per-test context so data flows through to TestContext.StateBag during execution

Test plan

• Added regression test TUnit.TestProject/Bugs/Issue5267/Tests.cs that reproduces the exact scenario from the issue
• Verified existing StateBagPropagationTests (bug 3723) still pass
• Verified UniqueBuilderContextsOnEnumerableDataGeneratorTests still pass

[claude]

Review: fix: preserve StateBag when creating per-test TestBuilderContext

The fix is correct and well-targeted. The root cause is accurately diagnosed: data written to TestBuilderContext.Current.StateBag inside an IEnumerable-based data source method runs during the await foreach over GetInitializedDataRowsAsync, which populates the outer testBuilderContext.StateBag. The per-test contexts created at line 282 were new, empty ConcurrentDictionary instances, so all those entries were silently discarded.

What's Good

• The fix is minimal and surgical — a single line at the one point in BuildTestsFromMetadataAsync where per-iteration contexts are created.
• The regression test (TUnit.TestProject/Bugs/Issue5267/Tests.cs) reproduces the exact scenario from the issue report, covers two keys, and uses two data items to confirm both iterations get the expected values.
• The comment // Preserve StateBag from discovery context is clear about intent.

Observations & Potential Issues

1. Copy-on-create semantics vs. reference sharing (line 286 vs. line 419 and 1856)

At line 282–289, the new per-iteration TestBuilderContext is given a copy of the discovery-time StateBag:

StateBag = new ConcurrentDictionary<string, object?>(testBuilderContext.StateBag)

But at line 415–423 (testSpecificContext in BuildTestsFromMetadataAsync) and at line 1852–1860 (testSpecificContext in BuildSingleTestAsync), the StateBag is directly referenced (not copied):

StateBag = contextAccessor.Current.StateBag,  // same ConcurrentDictionary instance

This means writes to StateBag after the per-iteration context is created (e.g., from post-data-factory steps) will mutate the shared dict and be visible across tests. The copy at line 286 makes the per-test context's StateBag independent, but the testSpecificContext at line 419 then reuses that same object. This is consistent internally — within one test's build path the copy is made once at line 286 and then forwarded by reference — but it is worth being explicit about the intended contract: "each test starts with a snapshot of the discovery-time bag, and any writes after that point accumulate only on that test's copy."

2. The SkipIfEmpty testSpecificContext paths (lines 481 and 540) do not copy StateBag

When a data source yields no items and SkipIfEmpty is set, a synthetic skipped test is created:

// Line 481 (method data, SkipIfEmpty)
var testSpecificContext = new TestBuilderContext
{
    TestMetadata = metadata.MethodMetadata,
    ClassConstructor = testBuilderContext.ClassConstructor,
    DataSourceAttribute = methodDataSource,
    InitializedAttributes = attributes
    // StateBag is NOT copied here
};

// Line 540 (class data, SkipIfEmpty)  — same pattern

These don't copy the StateBag from testBuilderContext. For a completely empty data source this is arguably moot (the StateBag couldn't have been populated if no items were yielded), but if a data source writes to StateBag in its generator body before yielding anything and then yields nothing, those entries would be silently lost in the skipped-test path. Whether this edge case matters in practice depends on usage patterns, but it's worth noting for consistency.

3. Streaming path (BuildSingleTestAsync / BuildTestsFromStreamAsync) is not affected

In the streaming path (BuildSingleTestAsync, line 1853) the contextAccessor.Current is the baseContext object itself — the context is never replaced with a per-iteration copy. So testSpecificContext.StateBag directly aliases baseContext.StateBag, and all tests in a stream share one dict. This is a pre-existing design choice (not introduced by this PR), but it is a semantic difference compared to the BuildTestsFromMetadataAsync path where each iteration now gets an independent copy.

4. Thread-safety of the copy

new ConcurrentDictionary<string, object?>(testBuilderContext.StateBag) takes a snapshot of the current entries. If data sources run concurrently and write to the dict simultaneously, the snapshot is not guaranteed to be perfectly consistent, but data source enumeration appears to be serial (await foreach), so in practice this is safe.

Minor Nit

The inline comment // Preserve StateBag from discovery context could mention that this is a shallow copy so readers know the values themselves are not deep-cloned. Object values that are mutable reference types will still be shared across per-test copies:

// Shallow-copy: value objects are shared, but key additions/removals are isolated per test
StateBag = new ConcurrentDictionary<string, object?>(testBuilderContext.StateBag),

Summary

The fix correctly resolves the reported bug. The only actionable items worth considering are:

1. Potential (non-critical): Apply the same StateBag copy to the SkipIfEmpty skipped-test context creation paths (lines 481 and 540) for consistency, even though the practical impact is minimal.
2. Awareness: The streaming path and the BuildTestsFromMetadataAsync path now have subtly different StateBag sharing semantics (reference vs. copy). This is unlikely to cause user-visible issues given how StateBag is typically used, but a short doc comment near TestBuilderContext.StateBag explaining the intended lifecycle would help future maintainers.

Overall this is a good, well-scoped fix. The regression test covers the reported scenario clearly.

[claude]

Code Review

The fix is well-targeted and correctly addresses the root cause: when contextAccessor.Current was replaced with a fresh TestBuilderContext at each data combination (line 282), any StateBag values set by the data source method were lost because they lived on the old context. The copy-on-fork pattern using RawStateBag is the right approach.

Positives

• RawStateBag internal property is a clean way to peek at the lazy field without materializing an empty dictionary.
• The is { IsEmpty: false } guard avoids an unnecessary allocation when the bag is unused — consistent with the "lazy-init for performance" philosophy in this file.
• Regression test correctly exercises the exact scenario from the issue.

Concerns

1. Three-way duplication of the same propagation pattern

The same 4-line pattern appears at lines ~290, ~492, and ~555:

if (testBuilderContext.RawStateBag is { IsEmpty: false } parentBag)
    targetContext.StateBag = new ConcurrentDictionary<string, object?>(parentBag);

Extracting this to a small helper would reduce noise and make future propagation sites consistent:

private static void PropagateStateBag(TestBuilderContext source, TestBuilderContext target)
{
    if (source.RawStateBag is { IsEmpty: false } bag)
        target.StateBag = new ConcurrentDictionary<string, object?>(bag);
}

2. Inconsistency between copy (new PR) and shared-reference (existing code)

The new fix correctly copies the parent StateBag into each per-iteration context. But then at lines ~419 and ~1868, gets a shared reference to :

StateBag = contextAccessor.Current.StateBag,  // shared reference, not a copy

This means the per-iteration context and the test-specific context share the same dictionary instance. Any writes to StateBag during test execution affect the per-iteration context too. This is pre-existing rather than introduced here, but the two strategies (copy vs shared reference) in the same file are worth documenting to avoid future confusion about the intended semantics.

3. Skip paths copy from the wrong source

In the skip paths (lines ~492, ~555), the copy sources from (the original discovery context). That's correct if the data source method never ran (which is true for SkipIfEmpty). However, if a scenario exists where the data source ran and modified StateBag and the test is still skipped, the fix would use stale state. For the current SkipIfEmpty use case this is fine — just worth a comment explaining why we copy from the parent rather than .

4. Streaming path (BuildSingleTestAsync) not updated

The second BuildTestsFromMethodDataAsync overload (returning IAsyncEnumerable, ~line 1535) calls BuildSingleTestAsync at line 1662 but never resets contextAccessor.Current between iterations the way the original path does at line 282. This means in the streaming path, all tests built from the same metadata share the same baseContext.StateBag instance. This is a pre-existing inconsistency but worth verifying it doesn't surface the same original bug under different scenarios.

5. Test coverage for skip paths

The regression test covers the happy path. The SkipIfEmpty StateBag propagation (also changed) isn't covered. A test verifying StateBag availability when a test is marked skipped-because-data-source-was-empty would lock in those code paths too.

Summary

The fix is correct for the primary reported issue and the RawStateBag internal API is a clean addition. The main improvement opportunity is extracting the repeated propagation pattern into a helper and adding a test for the skip paths. The streaming-path inconsistency is worth investigating but is out of scope for this fix.

[claude]

Code Review

The fix is well-targeted and correctly addresses the root cause: when contextAccessor.Current was replaced with a fresh TestBuilderContext at each data combination (line 282), any StateBag values set by the data source method were lost because they lived on the old context. The copy-on-fork pattern using RawStateBag is the right approach.

Positives

• RawStateBag internal property is a clean way to peek at the lazy field without materializing an empty dictionary.
• The is { IsEmpty: false } guard avoids an unnecessary allocation when the bag is unused — consistent with the "lazy-init for performance" philosophy in this file.
• Regression test correctly exercises the exact scenario from the issue.

Concerns

1. Three-way duplication of the same propagation pattern

The same 4-line pattern appears at three different sites in TestBuilder.cs. Extracting this to a small helper would reduce noise and make future propagation sites consistent:

private static void PropagateStateBag(TestBuilderContext source, TestBuilderContext target)
{
    if (source.RawStateBag is { IsEmpty: false } bag)
        target.StateBag = new ConcurrentDictionary<string, object?>(bag);
}

2. Inconsistency between copy (new PR) and shared-reference (existing code)

The new fix correctly copies the parent StateBag into each per-iteration context. But then at lines ~419 and ~1868, testSpecificContext gets a shared reference rather than a copy:

StateBag = contextAccessor.Current.StateBag,  // shared reference, not a copy

This means the per-iteration context and the test-specific context share the same dictionary instance. Any writes to StateBag during test execution affect the per-iteration context too. This is pre-existing rather than introduced here, but having two strategies (copy vs shared reference) in the same file is worth documenting to avoid future confusion about the intended semantics.

3. Skip paths copy from the original discovery context, not the per-iteration context

In the skip paths (lines ~492, ~555), the copy sources from testBuilderContext.RawStateBag (the original discovery context). That's correct for SkipIfEmpty since the data source method never ran. However, if a scenario exists where the data source ran and modified StateBag and the test is still skipped, the fix would use stale state. For the current SkipIfEmpty use case this is fine — but a comment explaining why we copy from the parent rather than contextAccessor.Current would prevent future confusion.

4. Streaming path (BuildSingleTestAsync) not updated

The second BuildTestsFromMethodDataAsync overload (returning IAsyncEnumerable, ~line 1535) calls BuildSingleTestAsync but never resets contextAccessor.Current between iterations the way the original path does at line 282. This means in the streaming path, all tests share the same baseContext.StateBag instance (a pre-existing inconsistency). Worth verifying this path doesn't manifest the same original bug under different scenarios.

5. Test coverage for skip paths

The regression test covers the happy path only. The SkipIfEmpty StateBag propagation (also changed) isn't covered. A test verifying StateBag availability when a test is skipped because its data source returned no data would lock in those code paths.

Summary

The fix is correct for the primary reported issue and RawStateBag is a clean addition. The main improvement opportunities are: extracting the repeated propagation pattern into a helper, adding tests for the skip paths, and documenting the copy-vs-reference semantics distinction.

[claude]

Follow-up Review

The author has responded to all three concerns from the previous review rounds. Here is an assessment of the current state of the PR.

What was addressed

Duplication eliminated — The repeated 4-line propagation pattern has been correctly extracted into PropagateStateBag(TestBuilderContext source, TestBuilderContext target). The helper is clean, focused, and uses the RawStateBag / IsEmpty guard to avoid a ConcurrentDictionary allocation in the common case where StateBag was never touched.

Skip-if-empty paths covered — PropagateStateBag is now called before the skipped-test contexts at lines ~492 and ~555, making all three creation sites consistent.

RawStateBag internal API — Exposing the backing field via internal ConcurrentDictionary<string, object?>? RawStateBag => _stateBag; is the right approach: it avoids forcing the lazy initializer to run just to do an existence check, keeping allocations minimal.

One remaining observation (non-blocking)

The streaming path (BuildSingleTestAsync, called from the IAsyncEnumerable-returning overload of BuildTestsFromMethodDataAsync) does not reset contextAccessor.Current between iterations and therefore never calls PropagateStateBag. All tests built from that path share the same baseContext.StateBag instance by reference rather than receiving an independent copy. This is a pre-existing design difference, not introduced by this PR, and is unlikely to cause problems given typical StateBag usage patterns. It is worth tracking as a known semantic difference between the two build paths.

Summary

The fix is correct, the three-commit iteration has addressed all review feedback cleanly, and the code is in good shape to merge. The streaming-path difference noted above is worth a follow-up issue rather than blocking this fix.

[claude]

Follow-up Review

Previous review rounds raised three concerns: (1) extraction of the repeated propagation pattern, (2) coverage of the SkipIfEmpty paths, and (3) the streaming-path inconsistency. Here is an assessment of the current diff against those concerns.

What was addressed

Pattern extraction — The repeated copy logic has been moved into TestBuilderContext.CopyStateBagTo(TestBuilderContext target). This is cleaner than the previously suggested static helper because it keeps the logic with the owning type and encapsulates access to _stateBag (no need to expose RawStateBag publicly). The is { IsEmpty: false } guard is preserved, so the allocation is still avoided when the bag was never touched.

All three context-creation sites are now consistent — CopyStateBagTo is called immediately after each of the three new TestBuilderContext { ... } blocks (line 290, 489, 549), making the propagation rule uniform across the normal iteration path and both SkipIfEmpty paths.

Regression test — The StateBagDataSourcePropagationTests class reproduces the exact failure scenario with two keys and two data rows, which is the right level of specificity.

Remaining observation (pre-existing, non-blocking)

The streaming build path (BuildSingleTestAsync, called from the IAsyncEnumerable-returning overload of BuildTestsFromMethodDataAsync) still does not call CopyStateBagTo between iterations. In that path, contextAccessor.Current is never replaced with a per-iteration copy, so all streaming tests share a single StateBag instance by reference. This is a pre-existing design difference that is out of scope for this fix, but it is worth tracking so a user who hits it doesn't have to go through the same diagnosis. A follow-up issue or a // Note: comment near that path would help future maintainers.

Minor nit

CopyStateBagTo creates a shallow copy (new ConcurrentDictionary<string, object?>(bag)). Mutable reference-type values stored in the bag will still be shared across all per-test copies. This is almost certainly fine for the intended use cases, but a short doc comment on the method (or on StateBag) noting that values are not deep-cloned would prevent surprises.

Summary

All actionable feedback from the previous review rounds has been addressed. The fix is correct, minimal, and consistent across all affected code paths. The CopyStateBagTo encapsulation on TestBuilderContext is a cleaner design than the static helper that was suggested earlier. The streaming-path difference is a known pre-existing issue and does not block this merge.