4 files changed, 102 insertions, 8 deletions

diff --git a/src/test/burn/WixToolsetTest.BurnE2E/ExePackageTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/ExePackageTests.cs
index a7438564..e7c7ee5b 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/ExePackageTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/ExePackageTests.cs
@@ -30,7 +30,7 @@ namespace WixToolsetTest.BurnE2E
             perMachineArpEntryExePackageBundle.VerifyUnregisteredAndRemovedFromPackageCache();
             arpEntryExePackage.VerifyRegistered(false);
 
-            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}"));
+            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}\""));
         }
 
         [RuntimeFact]
@@ -52,7 +52,7 @@ namespace WixToolsetTest.BurnE2E
             perMachineArpEntryExePackageBundle.VerifyUnregisteredAndRemovedFromPackageCache();
             arpEntryExePackage.VerifyRegistered(false);
 
-            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}"));
+            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}\""));
         }
 
         [RuntimeFact]
@@ -85,7 +85,7 @@ namespace WixToolsetTest.BurnE2E
             packageTestExe.VerifyInstalled(true);
 
             Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, "TESTBA: OnCachePackageNonVitalValidationFailure() - id: TestExe, default: None, requested: Acquire"));
-            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}"));
+            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}\""));
             Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"TestExe.exe\" /regw \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId},DisplayVersion,String,1.0.0.0\" /regw \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId},QuietUninstallString,String,\\\""));
         }
 
@@ -113,7 +113,7 @@ namespace WixToolsetTest.BurnE2E
             packageTestExe.VerifyInstalled(true);
 
             Assert.False(LogVerifier.MessageInLogFile(uninstallLogPath, "TESTBA: OnCachePackageNonVitalValidationFailure() - id: TestExe"));
-            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}"));
+            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}\""));
             Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"TestExe.exe\" /regw \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId},DisplayVersion,String,1.0.0.0\" /regw \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId},QuietUninstallString,String,\\\""));
         }
 
@@ -144,7 +144,7 @@ namespace WixToolsetTest.BurnE2E
             packageTestExe.VerifyInstalled(true);
 
             Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, "TESTBA: OnCachePackageNonVitalValidationFailure() - id: TestExe, default: None, requested: None"));
-            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}"));
+            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}\""));
         }
 
         [RuntimeFact]
@@ -161,7 +161,7 @@ namespace WixToolsetTest.BurnE2E
             arpEntryExePackage.VerifyRegistered(false);
 
             Assert.True(LogVerifier.MessageInLogFile(installLogPath, $"TestExe.exe\" /regw \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId},DisplayVersion,String,1.0.0.0\" /regw \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId},QuietUninstallString,String,\\\""));
-            Assert.True(LogVerifier.MessageInLogFile(installLogPath, $"testexe.exe\" /regd HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}"));
+            Assert.True(LogVerifier.MessageInLogFile(installLogPath, $"testexe.exe\" /regd \"HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}\""));
         }
 
         [RuntimeFact]
@@ -209,7 +209,7 @@ namespace WixToolsetTest.BurnE2E
             perUserArpEntryExePackageBundle.VerifyUnregisteredAndRemovedFromPackageCache();
             arpEntryExePackage.VerifyRegistered(false);
 
-            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd HKCU\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}"));
+            Assert.True(LogVerifier.MessageInLogFile(uninstallLogPath, $"testexe.exe\" /regd \"HKCU\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\{arpId}\""));
         }
 
         [RuntimeFact]
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/FailureTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/FailureTests.cs
index bd3762ed..7d399df0 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/FailureTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/FailureTests.cs
@@ -2,7 +2,6 @@
 
 namespace WixToolsetTest.BurnE2E
 {
-    using System.Threading;
     using WixTestTools;
     using WixToolset.BootstrapperApplicationApi;
     using Xunit;
@@ -13,6 +12,42 @@ namespace WixToolsetTest.BurnE2E
         public FailureTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
 
         [RuntimeFact]
+        public void CanSetDownloadSourceAndForceDownload()
+        {
+            var packageA = this.CreatePackageInstaller("PackageA");
+            var bundleA = this.CreateBundleInstaller("BundleA");
+            var testBAController = this.CreateTestBAController();
+
+            testBAController.SetPackageForceDownloadSource("PackageA", "https://1e1bf2be1c384fd1a0c4c0500eef971b/downloads/payloads/{0}");
+
+            packageA.VerifyInstalled(false);
+
+            bundleA.VerifyUnregisteredAndRemovedFromPackageCache();
+
+            bundleA.Install(0x2ee7/*ERROR_INTERNET_NAME_NOT_RESOLVED*/);
+
+            packageA.VerifyInstalled(false);
+        }
+
+        [RuntimeFact]
+        public void CanSetUpdateSource()
+        {
+            var packageA = this.CreatePackageInstaller("PackageA");
+            var bundleA = this.CreateBundleInstaller("BundleUpdate");
+            var testBAController = this.CreateTestBAController();
+
+            testBAController.SetForceUpdateSource("https://1e1bf2be1c384fd1a0c4c0500eef971b/update_feed.atom.xml");
+
+            packageA.VerifyInstalled(false);
+
+            bundleA.VerifyUnregisteredAndRemovedFromPackageCache();
+
+            bundleA.Install();
+
+            packageA.VerifyInstalled(true);
+        }
+
+        [RuntimeFact]
         public void CanCancelExePackageAndAbandonIt()
         {
             var bundleD = this.CreateBundleInstaller("BundleD");
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/OptionalUpdateRegistrationTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/OptionalUpdateRegistrationTests.cs
new file mode 100644
index 00000000..3731b0cc
--- /dev/null
+++ b/src/test/burn/WixToolsetTest.BurnE2E/OptionalUpdateRegistrationTests.cs
@@ -0,0 +1,44 @@
+// 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 WixTestTools;
+    using Xunit;
+    using Xunit.Abstractions;
+
+    public class OptionalUpdateRegistrationTests : BurnE2ETests
+    {
+        public OptionalUpdateRegistrationTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
+
+        [RuntimeFact]
+        public void BundleUpdateRegistrationIsStickyAndAccurateAcrossUpgrades()
+        {
+            var packageAv1 = this.CreatePackageInstaller("PackageAv1");
+            var packageAv2 = this.CreatePackageInstaller("PackageAv2");
+            var bundleAv1 = this.CreateBundleInstaller("BundleAv1");
+            var bundleAv2 = this.CreateBundleInstaller("BundleAv2");
+
+            bundleAv1.Install();
+            bundleAv1.VerifyRegisteredAndInPackageCache();
+            var gotV1Registration = bundleAv1.TryGetUpdateRegistration(out var v1Registration);
+
+            bundleAv2.Install();
+            bundleAv2.VerifyRegisteredAndInPackageCache();
+            var gotV2Registration = bundleAv2.TryGetUpdateRegistration(out var v2Registration);
+
+            bundleAv1.VerifyUnregisteredAndRemovedFromPackageCache();
+
+            Assert.True(gotV1Registration, "Missing update registration after v1 install.");
+            Assert.True(gotV2Registration, "Missing update registration after v2 upgrade.");
+
+            Assert.Equal("Acme", v1Registration.Publisher);
+            Assert.Equal("Acme", v2Registration.Publisher);
+            Assert.Equal("Setup Geeks", v1Registration.PublishingGroup);
+            Assert.Equal("Setup Geeks", v2Registration.PublishingGroup);
+            Assert.Equal("~OptionalUpdateRegistrationTests", v1Registration.PackageName);
+            Assert.Equal("~OptionalUpdateRegistrationTests", v2Registration.PackageName);
+            Assert.Equal("1.0.0.0", v1Registration.PackageVersion);
+            Assert.Equal("2.0.0.0", v2Registration.PackageVersion);
+        }
+    }
+}
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/Utilities/TestBAController.cs b/src/test/burn/WixToolsetTest.BurnE2E/Utilities/TestBAController.cs
index a378545e..c344ebce 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/Utilities/TestBAController.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/Utilities/TestBAController.cs
@@ -67,6 +67,11 @@ namespace WixToolsetTest.BurnE2E
             this.SetBurnTestValue("QuitAfterDetect", value);
         }
 
+        public void SetForceUpdateSource(string url)
+        {
+            this.SetBurnTestValue("ForceUpdateSource", url);
+        }
+
         /// <summary>
         /// Slows the cache progress of a package.
         /// </summary>
@@ -108,6 +113,16 @@ namespace WixToolsetTest.BurnE2E
         }
 
         /// <summary>
+        /// Forces a download action and sets the download source.
+        /// </summary>
+        /// <param name="packageId">Package identity.</param>
+        /// <param name="actionName">The URL format string.</param>
+        public void SetPackageForceDownloadSource(string packageId, string url)
+        {
+            this.SetPackageState(packageId, "ForceDownloadSource", url);
+        }
+
+        /// <summary>
         /// Cancels the execute of a package at the next progess after the specified MSI action start.
         /// </summary>
         /// <param name="packageId">Package identity.</param>