9 files changed, 114 insertions, 18 deletions

diff --git a/src/test/examples/TestEngine/Example.TestEngine.vcxproj b/src/test/examples/TestEngine/Example.TestEngine.vcxproj
index 554d54f6..56a536b4 100644
--- a/src/test/examples/TestEngine/Example.TestEngine.vcxproj
+++ b/src/test/examples/TestEngine/Example.TestEngine.vcxproj
@@ -46,12 +46,15 @@
     <ClCompile Include="precomp.cpp">
       <PrecompiledHeader>Create</PrecompiledHeader>
     </ClCompile>
+    <ClCompile Include="ReloadEngine.cpp" />
     <ClCompile Include="ShutdownEngine.cpp" />
     <ClCompile Include="ExampleTestEngine.cpp" />
     <ClCompile Include="TestEngine.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="precomp.h" />
+    <ClInclude Include="ReloadEngine.h" />
+    <ClInclude Include="ShutdownEngine.h" />
     <ClInclude Include="TestEngine.h" />
   </ItemGroup>
   <ItemGroup>
diff --git a/src/test/examples/TestEngine/ExampleTestEngine.cpp b/src/test/examples/TestEngine/ExampleTestEngine.cpp
index 848b385c..a378c9a3 100644
--- a/src/test/examples/TestEngine/ExampleTestEngine.cpp
+++ b/src/test/examples/TestEngine/ExampleTestEngine.cpp
@@ -5,16 +5,30 @@
 int __cdecl wmain(int argc, LPWSTR argv[])
 {
     HRESULT hr = E_INVALIDARG;
+    BOOL fShowUsage = FALSE;
 
     ConsoleInitialize();
 
-    if (argc != 3)
+    if (argc != 4)
     {
-        ConsoleWriteError(hr, CONSOLE_COLOR_RED, "Usage: Example.TestEngine.exe Bundle.exe BA.dll");
+        fShowUsage = TRUE;
+    }
+    else if (CSTR_EQUAL == ::CompareString(LOCALE_INVARIANT, NORM_IGNORECASE, argv[1], -1, L"reload", -1))
+    {
+        hr = RunReloadEngine(argv[2], argv[3]);
+    }
+    else if (CSTR_EQUAL == ::CompareString(LOCALE_INVARIANT, NORM_IGNORECASE, argv[1], -1, L"shutdown", -1))
+    {
+        hr = RunShutdownEngine(argv[2], argv[3]);
     }
     else
     {
-        hr = RunShutdownEngine(argv[1], argv[2]);
+        fShowUsage = TRUE;
+    }
+
+    if (fShowUsage)
+    {
+        ConsoleWriteError(hr, CONSOLE_COLOR_RED, "Usage: {reload|shutdown} Example.TestEngine.exe Bundle.exe BA.dll");
     }
 
     ConsoleUninitialize();
diff --git a/src/test/examples/TestEngine/ReloadEngine.cpp b/src/test/examples/TestEngine/ReloadEngine.cpp
new file mode 100644
index 00000000..83541672
--- /dev/null
+++ b/src/test/examples/TestEngine/ReloadEngine.cpp
@@ -0,0 +1,43 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+#include "precomp.h"
+
+HRESULT RunReloadEngine(
+    __in LPCWSTR wzBundleFilePath,
+    __in LPCWSTR wzBAFilePath
+    )
+{
+    HRESULT hr = S_OK;
+    TestEngine* pTestEngine = NULL;
+
+    pTestEngine = new TestEngine();
+    ConsoleExitOnNull(pTestEngine, hr, E_OUTOFMEMORY, CONSOLE_COLOR_RED, "Failed to create new test engine.");
+
+    hr = pTestEngine->Initialize(wzBundleFilePath);
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "Failed to initialize engine.");
+
+    hr = pTestEngine->LoadBA(wzBAFilePath);
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "Failed to load BA.");
+
+    hr = pTestEngine->SendStartupEvent();
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "BA returned failure for OnStartup.");
+
+    hr = pTestEngine->SendShutdownEvent(BOOTSTRAPPER_SHUTDOWN_ACTION_RELOAD_BOOTSTRAPPER);
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "BA returned failure for OnShutdown.");
+
+    pTestEngine->UnloadBA();
+
+    hr = pTestEngine->LoadBA(wzBAFilePath);
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "Failed to load BA.");
+
+    hr = pTestEngine->SendStartupEvent();
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "BA returned failure for OnStartup.");
+
+    hr = pTestEngine->SendShutdownEvent(BOOTSTRAPPER_SHUTDOWN_ACTION_RESTART);
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "BA returned failure for OnShutdown.");
+
+    pTestEngine->UnloadBA();
+
+LExit:
+    return hr;
+}
diff --git a/src/test/examples/TestEngine/ReloadEngine.h b/src/test/examples/TestEngine/ReloadEngine.h
new file mode 100644
index 00000000..0e8456af
--- /dev/null
+++ b/src/test/examples/TestEngine/ReloadEngine.h
@@ -0,0 +1,8 @@
+#pragma once
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+
+HRESULT RunReloadEngine(
+    __in LPCWSTR wzBundleFilePath,
+    __in LPCWSTR wzBAFilePath
+    );
diff --git a/src/test/examples/TestEngine/ShutdownEngine.cpp b/src/test/examples/TestEngine/ShutdownEngine.cpp
index 912d36ba..0dfbb429 100644
--- a/src/test/examples/TestEngine/ShutdownEngine.cpp
+++ b/src/test/examples/TestEngine/ShutdownEngine.cpp
@@ -13,7 +13,10 @@ HRESULT RunShutdownEngine(
     pTestEngine = new TestEngine();
     ConsoleExitOnNull(pTestEngine, hr, E_OUTOFMEMORY, CONSOLE_COLOR_RED, "Failed to create new test engine.");
 
-    hr = pTestEngine->LoadBA(wzBundleFilePath, wzBAFilePath);
+    hr = pTestEngine->Initialize(wzBundleFilePath);
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "Failed to initialize engine.");
+
+    hr = pTestEngine->LoadBA(wzBAFilePath);
     ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "Failed to load BA.");
 
     hr = pTestEngine->SendStartupEvent();
diff --git a/src/test/examples/TestEngine/ShutdownEngine.h b/src/test/examples/TestEngine/ShutdownEngine.h
new file mode 100644
index 00000000..0cfa147a
--- /dev/null
+++ b/src/test/examples/TestEngine/ShutdownEngine.h
@@ -0,0 +1,8 @@
+#pragma once
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+
+HRESULT RunShutdownEngine(
+    __in LPCWSTR wzBundleFilePath,
+    __in LPCWSTR wzBAFilePath
+    );
diff --git a/src/test/examples/TestEngine/TestEngine.cpp b/src/test/examples/TestEngine/TestEngine.cpp
index 9d8f8638..203df115 100644
--- a/src/test/examples/TestEngine/TestEngine.cpp
+++ b/src/test/examples/TestEngine/TestEngine.cpp
@@ -2,8 +2,22 @@
 
 #include "precomp.h"
 
+HRESULT TestEngine::Initialize(
+    __in LPCWSTR wzBundleFilePath
+    )
+{
+    HRESULT hr = S_OK;
+
+    LogInitialize(::GetModuleHandleW(NULL));
+
+    hr = LogOpen(NULL, PathFile(wzBundleFilePath), NULL, L"txt", FALSE, FALSE, NULL);
+    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "Failed to open log.");
+
+LExit:
+    return hr;
+}
+
 HRESULT TestEngine::LoadBA(
-    __in LPCWSTR wzBundleFilePath,
     __in LPCWSTR wzBAFilePath
     )
 {
@@ -12,16 +26,11 @@ HRESULT TestEngine::LoadBA(
     BOOTSTRAPPER_CREATE_ARGS args = { };
     PFN_BOOTSTRAPPER_APPLICATION_CREATE pfnCreate = NULL;
 
-    if (m_pCreateResults)
+    if (m_pCreateResults || m_hBAModule)
     {
         ExitFunction1(hr = E_INVALIDSTATE);
     }
 
-    LogInitialize(::GetModuleHandleW(NULL));
-
-    hr = LogOpen(NULL, PathFile(wzBundleFilePath), NULL, L"txt", FALSE, FALSE, NULL);
-    ConsoleExitOnFailure(hr, CONSOLE_COLOR_RED, "Failed to open log.");
-
     m_pCreateResults = static_cast<BOOTSTRAPPER_CREATE_RESULTS*>(MemAlloc(sizeof(BOOTSTRAPPER_CREATE_RESULTS), TRUE));
 
     command.cbSize = sizeof(BOOTSTRAPPER_COMMAND);
@@ -51,6 +60,7 @@ HRESULT TestEngine::Log(
     __in LPCWSTR wzMessage
     )
 {
+    LogStringLine(REPORT_STANDARD, "%ls", wzMessage);
     return ConsoleWriteLine(CONSOLE_COLOR_NORMAL, "%ls", wzMessage);
 }
 
@@ -83,12 +93,20 @@ void TestEngine::UnloadBA()
 {
     PFN_BOOTSTRAPPER_APPLICATION_DESTROY pfnDestroy = NULL;
 
+    ReleaseNullMem(m_pCreateResults);
+
     pfnDestroy = (PFN_BOOTSTRAPPER_APPLICATION_DESTROY)::GetProcAddress(m_hBAModule, "BootstrapperApplicationDestroy");
 
     if (pfnDestroy)
     {
         pfnDestroy();
     }
+
+    if (m_hBAModule)
+    {
+        ::FreeLibrary(m_hBAModule);
+        m_hBAModule = NULL;
+    }
 }
 
 HRESULT TestEngine::BAEngineLog(
diff --git a/src/test/examples/TestEngine/TestEngine.h b/src/test/examples/TestEngine/TestEngine.h
index e5db9480..cf1c8aac 100644
--- a/src/test/examples/TestEngine/TestEngine.h
+++ b/src/test/examples/TestEngine/TestEngine.h
@@ -1,13 +1,15 @@
 #pragma once
 // Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
 
-#include "precomp.h"
 
 class TestEngine
 {
 public:
+    HRESULT Initialize(
+        __in LPCWSTR wzBundleFilePath
+        );
+
     HRESULT LoadBA(
-        __in LPCWSTR wzBundleFilePath,
         __in LPCWSTR wzBAFilePath
         );
 
diff --git a/src/test/examples/TestEngine/precomp.h b/src/test/examples/TestEngine/precomp.h
index d0068747..0d2afb06 100644
--- a/src/test/examples/TestEngine/precomp.h
+++ b/src/test/examples/TestEngine/precomp.h
@@ -14,8 +14,5 @@
 #include "BootstrapperApplication.h"
 
 #include "TestEngine.h"
-
-HRESULT RunShutdownEngine(
-    __in LPCWSTR wzBundleFilePath,
-    __in LPCWSTR wzBAFilePath
-    );
+#include "ReloadEngine.h"
+#include "ShutdownEngine.h"