5 files changed, 20 insertions, 10 deletions

diff --git a/src/ext/Bal/test/examples/TestEngine/ReloadEngine.cpp b/src/ext/Bal/test/examples/TestEngine/ReloadEngine.cpp
index 46fd9afa..3f290b86 100644
--- a/src/ext/Bal/test/examples/TestEngine/ReloadEngine.cpp
+++ b/src/ext/Bal/test/examples/TestEngine/ReloadEngine.cpp
@@ -31,7 +31,7 @@ HRESULT RunReloadEngine(
     hr = pTestEngine->SendShutdownEvent(BOOTSTRAPPER_SHUTDOWN_ACTION_RELOAD_BOOTSTRAPPER);
     ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "BA returned failure for OnShutdown.");
 
-    pTestEngine->UnloadBA();
+    pTestEngine->UnloadBA(TRUE);
 
     hr = pTestEngine->LoadBA(wzBAFilePath);
     ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "Failed to load BA.");
@@ -48,7 +48,7 @@ HRESULT RunReloadEngine(
     hr = pTestEngine->SendShutdownEvent(BOOTSTRAPPER_SHUTDOWN_ACTION_RESTART);
     ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "BA returned failure for OnShutdown.");
 
-    pTestEngine->UnloadBA();
+    pTestEngine->UnloadBA(FALSE);
 
 LExit:
     return hr;
diff --git a/src/ext/Bal/test/examples/TestEngine/ShutdownEngine.cpp b/src/ext/Bal/test/examples/TestEngine/ShutdownEngine.cpp
index 3b876e4e..b5c8c462 100644
--- a/src/ext/Bal/test/examples/TestEngine/ShutdownEngine.cpp
+++ b/src/ext/Bal/test/examples/TestEngine/ShutdownEngine.cpp
@@ -31,7 +31,7 @@ HRESULT RunShutdownEngine(
     hr = pTestEngine->SendShutdownEvent(BOOTSTRAPPER_SHUTDOWN_ACTION_RELOAD_BOOTSTRAPPER);
     ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "BA returned failure for OnShutdown.");
 
-    pTestEngine->UnloadBA();
+    pTestEngine->UnloadBA(FALSE);
 
 LExit:
     return hr;
diff --git a/src/ext/Bal/test/examples/TestEngine/TestEngine.cpp b/src/ext/Bal/test/examples/TestEngine/TestEngine.cpp
index 4c7ec1c3..5c6ed398 100644
--- a/src/ext/Bal/test/examples/TestEngine/TestEngine.cpp
+++ b/src/ext/Bal/test/examples/TestEngine/TestEngine.cpp
@@ -147,10 +147,18 @@ HRESULT TestEngine::SimulateQuit(
     return BAEngineQuit(&args, &results);
 }
 
-void TestEngine::UnloadBA()
+void TestEngine::UnloadBA(
+    __in BOOL fReload
+    )
 {
     PFN_BOOTSTRAPPER_APPLICATION_DESTROY pfnDestroy = NULL;
-    BOOL fDisableUnloading = m_pCreateResults && m_pCreateResults->fDisableUnloading;
+    BOOTSTRAPPER_DESTROY_ARGS args = { };
+    BOOTSTRAPPER_DESTROY_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.fReload = fReload;
+
+    results.cbSize = sizeof(results);
 
     ReleaseNullMem(m_pCreateResults);
 
@@ -158,12 +166,12 @@ void TestEngine::UnloadBA()
 
     if (pfnDestroy)
     {
-        pfnDestroy();
+        pfnDestroy(&args, &results);
     }
 
     if (m_hBAModule)
     {
-        if (!fDisableUnloading)
+        if (!results.fDisableUnloading)
         {
             ::FreeLibrary(m_hBAModule);
         }
diff --git a/src/ext/Bal/test/examples/TestEngine/TestEngine.h b/src/ext/Bal/test/examples/TestEngine/TestEngine.h
index 44e813bd..248e979a 100644
--- a/src/ext/Bal/test/examples/TestEngine/TestEngine.h
+++ b/src/ext/Bal/test/examples/TestEngine/TestEngine.h
@@ -44,7 +44,9 @@ public:
         __in DWORD dwExitCode
         );
 
-    void UnloadBA();
+    void UnloadBA(
+        __in BOOL fReload
+        );
 
 private:
     HRESULT BAEngineLog(
@@ -77,4 +79,4 @@ private:
     HMODULE m_hBAModule;
     BOOTSTRAPPER_CREATE_RESULTS* m_pCreateResults;
     DWORD m_dwThreadId;
-};
\ No newline at end of file
+};
diff --git a/src/ext/Bal/test/examples/TestEngine/WaitForQuitEngine.cpp b/src/ext/Bal/test/examples/TestEngine/WaitForQuitEngine.cpp
index 2f80ba75..64b618f4 100644
--- a/src/ext/Bal/test/examples/TestEngine/WaitForQuitEngine.cpp
+++ b/src/ext/Bal/test/examples/TestEngine/WaitForQuitEngine.cpp
@@ -28,7 +28,7 @@ HRESULT RunWaitForQuitEngine(
     hr = pTestEngine->SendShutdownEvent(BOOTSTRAPPER_SHUTDOWN_ACTION_RELOAD_BOOTSTRAPPER);
     ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "BA returned failure for OnShutdown.");
 
-    pTestEngine->UnloadBA();
+    pTestEngine->UnloadBA(FALSE);
 
 LExit:
     return hr;