13 files changed, 217 insertions, 18 deletions

diff --git a/src/test/WixToolsetTest.ManagedHost/MbaHostFixture.cs b/src/test/WixToolsetTest.ManagedHost/MbaHostFixture.cs
index 8ef0320e..8b0a3691 100644
--- a/src/test/WixToolsetTest.ManagedHost/MbaHostFixture.cs
+++ b/src/test/WixToolsetTest.ManagedHost/MbaHostFixture.cs
@@ -78,5 +78,83 @@ namespace WixToolsetTest.ManagedHost
                 Assert.Equal("Shutdown,ReloadBootstrapper,0", logMessages[3]);
             }
         }
+
+        [Fact]
+        public void CanReloadFullFramework2MBA()
+        {
+            using (var fs = new DisposableFileSystem())
+            {
+                var baseFolder = fs.GetFolder();
+                var binFolder = Path.Combine(baseFolder, "bin");
+                var bundleFile = Path.Combine(binFolder, "FullFramework2MBA.exe");
+                var baSourceFolder = TestData.Get(@"..\examples");
+                var bundleSourceFolder = TestData.Get(@"TestData\FullFramework2MBA");
+                var intermediateFolder = Path.Combine(baseFolder, "obj");
+
+                var compileResult = WixRunner.Execute(new[]
+                {
+                    "build",
+                    Path.Combine(bundleSourceFolder, "Bundle.wxs"),
+                    "-ext", TestData.Get(@"WixToolset.Bal.wixext.dll"),
+                    "-ext", TestData.Get(@"WixToolset.NetFx.wixext.dll"),
+                    "-intermediateFolder", intermediateFolder,
+                    "-bindpath", baSourceFolder,
+                    "-burnStub", TestEngine.BurnStubFile,
+                    "-o", bundleFile,
+                });
+                compileResult.AssertSuccess();
+                var testEngine = new TestEngine();
+
+                var result = testEngine.RunReloadEngine(bundleFile, baseFolder);
+                var logMessages = result.Output;
+                Assert.Equal("Loading managed bootstrapper application.", logMessages[0]);
+                Assert.Equal("Creating BA thread to run asynchronously.", logMessages[1]);
+                Assert.Equal("FullFramework2BA", logMessages[2]);
+                Assert.Equal("Shutdown,ReloadBootstrapper,0", logMessages[3]);
+                Assert.Equal("Loading managed bootstrapper application.", logMessages[4]);
+                Assert.Equal("Creating BA thread to run asynchronously.", logMessages[5]);
+                Assert.Equal("FullFramework2BA", logMessages[6]);
+                Assert.Equal("Shutdown,Restart,0", logMessages[7]);
+            }
+        }
+
+        [Fact]
+        public void CanReloadFullFramework4MBA()
+        {
+            using (var fs = new DisposableFileSystem())
+            {
+                var baseFolder = fs.GetFolder();
+                var binFolder = Path.Combine(baseFolder, "bin");
+                var bundleFile = Path.Combine(binFolder, "FullFramework4MBA.exe");
+                var baSourceFolder = TestData.Get(@"..\examples");
+                var bundleSourceFolder = TestData.Get(@"TestData\FullFramework4MBA");
+                var intermediateFolder = Path.Combine(baseFolder, "obj");
+
+                var compileResult = WixRunner.Execute(new[]
+                {
+                    "build",
+                    Path.Combine(bundleSourceFolder, "Bundle.wxs"),
+                    "-ext", TestData.Get(@"WixToolset.Bal.wixext.dll"),
+                    "-ext", TestData.Get(@"WixToolset.NetFx.wixext.dll"),
+                    "-intermediateFolder", intermediateFolder,
+                    "-bindpath", baSourceFolder,
+                    "-burnStub", TestEngine.BurnStubFile,
+                    "-o", bundleFile,
+                });
+                compileResult.AssertSuccess();
+                var testEngine = new TestEngine();
+
+                var result = testEngine.RunReloadEngine(bundleFile, baseFolder);
+                var logMessages = result.Output;
+                Assert.Equal("Loading managed bootstrapper application.", logMessages[0]);
+                Assert.Equal("Creating BA thread to run asynchronously.", logMessages[1]);
+                Assert.Equal("FullFramework4BA", logMessages[2]);
+                Assert.Equal("Shutdown,ReloadBootstrapper,0", logMessages[3]);
+                Assert.Equal("Loading managed bootstrapper application.", logMessages[4]);
+                Assert.Equal("Creating BA thread to run asynchronously.", logMessages[5]);
+                Assert.Equal("FullFramework4BA", logMessages[6]);
+                Assert.Equal("Shutdown,Restart,0", logMessages[7]);
+            }
+        }
     }
 }
diff --git a/src/test/WixToolsetTest.ManagedHost/TestEngine.cs b/src/test/WixToolsetTest.ManagedHost/TestEngine.cs
index 751ed59c..cda32895 100644
--- a/src/test/WixToolsetTest.ManagedHost/TestEngine.cs
+++ b/src/test/WixToolsetTest.ManagedHost/TestEngine.cs
@@ -13,14 +13,25 @@ namespace WixToolsetTest.ManagedHost
         private static readonly string TestEngineFile = TestData.Get(@"..\Win32\examples\Example.TestEngine\Example.TestEngine.exe");
         public static readonly string BurnStubFile = TestData.Get(@"runtimes\win-x86\native\burn.x86.exe");
 
+        public TestEngineResult RunReloadEngine(string bundleFilePath, string tempFolderPath)
+        {
+            return this.RunTestEngine("reload", bundleFilePath, tempFolderPath);
+        }
+
         public TestEngineResult RunShutdownEngine(string bundleFilePath, string tempFolderPath)
         {
+            return this.RunTestEngine("shutdown", bundleFilePath, tempFolderPath);
+        }
+
+        private TestEngineResult RunTestEngine(string engineMode, string bundleFilePath, string tempFolderPath)
+        {
             var baFolderPath = Path.Combine(tempFolderPath, "ba");
             var extractFolderPath = Path.Combine(tempFolderPath, "extract");
             var extractResult = BundleExtractor.ExtractBAContainer(null, bundleFilePath, baFolderPath, extractFolderPath);
             extractResult.AssertSuccess();
 
             var args = new string[] {
+                engineMode,
                 '"' + bundleFilePath + '"',
                 '"' + extractResult.GetBAFilePath(baFolderPath) + '"',
             };
diff --git a/src/test/examples/FullFramework2MBA/FullFramework2BAFactory.cs b/src/test/examples/FullFramework2MBA/FullFramework2BAFactory.cs
index d3cafc70..40cff30f 100644
--- a/src/test/examples/FullFramework2MBA/FullFramework2BAFactory.cs
+++ b/src/test/examples/FullFramework2MBA/FullFramework2BAFactory.cs
@@ -6,8 +6,15 @@ namespace Example.FullFramework2MBA
 
     public class FullFramework2BAFactory : BaseBootstrapperApplicationFactory
     {
+        private static int loadCount = 0;
+
         protected override IBootstrapperApplication Create(IEngine engine, IBootstrapperCommand bootstrapperCommand)
         {
+            if (loadCount > 0)
+            {
+                engine.Log(LogLevel.Standard, $"Reloaded {loadCount} time(s)");
+            }
+            ++loadCount;
             return new FullFramework2BA(engine);
         }
     }
diff --git a/src/test/examples/FullFramework4MBA/FullFramework4BAFactory.cs b/src/test/examples/FullFramework4MBA/FullFramework4BAFactory.cs
index b7c8750d..6a571a54 100644
--- a/src/test/examples/FullFramework4MBA/FullFramework4BAFactory.cs
+++ b/src/test/examples/FullFramework4MBA/FullFramework4BAFactory.cs
@@ -7,8 +7,15 @@ namespace Example.FullFramework4MBA
 
     public class FullFramework4BAFactory : BaseBootstrapperApplicationFactory
     {
+        private static int loadCount = 0;
+
         protected override IBootstrapperApplication Create(IEngine engine, IBootstrapperCommand bootstrapperCommand)
         {
+            if (loadCount > 0)
+            {
+                engine.Log(LogLevel.Standard, $"Reloaded {loadCount} time(s)");
+            }
+            ++loadCount;
             return new FullFramework4BA(engine);
         }
     }
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"