From aeef939c75cb8b9c5c40da7ff94a46ccefdd594f Mon Sep 17 00:00:00 2001
From: Ben Warzeski <bwarzesk@gmail.com>
Date: Thu, 6 Jul 2023 14:08:13 +0000
Subject: [PATCH] fix: load mock values on init for mockUseKeyedState

---
 src/testUtils/mockUseKeyedState.js | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

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);