12 files changed, 193 insertions, 35 deletions

diff --git a/global.json b/global.json
index f227ec06..71e4088d 100644
--- a/global.json
+++ b/global.json
@@ -1,6 +1,6 @@
 {
   "msbuild-sdks": {
-    "WixToolset.Sdk": "4.0.0-build-0186"
+    "WixToolset.Sdk": "4.0.0-build-0187"
   },
   "sdk": {
     "allowPrerelease": false
diff --git a/src/TestData/PrereqBaTests/BundleA/BundleA.wxs b/src/TestData/PrereqBaTests/BundleA/BundleA.wxs
index 0c9e8d21..c903988b 100644
--- a/src/TestData/PrereqBaTests/BundleA/BundleA.wxs
+++ b/src/TestData/PrereqBaTests/BundleA/BundleA.wxs
@@ -16,7 +16,7 @@
   <Fragment>
     <PackageGroup Id="BundlePackages">
       <MsiPackage Id="PackageA" SourceFile="$(var.PackageA.TargetPath)" bal:PrereqPackage="yes" Permanent="yes" />
-      <MsiPackage Id="PackageF" SourceFile="$(var.PackageF.TargetPath)" />
+      <MsiPackage Id="PackageF" SourceFile="$(var.PackageF.TargetPath)" Cache="always" />
     </PackageGroup>
   </Fragment>
 </Wix>
diff --git a/src/TestData/RegistrationTests/BundleA/BundleA.wixproj b/src/TestData/RegistrationTests/BundleA/BundleA.wixproj
new file mode 100644
index 00000000..313ffec9
--- /dev/null
+++ b/src/TestData/RegistrationTests/BundleA/BundleA.wixproj
@@ -0,0 +1,18 @@
+<!-- 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. -->
+<Project Sdk="WixToolset.Sdk">
+  <PropertyGroup>
+    <OutputType>Bundle</OutputType>
+    <UpgradeCode>{17041020-8A61-4A3B-8FDB-4591CB900049}</UpgradeCode>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="..\..\Templates\Bundle.wxs" Link="Bundle.wxs" />
+  </ItemGroup>
+  <ItemGroup>
+    <ProjectReference Include="..\PackageA\PackageA.wixproj" />
+    <ProjectReference Include="..\..\TestBA\TestBAWixlib\testbawixlib.wixproj" />
+  </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.Bal.wixext" Version="4.0.82" />
+    <PackageReference Include="WixToolset.NetFx.wixext" Version="4.0.59" />
+  </ItemGroup>
+</Project>
\ No newline at end of file
diff --git a/src/TestData/RegistrationTests/BundleA/BundleA.wxs b/src/TestData/RegistrationTests/BundleA/BundleA.wxs
new file mode 100644
index 00000000..bd164a29
--- /dev/null
+++ b/src/TestData/RegistrationTests/BundleA/BundleA.wxs
@@ -0,0 +1,10 @@
+<!-- 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. -->
+
+
+<Wix xmlns="http://wixtoolset.org/schemas/v4/wxs">
+  <Fragment>
+    <PackageGroup Id="BundlePackages">
+      <MsiPackage Id="PackageA" SourceFile="$(var.PackageA.TargetPath)" />
+    </PackageGroup>
+  </Fragment>
+</Wix>
diff --git a/src/TestData/RegistrationTests/PackageA/PackageA.wixproj b/src/TestData/RegistrationTests/PackageA/PackageA.wixproj
new file mode 100644
index 00000000..7e5dbb8c
--- /dev/null
+++ b/src/TestData/RegistrationTests/PackageA/PackageA.wixproj
@@ -0,0 +1,9 @@
+<!-- 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. -->
+<Project Sdk="WixToolset.Sdk">
+  <PropertyGroup>
+    <UpgradeCode>{BF26D3E4-1D6B-480E-B312-3FECE6363E43}</UpgradeCode>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="..\..\Templates\Package.wxs" Link="Package.wxs" />
+  </ItemGroup>
+</Project>
\ No newline at end of file
diff --git a/src/Utilities/TestBA/TestBA.cs b/src/Utilities/TestBA/TestBA.cs
index c68c2b61..1348ce98 100644
--- a/src/Utilities/TestBA/TestBA.cs
+++ b/src/Utilities/TestBA/TestBA.cs
@@ -18,14 +18,17 @@ namespace WixToolset.Test.BA
     {
         private const string BurnBundleVersionVariable = "WixBundleVersion";
 
-        private ApplicationContext appContext;
         private Form dummyWindow;
+        private IntPtr windowHandle;
         private LaunchAction action;
+        private ManualResetEvent wait;
         private int result;
 
         private string updateBundlePath;
 
-        private int redetectCount;
+        private bool immediatelyQuit;
+        private bool quitAfterDetect;
+        private int redetectRemaining;
         private int sleepDuringCache;
         private int cancelCacheAtProgress;
         private int sleepDuringExecute;
@@ -45,6 +48,7 @@ namespace WixToolset.Test.BA
             : base(engine)
         {
             this.Command = bootstrapperCommand;
+            this.wait = new ManualResetEvent(false);
         }
 
         /// <summary>
@@ -60,8 +64,17 @@ namespace WixToolset.Test.BA
         /// <summary>
         /// UI Thread entry point for TestUX.
         /// </summary>
-        protected override void Run()
+        protected override void OnStartup(StartupEventArgs args)
         {
+            string immediatelyQuit = this.ReadPackageAction(null, "ImmediatelyQuit");
+            if (!String.IsNullOrEmpty(immediatelyQuit) && Boolean.TryParse(immediatelyQuit, out this.immediatelyQuit) && this.immediatelyQuit)
+            {
+                this.Engine.Quit(0);
+                return;
+            }
+
+            base.OnStartup(args);
+
             this.action = this.Command.Action;
             this.TestVariables();
 
@@ -102,20 +115,46 @@ namespace WixToolset.Test.BA
                 return;
             }
 
-            this.dummyWindow = new Form();
-            this.dummyWindow.CreateControl();
-            this.appContext = new ApplicationContext();
-
-            this.redetectCount = 0;
+            int redetectCount;
             string redetect = this.ReadPackageAction(null, "RedetectCount");
-            if (String.IsNullOrEmpty(redetect) || !Int32.TryParse(redetect, out this.redetectCount))
+            if (String.IsNullOrEmpty(redetect) || !Int32.TryParse(redetect, out redetectCount))
+            {
+                redetectCount = 0;
+            }
+
+            string quitAfterDetect = this.ReadPackageAction(null, "QuitAfterDetect");
+            if (String.IsNullOrEmpty(quitAfterDetect) || !Boolean.TryParse(quitAfterDetect, out this.quitAfterDetect))
             {
-                this.redetectCount = 0;
+                this.quitAfterDetect = false;
             }
 
-            this.Engine.Detect();
+            this.wait.WaitOne();
+
+            this.redetectRemaining = redetectCount;
+            for (int i = -1; i < redetectCount; i++)
+            {
+                this.Engine.Detect(this.windowHandle);
+            }
+        }
+
+        protected override void Run()
+        {
+            this.dummyWindow = new Form();
+            this.windowHandle = this.dummyWindow.Handle;
+
+            this.Log("Running TestBA application");
+            this.wait.Set();
+            Application.Run();
+        }
+
+        private void ShutdownUiThread()
+        {
+            if (this.dummyWindow != null)
+            {
+                this.dummyWindow.Invoke(new Action(Application.ExitThread));
+                this.dummyWindow.Dispose();
+            }
 
-            Application.Run(this.appContext);
             this.Engine.Quit(this.result & 0xFFFF); // return plain old Win32 error, not HRESULT.
         }
 
@@ -167,10 +206,13 @@ namespace WixToolset.Test.BA
             if (Hresult.Succeeded(this.result) &&
                 (this.UpdateAvailable || LaunchAction.UpdateReplaceEmbedded != this.action && LaunchAction.UpdateReplace != this.action))
             {
-                if (this.redetectCount > 0)
+                if (this.redetectRemaining > 0)
+                {
+                    this.Log("Completed detection phase: {0} re-runs remaining", this.redetectRemaining--);
+                }
+                else if (this.quitAfterDetect)
                 {
-                    this.Log("Completed detection phase: {0} re-runs remaining", this.redetectCount--);
-                    this.Engine.Detect();
+                    this.ShutdownUiThread();
                 }
                 else
                 {
@@ -179,7 +221,7 @@ namespace WixToolset.Test.BA
             }
             else
             {
-                this.appContext.ExitThread();
+                this.ShutdownUiThread();
             }
         }
 
@@ -218,11 +260,11 @@ namespace WixToolset.Test.BA
             this.result = args.Status;
             if (Hresult.Succeeded(this.result))
             {
-                this.Engine.Apply(this.dummyWindow.Handle);
+                this.Engine.Apply(this.windowHandle);
             }
             else
             {
-                this.appContext.ExitThread();
+                this.ShutdownUiThread();
             }
         }
 
@@ -402,7 +444,7 @@ namespace WixToolset.Test.BA
             this.Log("After elevation: WixBundleElevated = {0}", this.Engine.GetVariableNumeric("WixBundleElevated"));
 
             this.result = args.Status;
-            this.appContext.ExitThread();
+            this.ShutdownUiThread();
         }
 
         protected override void OnSystemShutdown(SystemShutdownEventArgs args)
@@ -411,7 +453,7 @@ namespace WixToolset.Test.BA
             this.Log("Disallowed system request to shut down the bootstrapper application.");
             args.Cancel = true;
 
-            this.appContext.ExitThread();
+            this.ShutdownUiThread();
         }
 
         private void TestVariables()
diff --git a/src/WixToolsetTest.BurnE2E/BundleVerifier.cs b/src/WixToolsetTest.BurnE2E/BundleVerifier.cs
index 94d51890..98ec96a0 100644
--- a/src/WixToolsetTest.BurnE2E/BundleVerifier.cs
+++ b/src/WixToolsetTest.BurnE2E/BundleVerifier.cs
@@ -33,14 +33,19 @@ namespace WixToolsetTest.BurnE2E
             return this.BundleSymbol;
         }
 
-        public string GetExpectedCachedBundlePath()
+        public string GetPackageCachePathForCacheId(string cacheId)
         {
-            var bundleSymbol = this.GetBundleSymbol();
-
             using var policyKey = Registry.LocalMachine.OpenSubKey(FULL_BURN_POLICY_REGISTRY_PATH);
             var redirectedCachePath = policyKey?.GetValue("PackageCache") as string;
             var cachePath = redirectedCachePath ?? Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.CommonApplicationData), PACKAGE_CACHE_FOLDER_NAME);
-            return Path.Combine(cachePath, bundleSymbol.BundleId, Path.GetFileName(this.Bundle));
+            return Path.Combine(cachePath, cacheId);
+        }
+
+        public string GetExpectedCachedBundlePath()
+        {
+            var bundleSymbol = this.GetBundleSymbol();
+            var cachePath = this.GetPackageCachePathForCacheId(bundleSymbol.BundleId);
+            return Path.Combine(cachePath, Path.GetFileName(this.Bundle));
         }
 
         public bool TryGetPerMachineRegistration(out BundleRegistration registration)
@@ -74,5 +79,18 @@ namespace WixToolsetTest.BurnE2E
             Assert.False(this.TryGetPerMachineRegistration(out _));
             Assert.False(File.Exists(cachedBundlePath));
         }
+
+        public void RemovePackageFromCache(string packageId)
+        {
+            using var wixOutput = WixOutput.Read(this.BundlePdb);
+            var intermediate = Intermediate.Load(wixOutput);
+            var section = intermediate.Sections.Single();
+            var packageSymbol = section.Symbols.OfType<WixBundlePackageSymbol>().Single(p => p.Id.Id == packageId);
+            var cachePath = this.GetPackageCachePathForCacheId(packageSymbol.CacheId);
+            if (Directory.Exists(cachePath))
+            {
+                Directory.Delete(cachePath, true);
+            }
+        }
     }
 }
diff --git a/src/WixToolsetTest.BurnE2E/BurnE2ETests.cs b/src/WixToolsetTest.BurnE2E/BurnE2ETests.cs
index f3b3e583..d3400658 100644
--- a/src/WixToolsetTest.BurnE2E/BurnE2ETests.cs
+++ b/src/WixToolsetTest.BurnE2E/BurnE2ETests.cs
@@ -12,32 +12,32 @@ namespace WixToolsetTest.BurnE2E
     {
         protected BurnE2ETests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
 
-        private Queue<IDisposable> Installers { get; } = new Queue<IDisposable>();
+        private Stack<IDisposable> Installers { get; } = new Stack<IDisposable>();
 
         protected BundleInstaller CreateBundleInstaller(string name)
         {
             var installer = new BundleInstaller(this.TestContext, name);
-            this.Installers.Enqueue(installer);
+            this.Installers.Push(installer);
             return installer;
         }
 
         protected PackageInstaller CreatePackageInstaller(string filename)
         {
             var installer = new PackageInstaller(this.TestContext, filename);
-            this.Installers.Enqueue(installer);
+            this.Installers.Push(installer);
             return installer;
         }
 
         protected TestBAController CreateTestBAController()
         {
             var controller = new TestBAController(this.TestContext);
-            this.Installers.Enqueue(controller);
+            this.Installers.Push(controller);
             return controller;
         }
 
         public void Dispose()
         {
-            while (this.Installers.TryDequeue(out var installer))
+            while (this.Installers.TryPop(out var installer))
             {
                 try
                 {
diff --git a/src/WixToolsetTest.BurnE2E/PackageInstaller.cs b/src/WixToolsetTest.BurnE2E/PackageInstaller.cs
index bd1e30cc..c3516fe7 100644
--- a/src/WixToolsetTest.BurnE2E/PackageInstaller.cs
+++ b/src/WixToolsetTest.BurnE2E/PackageInstaller.cs
@@ -40,7 +40,7 @@ namespace WixToolsetTest.BurnE2E
         /// <param name="expectedExitCode">Expected exit code</param>
         /// <param name="otherArguments">Other arguments to pass to MSIExec.</param>
         /// <returns>MSIExec log File</returns>
-        public string UninstallProduct(MSIExecReturnCode expectedExitCode, params string[] otherArguments)
+        public string UninstallProduct(MSIExecReturnCode expectedExitCode = MSIExecReturnCode.SUCCESS, params string[] otherArguments)
         {
             return this.RunMSIExec(MSIExecMode.Uninstall, otherArguments, expectedExitCode);
         }
@@ -51,7 +51,7 @@ namespace WixToolsetTest.BurnE2E
         /// <param name="expectedExitCode">Expected exit code</param>
         /// <param name="otherArguments">Other arguments to pass to msiexe.exe.</param>
         /// <returns>MSIExec log File</returns>
-        public string RepairProduct(MSIExecReturnCode expectedExitCode, params string[] otherArguments)
+        public string RepairProduct(MSIExecReturnCode expectedExitCode = MSIExecReturnCode.SUCCESS, params string[] otherArguments)
         {
             return this.RunMSIExec(MSIExecMode.Repair, otherArguments, expectedExitCode);
         }
diff --git a/src/WixToolsetTest.BurnE2E/PrereqBaTests.cs b/src/WixToolsetTest.BurnE2E/PrereqBaTests.cs
index b07b3aec..ced2e08e 100644
--- a/src/WixToolsetTest.BurnE2E/PrereqBaTests.cs
+++ b/src/WixToolsetTest.BurnE2E/PrereqBaTests.cs
@@ -18,7 +18,7 @@ namespace WixToolsetTest.BurnE2E
         ///   The preqba doesn't infinitely reload itself after failing to load the managed BA.
         ///   The engine automatically uninstalls the bundle since only permanent packages were installed.
         /// </summary>
-        [Fact(Skip = "https://github.com/wixtoolset/issues/issues/6297")]
+        [Fact]
         public void DncPreqBaDetectsInfiniteLoop()
         {
             var packageA = this.CreatePackageInstaller("PackageA");
@@ -49,7 +49,7 @@ namespace WixToolsetTest.BurnE2E
         ///   The preqba doesn't infinitely reload itself after failing to load the managed BA.
         ///   The engine automatically uninstalls the bundle since only permanent packages were installed.
         /// </summary>
-        [Fact(Skip = "https://github.com/wixtoolset/issues/issues/6297")]
+        [Fact]
         public void MbaPreqBaDetectsInfiniteLoop()
         {
             var packageB = this.CreatePackageInstaller("PackageB");
diff --git a/src/WixToolsetTest.BurnE2E/RegistrationTests.cs b/src/WixToolsetTest.BurnE2E/RegistrationTests.cs
new file mode 100644
index 00000000..640b5085
--- /dev/null
+++ b/src/WixToolsetTest.BurnE2E/RegistrationTests.cs
@@ -0,0 +1,51 @@
+// 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.
+
+namespace WixToolsetTest.BurnE2E
+{
+    using System;
+    using Xunit;
+    using Xunit.Abstractions;
+
+    public class RegistrationTests : BurnE2ETests
+    {
+        public RegistrationTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
+
+        [Fact]
+        public void AutomaticallyUninstallsBundleWithoutBADoingApply()
+        {
+            this.InstallBundleThenManuallyUninstallPackageAndRemovePackageFromCacheThenRunAndQuitWithoutApply(true);
+        }
+
+        [Fact]
+        public void AutomaticallyUninstallsBundleWithoutBADoingDetect()
+        {
+            this.InstallBundleThenManuallyUninstallPackageAndRemovePackageFromCacheThenRunAndQuitWithoutApply(false);
+        }
+
+        private void InstallBundleThenManuallyUninstallPackageAndRemovePackageFromCacheThenRunAndQuitWithoutApply(bool detect)
+        {
+            var packageA = this.CreatePackageInstaller("PackageA");
+            var bundleA = this.CreateBundleInstaller("BundleA");
+            var testBAController = this.CreateTestBAController();
+
+            bundleA.Install();
+            bundleA.VerifyRegisteredAndInPackageCache();
+            packageA.VerifyInstalled(true);
+
+            packageA.UninstallProduct();
+            bundleA.RemovePackageFromCache("PackageA");
+
+            if (detect)
+            {
+                testBAController.SetQuitAfterDetect();
+            }
+            else
+            {
+                testBAController.SetImmediatelyQuit();
+            }
+            bundleA.Install();
+            packageA.VerifyInstalled(false);
+            bundleA.VerifyUnregisteredAndRemovedFromPackageCache();
+        }
+    }
+}
diff --git a/src/WixToolsetTest.BurnE2E/TestBAController.cs b/src/WixToolsetTest.BurnE2E/TestBAController.cs
index 54a81b46..91d1b817 100644
--- a/src/WixToolsetTest.BurnE2E/TestBAController.cs
+++ b/src/WixToolsetTest.BurnE2E/TestBAController.cs
@@ -40,6 +40,16 @@ namespace WixToolsetTest.BurnE2E
             }
         }
 
+        public void SetImmediatelyQuit(string value = "true")
+        {
+            this.SetBurnTestValue("ImmediatelyQuit", value);
+        }
+
+        public void SetQuitAfterDetect(string value = "true")
+        {
+            this.SetBurnTestValue("QuitAfterDetect", value);
+        }
+
         /// <summary>
         /// Slows the cache progress of a package.
         /// </summary>