diff --git a/src/testUtils/mockUseKeyedState.js b/src/testUtils/mockUseKeyedState.js
index 839b358..3d161c0 100644
--- a/src/testUtils/mockUseKeyedState.js
+++ b/src/testUtils/mockUseKeyedState.js
@@ -122,8 +122,10 @@ export class MockUseKeyedState {
    * @param {any} val - new value to be returned by the useState call.
    */
   mockVal(mockKey, val) {
-    this.hookSpy.mockImplementationOnce((key) => {
+    this.hookSpy.mockImplementationOnce((key, newVal) => {
       if (key === mockKey) {
+        this.values[key] = val;
+        this.initValues[key] = newVal;
         return [val, this.setState[key]];
       }
       return this.mockHook(key);
@@ -136,8 +138,10 @@ export class MockUseKeyedState {
    * @param {object} mapping - { <stateKey>: <val to return> }
    */
   mockVals(mapping) {
-    this.hookSpy.mockImplementation((key) => {
+    this.hookSpy.mockImplementation((key, val) => {
       if (mapping[key]) {
+        this.values[key] = mapping[key];
+        this.initValues[key] = val;
         return [mapping[key], this.setState[key]];
       }
       return this.mockHook(key);