[PATCH v1 2/2] UnitTestFrameworkPkg/UnitTestLib: Update SaveFrameworkState() signature

["Michael Kubacki" <michael.kubacki@outlook.com>]

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

REF:https://bugzilla.tianocore.org/show_bug.cgi?id=2612

Removes the FrameworkHandle parameter from SaveFrameworkState() in the
UnitTestLib library instance and updates callers of the function in the
library to use the new function signature.

Cc: Bret Barkelew <Bret.Barkelew@microsoft.com>
Cc: Sean Brogan <sean.brogan@microsoft.com>
Cc: Michael D Kinney <michael.d.kinney@intel.com>
Signed-off-by: Michael Kubacki <michael.kubacki@microsoft.com>
---
 UnitTestFrameworkPkg/Library/UnitTestLib/RunTests.c    |  2 +-
 UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c | 25 ++++++++++++--------
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/UnitTestFrameworkPkg/Library/UnitTestLib/RunTests.c b/UnitTestFrameworkPkg/Library/UnitTestLib/RunTests.c
index b053e0495912..793335fd0f05 100644
--- a/UnitTestFrameworkPkg/Library/UnitTestLib/RunTests.c
+++ b/UnitTestFrameworkPkg/Library/UnitTestLib/RunTests.c
@@ -162,7 +162,7 @@ RunAllTestSuites (
   //
   // Save current state so if test is started again it doesn't have to run.  It will just report
   //
-  SaveFrameworkState (FrameworkHandle, NULL, 0);
+  SaveFrameworkState (NULL, 0);
   OutputUnitTestFrameworkReport (FrameworkHandle);
 
   mFrameworkHandle = NULL;
diff --git a/UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c b/UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c
index b136992d99a0..e2405251f0f4 100644
--- a/UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c
+++ b/UnitTestFrameworkPkg/Library/UnitTestLib/UnitTestLib.c
@@ -783,11 +783,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
@@ -795,7 +793,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.
@@ -808,21 +806,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;
   }
-- 
2.16.3.windows.1