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]