[PATCH v1 30/41] PrmPkg/Library/DxePrmContextBufferLib: Add host-based unit tests

["Michael Kubacki" <mikuback@linux.microsoft.com>]

From: Michael Kubacki <michael.kubacki@microsoft.com>

Adds host-based unit tests for DxePrmContextBufferLib. This tests
each function in DxePrmContextBufferLib.

Any changes to DxePrmContextBufferLib moving forward must pass the unit
tests. The unit tests should be updated as needed if the library API
changes.

Cc: Andrew Fish <afish@apple.com>
Cc: Kang Gao <kang.gao@intel.com>
Cc: Michael D Kinney <michael.d.kinney@intel.com>
Cc: Michael Kubacki <michael.kubacki@microsoft.com>
Cc: Leif Lindholm <leif@nuviainc.com>
Cc: Benjamin You <benjamin.you@intel.com>
Cc: Liu Yun <yun.y.liu@intel.com>
Cc: Ankit Sinha <ankit.sinha@intel.com>
Cc: Nate DeSimone <nathaniel.l.desimone@intel.com>
Signed-off-by: Michael Kubacki <michael.kubacki@microsoft.com>
---
 PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTest.c       | 649 ++++++++++++++++++++
 PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTestHost.inf |  46 ++
 PrmPkg/Test/PrmPkgHostTest.dsc                                                        |   6 +
 3 files changed, 701 insertions(+)

diff --git a/PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTest.c b/PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTest.c
new file mode 100644
index 000000000000..07e247305880
--- /dev/null
+++ b/PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTest.c
@@ -0,0 +1,649 @@
+/** @file
+
+  Unit tests for the PRM Context Buffer Library.
+
+  Copyright (c) Microsoft Corporation
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <stdio.h>
+#include <string.h>
+#include <stdarg.h>
+#include <stddef.h>
+#include <setjmp.h>
+#include <cmocka.h>
+
+#include <Uefi.h>
+
+#include <Guid/ZeroGuid.h>
+#include <Library/BaseLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/DebugLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/PrintLib.h>
+#include <Library/PrmContextBufferLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/UnitTestLib.h>
+#include <Protocol/PrmConfig.h>
+
+#define UNIT_TEST_NAME          "PRM Context Buffer Library Unit Test"
+#define UNIT_TEST_VERSION       "0.1"
+
+///=== TEST DATA ==================================================================================
+
+EFI_HANDLE  mTestProtocolHandle;
+
+//*----------------------------------------------------------------------------------*
+//* Context Structures                                                               *
+//*----------------------------------------------------------------------------------*
+typedef struct {
+  PRM_MODULE_CONTEXT_BUFFERS  *InstallationStructure;
+  EFI_HANDLE                  Handle;
+  PRM_GUID_SEARCH_TYPE        GuidSearchType;
+  EFI_GUID                    *Guid;
+  EFI_GUID                    *ExpectedModuleGuid;
+  EFI_STATUS                  ExpectedStatus;
+} PRM_CONTEXT_BUFFERS_TEST_CONTEXT;
+
+typedef struct {
+  EFI_GUID                    *HandlerGuid;
+  PRM_MODULE_CONTEXT_BUFFERS  *ContextBuffers;
+  PRM_CONTEXT_BUFFER          *ExpectedContextBuffer;
+  EFI_STATUS                  ExpectedStatus;
+} PRM_CONTEXT_BUFFER_TEST_CONTEXT;
+
+//*----------------------------------------------------------------------------------*
+//* Test GUIDs                                                                       *
+//*----------------------------------------------------------------------------------*
+
+// {52960b90-2f3a-4917-b91a-ed5f599a8809}
+#define HANDLER_TEST_GUID_1 {0x52960b90, 0x2f3a, 0x4917, { 0xb9, 0x1a, 0xed, 0x5f, 0x59, 0x9a, 0x88, 0x09 }}
+EFI_GUID mHandlerTestGuid1 = HANDLER_TEST_GUID_1;
+
+// {9316a80d-06dc-417b-b21d-6b3c2ae4ed6f}
+#define HANDLER_TEST_GUID_2 {0x9316a80d, 0x06dc, 0x417b, { 0xb2, 0x1d, 0x6b, 0x3c, 0x2a, 0xe4, 0xed, 0x6f }}
+EFI_GUID mHandlerTestGuid2 = HANDLER_TEST_GUID_2;
+
+// {d32ac8ba-6cc6-456f-9ed9-9233fa310434}
+#define HANDLER_TEST_GUID_3 {0xd32ac8ba, 0x6cc6, 0x456f, { 0x9e, 0xd9, 0x92, 0x33, 0xfa, 0x31, 0x04, 0x34 }}
+EFI_GUID mHandlerTestGuid3 = HANDLER_TEST_GUID_3;
+
+// {faadaa95-070b-4a34-a919-18305dc07370}
+#define MODULE_TEST_GUID_1 {0xfaadaa95, 0x070b, 0x4a34, { 0xa9, 0x19, 0x18, 0x30, 0x5d, 0xc0, 0x73, 0x70 }}
+EFI_GUID mModuleTestGuid1 = MODULE_TEST_GUID_1;
+
+// {0ea24584-731c-4863-9100-75780af509a7}
+#define MODULE_TEST_GUID_2 {0x0ea24584, 0x731c, 0x4863, { 0x91, 0x00, 0x75, 0x78, 0x0a, 0xf5, 0x09, 0xa7 }}
+EFI_GUID mModuleTestGuid2 = MODULE_TEST_GUID_2;
+
+// {f456b7a1-82a6-4427-8486-87e3a602df43}
+#define MODULE_TEST_GUID_3 {0xf456b7a1, 0x82a6, 0x4427, { 0x84, 0x86, 0x87, 0xe3, 0xa6, 0x02, 0xdf, 0x43 }}
+EFI_GUID mModuleTestGuid3 = MODULE_TEST_GUID_3;
+
+// {4a941a9c-9dcf-471b-94b5-d9e2d8c64a1b}
+#define NEGATIVE_TEST_GUID {0x4a941a9c, 0x9dcf, 0x471b,  {0x94, 0xb5, 0xd9, 0xe2, 0xd8, 0xc6, 0x4a, 0x1b }}
+EFI_GUID mNegativeTestGuid = NEGATIVE_TEST_GUID;
+
+//*----------------------------------------------------------------------------------*
+//* PRM Static Test Structures                                                       *
+//*----------------------------------------------------------------------------------*
+
+PRM_DATA_BUFFER mTestStaticDataBuffer1 = {
+  {
+    PRM_DATA_BUFFER_HEADER_SIGNATURE,
+    sizeof (PRM_DATA_BUFFER)
+  }
+  // No data in the buffer (only a header)
+};
+
+PRM_CONTEXT_BUFFER mTestPrmContextBuffer1 = {
+  PRM_CONTEXT_BUFFER_SIGNATURE,             // Signature
+  PRM_CONTEXT_BUFFER_INTERFACE_VERSION,     // Version
+  0,                                        // Reserved
+  HANDLER_TEST_GUID_1,                      // HandlerGuid
+  &mTestStaticDataBuffer1                   // StaticDataBuffer
+};
+
+PRM_CONTEXT_BUFFER mTestPrmContextBuffer2[2] = {
+  // Context buffer #1
+  {
+    PRM_CONTEXT_BUFFER_SIGNATURE,           // Signature
+    PRM_CONTEXT_BUFFER_INTERFACE_VERSION,   // Version
+    0,                                      // Reserved
+    HANDLER_TEST_GUID_2,                    // HandlerGuid
+    NULL                                    // StaticDataBuffer
+  },
+  // Context buffer #2
+  {
+    PRM_CONTEXT_BUFFER_SIGNATURE,           // Signature
+    PRM_CONTEXT_BUFFER_INTERFACE_VERSION,   // Version
+    0,                                      // Reserved
+    HANDLER_TEST_GUID_3,                    // HandlerGuid
+    &mTestStaticDataBuffer1                 // StaticDataBuffer (reuse buffer StaticDataBuffer1)
+  }
+};
+
+PRM_MODULE_CONTEXT_BUFFERS mTestPrmModuleContextBuffers1 = {
+  MODULE_TEST_GUID_1,
+  1,
+  &mTestPrmContextBuffer1,
+  NULL
+};
+
+PRM_MODULE_CONTEXT_BUFFERS mTestPrmModuleContextBuffers2 = {
+  MODULE_TEST_GUID_2,
+  1,
+  &mTestPrmContextBuffer1,
+  NULL
+};
+
+PRM_MODULE_CONTEXT_BUFFERS mTestPrmModuleContextBuffers3 = {
+  MODULE_TEST_GUID_3,
+  2,
+  &mTestPrmContextBuffer2[0],
+  NULL
+};
+
+
+//*----------------------------------------------------------------------------------*
+//* Test Contexts                                                                    *
+//*----------------------------------------------------------------------------------*
+
+//* Searches by module GUID *
+//                                                   +--------------------------------+--------+----------------+--------------------+--------------------+--------------------+
+//                                                   + InstallationStructure          | Handle | GuidSearchType | Guid               | ExpectedModuleGuid | ExpectedStatus     |
+//                                                   +--------------------------------+--------+----------------+--------------------+--------------------+--------------------+
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers00 = { &mTestPrmModuleContextBuffers1,  NULL,    ByModuleGuid,    &mModuleTestGuid1,   &mModuleTestGuid1,   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers01 = { &mTestPrmModuleContextBuffers2,  NULL,    ByModuleGuid,    &mModuleTestGuid2,   &mModuleTestGuid2,   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers02 = { &mTestPrmModuleContextBuffers3,  NULL,    ByModuleGuid,    &mModuleTestGuid3,   &mModuleTestGuid3,   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers03 = { &mTestPrmModuleContextBuffers3,  NULL,    ByModuleGuid,    &mNegativeTestGuid,  &gZeroGuid,          EFI_NOT_FOUND      };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers04 = { &mTestPrmModuleContextBuffers1,  NULL,    ByModuleGuid,    &gZeroGuid,          &gZeroGuid,          EFI_NOT_FOUND      };
+
+//* Searches by handler GUID *
+//                                                   +--------------------------------+--------+----------------+--------------------+--------------------+--------------------+
+//                                                   + InstallationStructure          | Handle | GuidSearchType | Guid               | ExpectedModuleGuid | ExpectedStatus     |
+//                                                   +--------------------------------+--------+----------------+--------------------+--------------------+--------------------+
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers05 = { &mTestPrmModuleContextBuffers1,  NULL,    ByHandlerGuid,   &mHandlerTestGuid1,  &mModuleTestGuid1,   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers06 = { &mTestPrmModuleContextBuffers1,  NULL,    ByHandlerGuid,   &gZeroGuid,          &gZeroGuid,          EFI_NOT_FOUND      };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers07 = { &mTestPrmModuleContextBuffers2,  NULL,    ByHandlerGuid,   &mHandlerTestGuid1,  &mModuleTestGuid2,   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers08 = { &mTestPrmModuleContextBuffers2,  NULL,    ByHandlerGuid,   &mNegativeTestGuid,  &gZeroGuid,          EFI_NOT_FOUND      };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers09 = { &mTestPrmModuleContextBuffers3,  NULL,    ByHandlerGuid,   &mHandlerTestGuid1,  &gZeroGuid,          EFI_NOT_FOUND      };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers10 = { &mTestPrmModuleContextBuffers3,  NULL,    ByHandlerGuid,   &mHandlerTestGuid2,  &mModuleTestGuid3,   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers11 = { &mTestPrmModuleContextBuffers3,  NULL,    ByHandlerGuid,   &mHandlerTestGuid3,  &mModuleTestGuid3,   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT mContextBuffers12 = { &mTestPrmModuleContextBuffers3,  NULL,    ByHandlerGuid,   &gZeroGuid,          &gZeroGuid,          EFI_NOT_FOUND      };
+
+PRM_CONTEXT_BUFFERS_TEST_CONTEXT *mContextBuffersArray[] = {
+  &mContextBuffers00,
+  &mContextBuffers01,
+  &mContextBuffers02,
+  &mContextBuffers03,
+  &mContextBuffers04,
+  &mContextBuffers05,
+  &mContextBuffers06,
+  &mContextBuffers07,
+  &mContextBuffers08,
+  &mContextBuffers09,
+  &mContextBuffers10,
+  &mContextBuffers11,
+  &mContextBuffers12
+};
+
+//                                                 +----------------------+----------------------------------+------------------------------------------+--------------------+
+//                                                 + HandlerGuid          | ContextBuffers                   | ExpectedContextBuffer                    | ExpectedStatus     |
+//                                                 +----------------------+----------------------------------+------------------------------------------+--------------------+
+PRM_CONTEXT_BUFFER_TEST_CONTEXT mContextBuffer00 = { &mHandlerTestGuid1,    &mTestPrmModuleContextBuffers1,    &mTestPrmContextBuffer1,                   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFER_TEST_CONTEXT mContextBuffer01 = { &mHandlerTestGuid1,    &mTestPrmModuleContextBuffers2,    &mTestPrmContextBuffer1,                   EFI_SUCCESS        };
+PRM_CONTEXT_BUFFER_TEST_CONTEXT mContextBuffer02 = { &mHandlerTestGuid2,    &mTestPrmModuleContextBuffers3,    &mTestPrmContextBuffer2[0],                EFI_SUCCESS        };
+PRM_CONTEXT_BUFFER_TEST_CONTEXT mContextBuffer03 = { &mHandlerTestGuid3,    &mTestPrmModuleContextBuffers3,    &mTestPrmContextBuffer2[1],                EFI_SUCCESS        };
+PRM_CONTEXT_BUFFER_TEST_CONTEXT mContextBuffer04 = { &mNegativeTestGuid,    &mTestPrmModuleContextBuffers1,    NULL,                                      EFI_NOT_FOUND      };
+PRM_CONTEXT_BUFFER_TEST_CONTEXT mContextBuffer05 = { &gZeroGuid,            &mTestPrmModuleContextBuffers3,    NULL,                                      EFI_NOT_FOUND      };
+
+PRM_CONTEXT_BUFFER_TEST_CONTEXT *mContextBufferArray[] = {
+  &mContextBuffer00,
+  &mContextBuffer01,
+  &mContextBuffer02,
+  &mContextBuffer03,
+  &mContextBuffer04,
+  &mContextBuffer05
+};
+
+///=== HELPER FUNCTIONS ===========================================================================
+
+// None
+
+///=== TEST CASES =================================================================================
+
+///===== BASIC SUITE ==================================================
+
+/**
+  Verifies that passing NULL arguments to all library functions fails with EFI_INVALID_PARAMETER.
+
+  @param[in]  Context             [Optional] An optional context parameter.
+                                  Not used in this unit test.
+
+  @retval  UNIT_TEST_PASSED                      Unit test case prerequisites are met.
+  @retval  UNIT_TEST_ERROR_PREREQUISITE_NOT_MET  Test case should be skipped..
+
+**/
+UNIT_TEST_STATUS
+EFIAPI
+NullPointerArgumentsShouldFailGracefully (
+  IN UNIT_TEST_CONTEXT            Context
+  )
+{
+  EFI_GUID                        Guid;
+  PRM_CONTEXT_BUFFER              *ContextBufferPtr;
+  PRM_MODULE_CONTEXT_BUFFERS      ModuleContextBuffers;
+  PRM_MODULE_CONTEXT_BUFFERS      *ModuleContextBuffersPtr;
+
+  UT_ASSERT_EQUAL (FindContextBufferInModuleBuffers (NULL, NULL, NULL), EFI_INVALID_PARAMETER);
+  UT_ASSERT_EQUAL (FindContextBufferInModuleBuffers (NULL, &ModuleContextBuffers, &ContextBufferPtr), EFI_INVALID_PARAMETER);
+  UT_ASSERT_EQUAL (FindContextBufferInModuleBuffers (&Guid, NULL, &ContextBufferPtr), EFI_INVALID_PARAMETER);
+  UT_ASSERT_EQUAL (FindContextBufferInModuleBuffers (&Guid, &ModuleContextBuffers, NULL), EFI_INVALID_PARAMETER);
+
+  UT_ASSERT_EQUAL (GetModuleContextBuffers (ByModuleGuid, NULL, NULL), EFI_INVALID_PARAMETER);
+  UT_ASSERT_EQUAL (GetModuleContextBuffers (ByModuleGuid, NULL, &ModuleContextBuffersPtr), EFI_INVALID_PARAMETER);
+  UT_ASSERT_EQUAL (GetModuleContextBuffers (ByModuleGuid, &Guid, NULL), EFI_INVALID_PARAMETER);
+
+  UT_ASSERT_EQUAL (GetContextBuffer (NULL, NULL, NULL), EFI_INVALID_PARAMETER);
+  UT_ASSERT_EQUAL (GetContextBuffer (NULL, &ModuleContextBuffers, &ContextBufferPtr), EFI_INVALID_PARAMETER);
+  UT_ASSERT_EQUAL (GetContextBuffer (&Guid, NULL, &ContextBufferPtr), EFI_NOT_FOUND);
+  UT_ASSERT_EQUAL (GetContextBuffer (&Guid, &ModuleContextBuffers, NULL), EFI_INVALID_PARAMETER);
+
+  return UNIT_TEST_PASSED;
+}
+
+///===== FUNCTIONAL CORRECTNESS SUITE ==================================================
+
+/**
+  Functional Correctness pre-requisite function.
+
+  Installs a gPrmConfigProtocolGuid protocol instance as specified by the provided
+  context in preparation for unit test execution
+
+  @param[in]  Context         [Optional] An optional parameter that enables:
+                              A pointer to a PRM_CONTEXT_BUFFERS_TEST_CONTEXT structure with
+                              context information for this unit test.
+
+  @retval  UNIT_TEST_PASSED                      Unit test case prerequisites
+                                                 are met.
+  @retval  UNIT_TEST_ERROR_PREREQUISITE_NOT_MET  Test case should be skipped.
+
+**/
+STATIC
+UNIT_TEST_STATUS
+EFIAPI
+InitializeFunctionalCorrectness (
+  IN  UNIT_TEST_CONTEXT                 Context
+  )
+{
+  EFI_STATUS                            Status;
+  PRM_CONFIG_PROTOCOL                   *PrmConfigProtocol;
+  PRM_MODULE_CONTEXT_BUFFERS            *ModuleContextBuffers;
+  PRM_CONTEXT_BUFFERS_TEST_CONTEXT      *TestContext;
+
+  UT_ASSERT_NOT_NULL (Context);
+  TestContext = (PRM_CONTEXT_BUFFERS_TEST_CONTEXT *) Context;
+  ModuleContextBuffers = TestContext->InstallationStructure;
+
+  PrmConfigProtocol = AllocateZeroPool (sizeof (*PrmConfigProtocol));
+  if (PrmConfigProtocol == NULL) {
+    return UNIT_TEST_ERROR_PREREQUISITE_NOT_MET;
+  }
+
+  CopyGuid (&PrmConfigProtocol->ModuleContextBuffers.ModuleGuid, &ModuleContextBuffers->ModuleGuid);
+  PrmConfigProtocol->ModuleContextBuffers.BufferCount = ModuleContextBuffers->BufferCount;
+  PrmConfigProtocol->ModuleContextBuffers.Buffer = ModuleContextBuffers->Buffer;
+
+  Status =  gBS->InstallProtocolInterface (
+                  &TestContext->Handle,
+                  &gPrmConfigProtocolGuid,
+                  EFI_NATIVE_INTERFACE,
+                  (VOID *) PrmConfigProtocol
+                  );
+  UT_ASSERT_NOT_EFI_ERROR (Status);
+
+  return UNIT_TEST_PASSED;
+}
+
+/**
+  Functional Correctness cleanup function.
+
+  Uninstalls the gPrmConfigProtocolGuid protocol instance as specified by the
+  provided context. This is used to clean up the mocked protocol database after
+  unit test execution.
+
+  @param[in]  Context       [Optional] An optional parameter that enables:
+                            A pointer to a PRM_CONTEXT_BUFFERS_TEST_CONTEXT structure with
+                            context information for this unit test.
+
+  @retval  UNIT_TEST_PASSED                Test case cleanup succeeded.
+  @retval  UNIT_TEST_ERROR_CLEANUP_FAILED  Test case cleanup failed.
+
+**/
+STATIC
+UNIT_TEST_STATUS
+EFIAPI
+DeInitializeFunctionalCorrectness (
+  IN  UNIT_TEST_CONTEXT               Context
+  )
+{
+  EFI_STATUS                          Status;
+  PRM_CONFIG_PROTOCOL                 *PrmConfigProtocol;
+  PRM_CONTEXT_BUFFERS_TEST_CONTEXT    *TestContext;
+
+  UT_ASSERT_NOT_NULL (Context);
+  TestContext = (PRM_CONTEXT_BUFFERS_TEST_CONTEXT *) Context;
+
+  Status = gBS->HandleProtocol (
+                  TestContext->Handle,
+                  &gPrmConfigProtocolGuid,
+                  (VOID **) &PrmConfigProtocol
+                  );
+  UT_ASSERT_NOT_EFI_ERROR (Status);
+
+  if (!EFI_ERROR (Status)) {
+    Status =  gBS->UninstallProtocolInterface (
+                     TestContext->Handle,
+                     &gPrmConfigProtocolGuid,
+                     PrmConfigProtocol
+                    );
+    UT_ASSERT_NOT_EFI_ERROR (Status);
+    if (!EFI_ERROR (Status)) {
+      FreePool (PrmConfigProtocol);
+    }
+  }
+
+  return UNIT_TEST_PASSED;
+}
+
+/**
+  Verifies that the correct PRM_MODULE_CONTEXT_BUFFERS structure instance is found
+  for a given PRM module or PRM handler GUID.
+
+  @param[in]  Context       [Optional] An optional context parameter.
+                            A pointer to a PRM_CONTEXT_BUFFERS_TEST_CONTEXT structure with
+                            context information for this unit test.
+
+  @retval  UNIT_TEST_PASSED                      Unit test case prerequisites are met.
+  @retval  UNIT_TEST_ERROR_PREREQUISITE_NOT_MET  Test case should be skipped..
+
+**/
+UNIT_TEST_STATUS
+EFIAPI
+VerifyGetModuleContextBuffers (
+  IN UNIT_TEST_CONTEXT                Context
+  )
+{
+  EFI_STATUS                          Status;
+  PRM_MODULE_CONTEXT_BUFFERS          *ContextBuffers;
+  PRM_CONTEXT_BUFFERS_TEST_CONTEXT    *TestContext;
+
+  ContextBuffers = NULL;
+  TestContext = (PRM_CONTEXT_BUFFERS_TEST_CONTEXT *) Context;
+
+  Status = GetModuleContextBuffers (TestContext->GuidSearchType, TestContext->Guid, &ContextBuffers);
+  UT_ASSERT_STATUS_EQUAL (Status, TestContext->ExpectedStatus);
+
+  if (!EFI_ERROR (TestContext->ExpectedStatus)) {
+    UT_ASSERT_TRUE (CompareGuid (TestContext->ExpectedModuleGuid, &ContextBuffers->ModuleGuid));
+    UT_LOG_INFO (
+      "%a: Searching by %a GUID ({%g}) returned ContextBuffers at 0x%x\n",
+      __FUNCTION__,
+      ((TestContext->GuidSearchType == ByModuleGuid) ? "module" : "handler"),
+      TestContext->Guid,
+      (UINTN) ContextBuffers
+      );
+  }
+
+  return UNIT_TEST_PASSED;
+}
+
+/**
+  Verifies that the expected PRM_CONTEXT_BUFFER instance is found for the given HandlerGuid
+  in the provided PRM_MODULE_CONTEXT_BUFFERS structure.
+
+  @param[in]  Context       [Optional] An optional context parameter.
+                            A pointer to a PRM_CONTEXT_BUFFERS_TEST_CONTEXT structure with
+                            context information for this unit test.
+
+  @retval  UNIT_TEST_PASSED                      Unit test case prerequisites are met.
+  @retval  UNIT_TEST_ERROR_PREREQUISITE_NOT_MET  Test case should be skipped..
+
+**/
+UNIT_TEST_STATUS
+EFIAPI
+VerifyFindContextBufferInModuleBuffers (
+  IN UNIT_TEST_CONTEXT                Context
+  )
+{
+  EFI_STATUS                          Status;
+  PRM_CONTEXT_BUFFER                  *FoundContextBuffer;
+  PRM_MODULE_CONTEXT_BUFFERS          *ContextBuffers;
+  PRM_CONTEXT_BUFFER_TEST_CONTEXT     *TestContext;
+
+  ContextBuffers = NULL;
+  FoundContextBuffer = NULL;
+  TestContext = (PRM_CONTEXT_BUFFER_TEST_CONTEXT *) Context;
+
+  Status = FindContextBufferInModuleBuffers (TestContext->HandlerGuid, TestContext->ContextBuffers, &FoundContextBuffer);
+  UT_ASSERT_STATUS_EQUAL (Status, TestContext->ExpectedStatus);
+
+  if (!EFI_ERROR (TestContext->ExpectedStatus)) {
+    UT_ASSERT_NOT_NULL (FoundContextBuffer);
+    UT_ASSERT_TRUE (FoundContextBuffer == TestContext->ExpectedContextBuffer);
+  }
+
+  return UNIT_TEST_PASSED;
+}
+
+/**
+  Verifies that the expected PRM_CONTEXT_BUFFER instance is found for the given HandlerGuid.
+
+  This function checks both the case when a PRM_MODULE_CONTEXT_BUFFERS structure pointer is provided and
+  not provided.
+
+  NOTES:
+  - In the future, this function should mock the internal calls to other library functions but the direct
+    calls are left in place for now.
+  - The PrmModuleContextBuffers being NULL is not actually tested at the moment. In the future, that case
+    should also be added.
+
+  @param[in]  Context       [Optional] An optional context parameter.
+                            A pointer to a PRM_CONTEXT_BUFFERS_TEST_CONTEXT structure with
+                            context information for this unit test.
+
+  @retval  UNIT_TEST_PASSED                      Unit test case prerequisites are met.
+  @retval  UNIT_TEST_ERROR_PREREQUISITE_NOT_MET  Test case should be skipped..
+
+**/
+UNIT_TEST_STATUS
+EFIAPI
+VerifyGetContextBuffer (
+  IN UNIT_TEST_CONTEXT                Context
+  )
+{
+  EFI_STATUS                          Status;
+  PRM_CONTEXT_BUFFER                  *FoundContextBuffer;
+  PRM_MODULE_CONTEXT_BUFFERS          *ContextBuffers;
+  PRM_CONTEXT_BUFFER_TEST_CONTEXT     *TestContext;
+
+  ContextBuffers = NULL;
+  FoundContextBuffer = NULL;
+  TestContext = (PRM_CONTEXT_BUFFER_TEST_CONTEXT *) Context;
+
+  Status = GetContextBuffer (TestContext->HandlerGuid, TestContext->ContextBuffers, &FoundContextBuffer);
+  UT_ASSERT_STATUS_EQUAL (Status, TestContext->ExpectedStatus);
+
+  if (!EFI_ERROR (TestContext->ExpectedStatus)) {
+    UT_ASSERT_NOT_NULL (FoundContextBuffer);
+    UT_ASSERT_TRUE (FoundContextBuffer == TestContext->ExpectedContextBuffer);
+  }
+
+  return UNIT_TEST_PASSED;
+}
+
+///=== TEST ENGINE ================================================================================
+
+/**
+  Entry point for the PRM Context Buffer Library unit tests.
+
+  @param[in] ImageHandle  The firmware allocated handle for the EFI image.
+  @param[in] SystemTable  A pointer to the EFI System Table.
+
+  @retval EFI_SUCCESS     The entry point executed successfully.
+  @retval other           Some error occurred when executing this entry point.
+
+**/
+int main ()
+{
+  EFI_STATUS                  Status;
+  UINTN                       Index;
+  UNIT_TEST_FRAMEWORK_HANDLE  Framework;
+  UNIT_TEST_SUITE_HANDLE      BasicTests;
+  UNIT_TEST_SUITE_HANDLE      FunctionalCorrectnessTests;
+  CHAR8                       TestCaseClassNameString[256];
+  CHAR8                       TestCaseDescriptionString[256];
+
+  Framework = NULL;
+
+  DEBUG ((DEBUG_INFO, "%a v%a\n", UNIT_TEST_NAME, UNIT_TEST_VERSION));
+
+  //
+  // Start setting up the test framework for running the tests.
+  //
+  Status = InitUnitTestFramework (&Framework, UNIT_TEST_NAME, gEfiCallerBaseName, UNIT_TEST_VERSION);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Failed in InitUnitTestFramework. Status = %r\n", Status));
+    goto EXIT;
+  }
+
+  //
+  // Add all test suites and tests.
+  //
+  Status = CreateUnitTestSuite (&BasicTests, Framework, "Basic Context Buffer Tests", "PrmContextBufferLib.Basic", NULL, NULL);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Failed in CreateUnitTestSuite for PrmContextBufferLib.Basic\n"));
+    Status = EFI_OUT_OF_RESOURCES;
+    goto EXIT;
+  }
+
+  AddTestCase (
+    BasicTests,
+    "",
+    "PrmContextBufferLib.Basic.NullPointerGracefulFailure",
+    NullPointerArgumentsShouldFailGracefully,
+    NULL,
+    NULL,
+    NULL
+    );
+
+  Status = CreateUnitTestSuite (&FunctionalCorrectnessTests, Framework, "Functional Correctness Tests", "PrmContextBufferLib.Functional", NULL, NULL);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Failed in CreateUnitTestSuite for PrmContextBufferLib.Functional\n"));
+    Status = EFI_OUT_OF_RESOURCES;
+    goto EXIT;
+  }
+
+  //
+  // Add Functional Correctness unit tests
+  //
+  for (Index = 0; Index < ARRAY_SIZE (mContextBuffersArray); Index++) {
+    ZeroMem (&TestCaseClassNameString[0], ARRAY_SIZE (TestCaseClassNameString));
+    ZeroMem (&TestCaseDescriptionString[0], ARRAY_SIZE (TestCaseDescriptionString));
+
+    AsciiSPrint (
+      &TestCaseClassNameString[0],
+      ARRAY_SIZE (TestCaseClassNameString),
+      "PrmContextBufferLib.Functional.VerifyGetModuleContextBuffers%d",
+      Index + 1
+      );
+    AsciiSPrint (
+      &TestCaseDescriptionString[0],
+      ARRAY_SIZE (TestCaseDescriptionString),
+      "Verify Get PRM Module Context Buffers Structure by %a GUID %d\n",
+      ((mContextBuffersArray[Index]->GuidSearchType == ByModuleGuid) ? "module" : "handler"),
+      Index + 1
+      );
+
+    AddTestCase (
+      FunctionalCorrectnessTests,
+      &TestCaseDescriptionString[0],
+      &TestCaseClassNameString[0],
+      VerifyGetModuleContextBuffers,
+      InitializeFunctionalCorrectness,
+      DeInitializeFunctionalCorrectness,
+      mContextBuffersArray[Index]
+      );
+  }
+
+  for (Index = 0; Index < ARRAY_SIZE (mContextBufferArray); Index++) {
+    ZeroMem (&TestCaseClassNameString[0], ARRAY_SIZE (TestCaseClassNameString));
+    ZeroMem (&TestCaseDescriptionString[0], ARRAY_SIZE (TestCaseDescriptionString));
+
+    AsciiSPrint (
+      &TestCaseClassNameString[0],
+      ARRAY_SIZE (TestCaseClassNameString),
+      "PrmContextBufferLib.Functional.VerifyFindContextBufferInModuleBuffers%d",
+      Index + 1
+      );
+    AsciiSPrint (
+      &TestCaseDescriptionString[0],
+      ARRAY_SIZE (TestCaseDescriptionString),
+      "Verify Find PRM Context Buffer by Handler GUID %d\n",
+      Index + 1
+      );
+
+    AddTestCase (
+      FunctionalCorrectnessTests,
+      &TestCaseDescriptionString[0],
+      &TestCaseClassNameString[0],
+      VerifyFindContextBufferInModuleBuffers,
+      NULL,
+      NULL,
+      mContextBufferArray[Index]
+      );
+  }
+
+  for (Index = 0; Index < ARRAY_SIZE (mContextBufferArray); Index++) {
+    ZeroMem (&TestCaseClassNameString[0], ARRAY_SIZE (TestCaseClassNameString));
+    ZeroMem (&TestCaseDescriptionString[0], ARRAY_SIZE (TestCaseDescriptionString));
+
+    AsciiSPrint (
+      &TestCaseClassNameString[0],
+      ARRAY_SIZE (TestCaseClassNameString),
+      "PrmContextBufferLib.Functional.VerifyGetContextBuffer%d",
+      Index + 1
+      );
+    AsciiSPrint (
+      &TestCaseDescriptionString[0],
+      ARRAY_SIZE (TestCaseDescriptionString),
+      "Verify Get PRM Context Buffer by Handler GUID %d\n",
+      Index + 1
+      );
+
+    AddTestCase (
+      FunctionalCorrectnessTests,
+      &TestCaseDescriptionString[0],
+      &TestCaseClassNameString[0],
+      VerifyGetContextBuffer,
+      NULL,
+      NULL,
+      mContextBufferArray[Index]
+      );
+  }
+
+  //
+  // Execute the tests.
+  //
+  Status = RunAllTestSuites (Framework);
+
+EXIT:
+  if (Framework)
+  {
+    FreeUnitTestFramework (Framework);
+  }
+
+  return Status;
+}
diff --git a/PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTestHost.inf b/PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTestHost.inf
new file mode 100644
index 000000000000..7cf6a16867d4
--- /dev/null
+++ b/PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTestHost.inf
@@ -0,0 +1,46 @@
+## @file
+#  PRM Context Buffer Library Host-Based Unit Tests
+#
+#  Copyright (c) Microsoft Corporation
+#
+#  SPDX-License-Identifier: BSD-2-Clause-Patent
+#
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010006
+  BASE_NAME                      = PrmContextBufferLibUnitTestHost
+  FILE_GUID                      = F1FB5F32-BDB5-4391-BD6D-979E90EE2DC3
+  MODULE_TYPE                    = HOST_APPLICATION
+  VERSION_STRING                 = 1.0
+
+#
+# The following information is for reference only and not required by the build tools.
+#
+#  VALID_ARCHITECTURES           = IA32 X64 ARM AARCH64
+#
+
+[Sources]
+  DxePrmContextBufferLibUnitTest.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  MdeModulePkg/MdeModulePkg.dec
+  UnitTestFrameworkPkg/UnitTestFrameworkPkg.dec
+  PrmPkg/PrmPkg.dec
+
+[Guids]
+  gZeroGuid
+
+[Protocols]
+  gPrmConfigProtocolGuid
+
+[LibraryClasses]
+  BaseLib
+  BaseMemoryLib
+  DebugLib
+  MemoryAllocationLib
+  PrintLib
+  PrmContextBufferLib
+  UefiBootServicesTableLib
+  UnitTestLib
diff --git a/PrmPkg/Test/PrmPkgHostTest.dsc b/PrmPkg/Test/PrmPkgHostTest.dsc
index 2d718f8888b8..1f44037d0102 100644
--- a/PrmPkg/Test/PrmPkgHostTest.dsc
+++ b/PrmPkg/Test/PrmPkgHostTest.dsc
@@ -19,6 +19,7 @@ [Defines]
 !include UnitTestFrameworkPkg/UnitTestFrameworkPkgHost.dsc.inc
 
 [LibraryClasses]
+  PrmContextBufferLib|PrmPkg/Library/DxePrmContextBufferLib/DxePrmContextBufferLib.inf
   UefiBootServicesTableLib|PrmPkg/Test/UnitTest/Library/UefiBootServicesTableLibUnitTest/UefiBootServicesTableLibUnitTest.inf
 
 [Components]
@@ -26,3 +27,8 @@ [Components]
   # Unit test helper libraries
   #
   PrmPkg/Test/UnitTest/Library/UefiBootServicesTableLibUnitTest/UefiBootServicesTableLibUnitTest.inf
+
+  #
+  # Unit test host applications
+  #
+  PrmPkg/Library/DxePrmContextBufferLib/UnitTest/DxePrmContextBufferLibUnitTestHost.inf
-- 
2.28.0.windows.1