1 files changed, 15 insertions, 10 deletions

diff --git a/UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c b/UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c
index ba4b18568d..4a6bbd752c 100644
--- a/UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c
+++ b/UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c
@@ -781,11 +781,9 @@ SerializeState (
   at least the current execution count) which will be saved by the framework and

   passed to the test case upon resume.

 

-  Generally called from within a test case prior to quitting or rebooting.

+  This should be called while the current test framework is valid and active. It is

+  generally called from within a test case prior to quitting or rebooting.

 

-  @param[in]  FrameworkHandle    A handle to the current running framework that

-                                 dispatched the test.  Necessary for recording

-                                 certain test events with the framework.

   @param[in]  ContextToSave      A buffer of test case-specific data to be saved

                                  along with framework state.  Will be passed as

                                  "Context" to the test case upon resume.  This

@@ -793,7 +791,7 @@ SerializeState (
   @param[in]  ContextToSaveSize  Size of the ContextToSave buffer.

 

   @retval  EFI_SUCCESS            The framework state and context were saved.

-  @retval  EFI_INVALID_PARAMETER  FrameworkHandle is NULL.

+  @retval  EFI_NOT_FOUND          An active framework handle was not found.

   @retval  EFI_INVALID_PARAMETER  ContextToSave is not NULL and

                                   ContextToSaveSize is 0.

   @retval  EFI_INVALID_PARAMETER  ContextToSave is >= 4GB.

@@ -806,21 +804,28 @@ SerializeState (
 EFI_STATUS

 EFIAPI

 SaveFrameworkState (

-  IN UNIT_TEST_FRAMEWORK_HANDLE  FrameworkHandle,

   IN UNIT_TEST_CONTEXT           ContextToSave     OPTIONAL,

   IN UINTN                       ContextToSaveSize

   )

 {

-  EFI_STATUS             Status;

-  UNIT_TEST_SAVE_HEADER  *Header;

+  EFI_STATUS                  Status;

+  UNIT_TEST_FRAMEWORK_HANDLE  FrameworkHandle;

+  UNIT_TEST_SAVE_HEADER       *Header;

 

   Header = NULL;

+  FrameworkHandle = GetActiveFrameworkHandle ();

+

+  //

+  // Return a unique error code if the framework is not set.

+  //

+  if (FrameworkHandle == NULL) {

+    return EFI_NOT_FOUND;

+  }

 

   //

   // First, let's not make assumptions about the parameters.

   //

-  if (FrameworkHandle == NULL ||

-      (ContextToSave != NULL && ContextToSaveSize == 0) ||

+  if ((ContextToSave != NULL && ContextToSaveSize == 0) ||

       ContextToSaveSize > MAX_UINT32) {

     return EFI_INVALID_PARAMETER;

   }