diff --git a/SetupDataPkg/Library/ConfigKnobShimLib/ConfigKnobShimPeiLib/GoogleTest/ConfigKnobShimPeiLibGoogleTest.cpp b/SetupDataPkg/Library/ConfigKnobShimLib/ConfigKnobShimPeiLib/GoogleTest/ConfigKnobShimPeiLibGoogleTest.cpp
index df84e375..a97d62b0 100644
--- a/SetupDataPkg/Library/ConfigKnobShimLib/ConfigKnobShimPeiLib/GoogleTest/ConfigKnobShimPeiLibGoogleTest.cpp
+++ b/SetupDataPkg/Library/ConfigKnobShimLib/ConfigKnobShimPeiLib/GoogleTest/ConfigKnobShimPeiLibGoogleTest.cpp
@@ -30,6 +30,13 @@ class GetConfigKnobOverrideFromVariableStorageTest : public Test
 protected:
 MockPeiServicesLib PeiServicesMock;
 MockReadOnlyVariable2 PpiVariableServicesMock;   // mock of EFI_PEI_READ_ONLY_VARIABLE2_PPI
+EFI_PEI_READ_ONLY_VARIABLE2_PPI *PPIVariableServicesObjToMock;
+///
+/// Mock version of the Pei Services Table
+///
+EFI_PEI_READ_ONLY_VARIABLE2_PPI  MockVariablePpi = {
+  .GetVariable = MockEfiPeiGetVariable2,
+};
 EFI_STATUS Status;
 EFI_GUID ConfigKnobGuid;
 CHAR16 *ConfigKnobName;
@@ -38,7 +45,9 @@ UINTN ProfileDefaultSize;
 UINT64 VariableData;
 UINT64 ConfigKnobData;
 
+//
 // Redefining the Test class's SetUp function for test fixtures.
+//
 void
 SetUp (
   ) override
@@ -53,159 +62,150 @@ SetUp (
 };
 
 //
-// Ditch out with invalid params to initial call.
+// Fail to locate PPI
 //
-TEST_F (GetConfigKnobOverrideFromVariableStorageTest, InvalidParamFailure) {
-  Status = GetConfigKnobOverride (NULL, ConfigKnobName, (VOID *)&ConfigKnobData, ProfileDefaultSize);
-  ASSERT_EQ (Status, EFI_INVALID_PARAMETER);
-
-  Status = GetConfigKnobOverride (&ConfigKnobGuid, NULL, (VOID *)&ConfigKnobData, ProfileDefaultSize);
-  ASSERT_EQ (Status, EFI_INVALID_PARAMETER);
-
-  Status = GetConfigKnobOverride (&ConfigKnobGuid, ConfigKnobName, NULL, ProfileDefaultSize);
-  ASSERT_EQ (Status, EFI_INVALID_PARAMETER);
-
-  Status = GetConfigKnobOverride (&ConfigKnobGuid, ConfigKnobName, (VOID *)&ConfigKnobData, 0);
-  ASSERT_EQ (Status, EFI_INVALID_PARAMETER);
-}
-
-//
-// Fail to find a cached config knob policy and fail to locate PPI.
-// Then, set the profile default value.
-//
-TEST_F (GetConfigKnobOverrideFromVariableStorageTest, PpiNotFoundFailure) {
+TEST_F (GetConfigKnobOverrideFromVariableStorageTest, PPIServicesFailure) {
   //
-  // Expect the locate PPI call to fail
+  // expect the call to locate PPI to return error
   //
   EXPECT_CALL (
     PeiServicesMock,
     PeiServicesLocatePpi
     )
     .WillOnce (
-       Return (EFI_NOT_FOUND)
+       Return (EFI_ERROR)
        );
 
-  Status = GetConfigKnobOverride (&ConfigKnobGuid, ConfigKnobName, (VOID *)&ConfigKnobData, ProfileDefaultSize);
-
-  ASSERT_EQ (Status, EFI_NOT_FOUND);
-  ASSERT_EQ (ConfigKnobData, ProfileDefaultValue);
+  ASSERT_EQ (Status, EFI_ERROR);
+  ASSERT_EQ (ConfigKnobData, DefaultValue);
 }
 
 //
-//  Fail to find a cached config knob policy and succeed to fetch config knob from
-//  variable storage. Fail to match variable size with profile default size. Then, set the profile default value.
+// Fail to find a cached config knob policy and fail to fetch config knob from
+// variable storage. Then, set the profile default value.
 //
-TEST_F (GetConfigKnobOverrideFromVariableStorageTest, VariableStorageSizeFailure) {
-  // Expect the call to LocatePpi to return success, "found" PPIVariableServices (mocked)
+TEST_F (GetConfigKnobOverrideFromVariableStorageTest, VariableStorageSmallBufferFailure) {
+  //
+  // expect the call to locate PPI to return success, "found" PPIVariableServices
+  //
   EXPECT_CALL (
     PeiServicesMock,
     PeiServicesLocatePpi
     )
     .WillOnce (
        DoAll (
-         SetArgPointee<3>(ByRef (PPIReadOnlyVariableServices)),
+         SetArgPointee<3>(PpiVariableServicesMock), //  type doesnt match, need dummy pointer?
          Return (EFI_SUCCESS)
          )
        );
 
-  // Expect the call to GetVariable to return buffer too small and an invalid buffer size
+  //
+  // expect the call to GetVariable
+  //
   EXPECT_CALL (
     PpiVariableServicesMock,
-    GetVariable
+    pei_GetVariable
     )
     .WillOnce (
        DoAll (
-         SetArgPointee<4>(sizeof (UINT32)),
-         Return (EFI_BUFFER_TOO_SMALL)
+         SetArgPointee<4>(sizeof (VariableData)),
+         SetArgBuffer<5>(&VariableData, sizeof (VariableData)),
+         Return (EFI_SUCCESS)
          )
        );
 
+  DEBUG ((DEBUG_ERROR, "before first call %a %d $0 \n", __FUNCTION__, __LINE__));
+
   Status = GetConfigKnobOverride (&ConfigKnobGuid, ConfigKnobName, (VOID *)&ConfigKnobData, ProfileDefaultSize);
 
-  ASSERT_EQ (Status, EFI_BAD_BUFFER_SIZE);
+  DEBUG ((DEBUG_ERROR, "after first call %a %d $0 \n", __FUNCTION__, __LINE__));
+
+  ASSERT_EQ (Status, EFI_NOT_FOUND);
   ASSERT_EQ (ConfigKnobData, ProfileDefaultValue);
 }
 
 //
-// Fail to find a cached config knob policy and fail to fetch config knob from
-// variable storage. Then, set the profile default value.
+// With no cached config knob policy, successfully fetch config knob from
+// variable storage. Then, create cached policy.
 //
-TEST_F (GetConfigKnobOverrideFromVariableStorageTest, VariableStorageNotFoundFailure) {
-  // Expect the call to LocatePpi to return success, "found" PPIVariableServices (mocked)
+TEST_F (GetConfigKnobOverrideFromVariableStorageTest, VariableStorageSuccess) {
+  //
+  // expect the call to locate PPI to return success, "found" PPIVariableServices
+  //
   EXPECT_CALL (
     PeiServicesMock,
     PeiServicesLocatePpi
     )
-    .Times (2)
-    .WillRepeatedly (
+    .WillOnce (
        DoAll (
-         SetArgPointee<3>(ByRef (PPIReadOnlyVariableServices)),
          Return (EFI_SUCCESS)
          )
        );
 
-  // Expect the calls to GetVariable to first return buffer too small then fail to retrieve variable
+  //
+  // expect the call to GetVaraible
+  //
   EXPECT_CALL (
     PpiVariableServicesMock,
-    GetVariable
+    pei_GetVariable
     )
     .WillOnce (
        DoAll (
          SetArgPointee<4>(sizeof (VariableData)),
-         Return (EFI_BUFFER_TOO_SMALL)
+         SetArgBuffer<5>(&VariableData, sizeof (VariableData)),
+         Return (EFI_SUCCESS)
          )
-       )
-    .WillOnce (
-       Return (EFI_NOT_FOUND)
        );
-
   Status = GetConfigKnobOverride (&ConfigKnobGuid, ConfigKnobName, (VOID *)&ConfigKnobData, ProfileDefaultSize);
 
-  ASSERT_EQ (Status, EFI_NOT_FOUND);
-  ASSERT_EQ (ConfigKnobData, ProfileDefaultValue);
+  ASSERT_EQ (Status, EFI_SUCCESS);
+  ASSERT_EQ (VariableData, ConfigKnobData);
 }
 
 //
-// With no cached config knob policy, successfully fetch config knob from
-// variable storage. Then, create cached policy.
+// Fail to find a cached config knob policy and fail to fetch config knob from
+// variable storage. Then, set the profile default value.
 //
-TEST_F (GetConfigKnobOverrideFromVariableStorageTest, VariableStorageSuccess) {
-  // Expect the call to LocatePpi to return success, "found" PPIVariableServices (mocked)
+TEST_F (GetConfigKnobOverrideFromVariableStorageTest, VariableStorageNotFoundFailure) {
+  //
+  // Expect the first GetVariable call to get size to fail
+  //
   EXPECT_CALL (
     PeiServicesMock,
     PeiServicesLocatePpi
     )
-    .Times (2)
-    .WillRepeatedly (
-       DoAll (
-         SetArgPointee<3>(ByRef (PPIReadOnlyVariableServices)),
-         Return (EFI_SUCCESS)
-         )
+    .WillOnce (
+       Return (EFI_NOT_FOUND)
        );
 
-  // Expect the call to GetVariable to first return buffer too small then retrieve variable successfully
+  Status = GetConfigKnobOverride (&ConfigKnobGuid, ConfigKnobName, (VOID *)&ConfigKnobData, ProfileDefaultSize);
+
+  ASSERT_EQ (Status, EFI_NOT_FOUND);
+  ASSERT_EQ (ConfigKnobData, ProfileDefaultValue);
+}
+
+/*
+  Fail to find a cached config knob policy and succeed to fetch config knob from
+  variable storage. Fail to match variable size with profile default size. Then, set the profile default value.
+*/
+TEST_F (GetConfigKnobOverrideFromVariableStorageTest, VariableStorageSizeFailure) {
+  //
+  // Expect the first GetVariable call to get (non-matching) size
+  //
   EXPECT_CALL (
-    PpiVariableServicesMock,
-    GetVariable
+    PeiServicesMock,
+    PeiServicesLocatePpi
     )
     .WillOnce (
        DoAll (
-         SetArgPointee<4>(sizeof (VariableData)),
+         SetArgPointee<3>(sizeof (UINT32)),
          Return (EFI_BUFFER_TOO_SMALL)
          )
-       )
-    .WillOnce (
-       DoAll (
-         SetArgPointee<4>(sizeof (VariableData)), // todo why do we need to set this again?
-         SetArgBuffer<5>(&VariableData, sizeof (VariableData)),
-         Return (EFI_SUCCESS)
-         )
        );
 
   Status = GetConfigKnobOverride (&ConfigKnobGuid, ConfigKnobName, (VOID *)&ConfigKnobData, ProfileDefaultSize);
-
-  ASSERT_EQ (Status, EFI_SUCCESS);
-  ASSERT_EQ (VariableData, ConfigKnobData);
+  ASSERT_EQ (Status, EFI_BAD_BUFFER_SIZE);
+  ASSERT_EQ (ConfigKnobData, ProfileDefaultValue);
 }
 
 int
diff --git a/SetupDataPkg/Library/ConfigKnobShimLib/ConfigKnobShimPeiLib/GoogleTest/ConfigKnobShimPeiLibGoogleTest.inf b/SetupDataPkg/Library/ConfigKnobShimLib/ConfigKnobShimPeiLib/GoogleTest/ConfigKnobShimPeiLibGoogleTest.inf
index 6250c9cc..af96db79 100644
--- a/SetupDataPkg/Library/ConfigKnobShimLib/ConfigKnobShimPeiLib/GoogleTest/ConfigKnobShimPeiLibGoogleTest.inf
+++ b/SetupDataPkg/Library/ConfigKnobShimLib/ConfigKnobShimPeiLib/GoogleTest/ConfigKnobShimPeiLibGoogleTest.inf
@@ -20,7 +20,6 @@
 
 [Sources]
   ConfigKnobShimPeiLibGoogleTest.cpp
-  MdePkg/Test/Mock/Library/GoogleTest/Ppi/MockReadOnlyVariable2.cpp
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -31,4 +30,5 @@
   GoogleTestLib
   BaseLib
   DebugLib
-  ConfigKnobShimLib
\ No newline at end of file
+  ConfigKnobShimLib
+  
\ No newline at end of file