From f683e8484cbe68d3cd43c4e90efe0e54ece5e9f9 Mon Sep 17 00:00:00 2001
From: Addam Boord <addam.boord@ketiv.com>
Date: Fri, 22 May 2026 15:47:42 -0500
Subject: [PATCH] Fix key validation and align failing test expectations

KeyHelper.ValidateKey wrapped ArgumentNullException.ThrowIfNull inside an
IsNullOrWhiteSpace check, so empty and whitespace keys were silently
accepted while null keys threw the wrong exception type. Split the
validation: null now throws ArgumentNullException and empty/whitespace
throws ArgumentException with nameof(key) -- matching the .NET BCL
convention -- and the XML documentation now lists both exception types.

Test fixes:
- StateStoreTests: two null-key cases now expect ArgumentNullException
  (and are renamed to reflect that), matching NullReferenceTests and the
  new contract.
- NullReferenceTests.UpsertAsync_NullKey_ThrowsAsync passed "" instead of
  null despite its name; corrected to null! so the assertion matches the
  test's intent.
- EdgeCaseTests.Get_NonExistentKey_ThrowsAsync asserted a KeyNotFoundException
  that StateStoreImplementation.GetAsync has never thrown -- missing keys
  return default, as verified by GetAsync_ReturnsDefault_WhenKeyDoesNotExist.
  Renamed and rewritten to verify default(int) is returned, preserving the
  value-type-specific coverage.

All 72 tests pass on net10.0.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 src/StateStore/Internal/KeyHelper.cs         | 6 ++++--
 tests/StateStore.Tests/EdgeCaseTests.cs      | 6 ++++--
 tests/StateStore.Tests/NullReferenceTests.cs | 2 +-
 tests/StateStore.Tests/StateStoreTests.cs    | 8 ++++----
 4 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/StateStore/Internal/KeyHelper.cs b/src/StateStore/Internal/KeyHelper.cs
index b863b20..118c24c 100644
--- a/src/StateStore/Internal/KeyHelper.cs
+++ b/src/StateStore/Internal/KeyHelper.cs
@@ -9,12 +9,14 @@ internal static class KeyHelper
     /// Validates that the key is non-null, non-empty, and non-whitespace.
     /// </summary>
     /// <param name="key">The key to validate.</param>
-    /// <exception cref="ArgumentException">Thrown when the key is null, empty, or whitespace.</exception>
+    /// <exception cref="ArgumentNullException">Thrown when the key is <see langword="null"/>.</exception>
+    /// <exception cref="ArgumentException">Thrown when the key is empty or whitespace.</exception>
     public static void ValidateKey(string key)
     {
+        ArgumentNullException.ThrowIfNull(key);
         if (string.IsNullOrWhiteSpace(key))
         {
-            ArgumentNullException.ThrowIfNull(key, "State store key must not be null, empty, or whitespace.");
+            throw new ArgumentException("State store key must not be empty or whitespace.", nameof(key));
         }
     }
 
diff --git a/tests/StateStore.Tests/EdgeCaseTests.cs b/tests/StateStore.Tests/EdgeCaseTests.cs
index 3adc9e8..2bd3d6c 100644
--- a/tests/StateStore.Tests/EdgeCaseTests.cs
+++ b/tests/StateStore.Tests/EdgeCaseTests.cs
@@ -48,13 +48,15 @@ public async Task SetAndGet_EmptyArray_SucceedsAsync()
     }
 
     [Fact]
-    public async Task Get_NonExistentKey_ThrowsAsync()
+    public async Task Get_NonExistentKey_ReturnsDefault_Async()
     {
         var provider = new InMemoryStorageProvider();
         var serializer = new JsonStateSerializer();
         var pipeline = new MiddlewarePipeline([], provider);
         var store = new StateStoreImplementation(serializer, pipeline);
 
-        await Assert.ThrowsAsync<KeyNotFoundException>(async () => await store.GetAsync<int>("does_not_exist", TestContext.Current.CancellationToken));
+        var result = await store.GetAsync<int>("does_not_exist", TestContext.Current.CancellationToken);
+
+        Assert.Equal(0, result);
     }
 }
diff --git a/tests/StateStore.Tests/NullReferenceTests.cs b/tests/StateStore.Tests/NullReferenceTests.cs
index ad3aae3..8498806 100644
--- a/tests/StateStore.Tests/NullReferenceTests.cs
+++ b/tests/StateStore.Tests/NullReferenceTests.cs
@@ -36,7 +36,7 @@ public async Task UpsertAsync_NullKey_ThrowsAsync()
         var pipeline = new MiddlewarePipeline([], provider);
         var store = new StateStoreImplementation(serializer, pipeline);
 
-        await Assert.ThrowsAsync<ArgumentNullException>(async () => await store.UpsertAsync<int>("", 1, x => x + 1, TestContext.Current.CancellationToken));
+        await Assert.ThrowsAsync<ArgumentNullException>(async () => await store.UpsertAsync<int>(null!, 1, x => x + 1, TestContext.Current.CancellationToken));
     }
 
     [Fact]
diff --git a/tests/StateStore.Tests/StateStoreTests.cs b/tests/StateStore.Tests/StateStoreTests.cs
index ca5089a..3b9534c 100644
--- a/tests/StateStore.Tests/StateStoreTests.cs
+++ b/tests/StateStore.Tests/StateStoreTests.cs
@@ -117,9 +117,9 @@ public async Task SetAsync_WorksWithComplexTypes_Async()
     }
 
     [Fact]
-    public async Task GetAsync_ThrowsArgumentException_ForNullKey_Async()
+    public async Task GetAsync_ThrowsArgumentNullException_ForNullKey_Async()
     {
-        await Assert.ThrowsAsync<ArgumentException>(() => _store.GetAsync<string>(null!, TestContext.Current.CancellationToken).AsTask());
+        await Assert.ThrowsAsync<ArgumentNullException>(() => _store.GetAsync<string>(null!, TestContext.Current.CancellationToken).AsTask());
     }
 
     [Fact]
@@ -135,9 +135,9 @@ public async Task GetAsync_ThrowsArgumentException_ForWhitespaceKey_Async()
     }
 
     [Fact]
-    public async Task SetAsync_ThrowsArgumentException_ForNullKey_Async()
+    public async Task SetAsync_ThrowsArgumentNullException_ForNullKey_Async()
     {
-        await Assert.ThrowsAsync<ArgumentException>(() => _store.SetAsync(null!, "value", TestContext.Current.CancellationToken).AsTask());
+        await Assert.ThrowsAsync<ArgumentNullException>(() => _store.SetAsync(null!, "value", TestContext.Current.CancellationToken).AsTask());
     }
 
     [Fact]