5 files changed, 90 insertions, 0 deletions

diff --git a/src/test/burn/TestBA/TestBA.cs b/src/test/burn/TestBA/TestBA.cs
index 4f75a055..51fc1cbd 100644
--- a/src/test/burn/TestBA/TestBA.cs
+++ b/src/test/burn/TestBA/TestBA.cs
@@ -42,6 +42,7 @@ namespace WixToolset.Test.BA
         private int retryExecuteFilesInUse;
         private bool rollingBack;
         private string forceDownloadSource;
+        private string forceUpdateSource;
 
         private IBootstrapperCommand Command { get; set; }
 
@@ -217,9 +218,28 @@ namespace WixToolset.Test.BA
             }
         }
 
+        protected override void OnDetectBegin(DetectBeginEventArgs args)
+        {
+            this.Log("OnDetectBegin");
+
+            this.forceUpdateSource = this.ReadPackageAction(null, "ForceUpdateSource");
+            if (!String.IsNullOrEmpty(this.forceUpdateSource))
+            {
+                this.Log("    OnDetectBegin::ForceUpdateSource: {0}", this.forceUpdateSource);
+            }
+
+        }
+
         protected override void OnDetectUpdateBegin(DetectUpdateBeginEventArgs args)
         {
             this.Log("OnDetectUpdateBegin");
+
+            if (!String.IsNullOrEmpty(this.forceUpdateSource))
+            {
+                this.Log("    OnDetectUpdateBegin::ForceUpdateSource: {0}", this.forceUpdateSource);
+                this.Engine.SetUpdateSource(this.forceUpdateSource, String.Empty);
+            }
+
             if (LaunchAction.UpdateReplaceEmbedded == this.action || LaunchAction.UpdateReplace == this.action)
             {
                 args.Skip = false;
diff --git a/src/test/burn/TestData/FailureTests/BundleUpdate/Bundle.wxs b/src/test/burn/TestData/FailureTests/BundleUpdate/Bundle.wxs
new file mode 100644
index 00000000..0af817c1
--- /dev/null
+++ b/src/test/burn/TestData/FailureTests/BundleUpdate/Bundle.wxs
@@ -0,0 +1,31 @@
+<!-- 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. -->
+
+<?ifndef TestVersion?>
+<?define TestVersion = 1.0.0.0?>
+<?endif?>
+
+<Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
+  <Bundle 
+    Name="~$(var.TestGroupName) - $(var.BundleName)" 
+    Version="$(var.TestVersion)" 
+    Id="WixToolset.Burn.SetUpdateSourceTest" 
+    Compressed="yes">
+    
+    <Update Location="updateurl" />
+    
+    <Log Prefix="~$(var.TestGroupName)_$(var.BundleName)" />
+
+    <Variable Name="TestGroupName" Value="$(var.TestGroupName)" />
+
+    <Chain>
+      <PackageGroupRef Id="TestBA" />
+      <PackageGroupRef Id="BundlePackages" />
+    </Chain>
+  </Bundle>
+
+  <Fragment>
+    <PackageGroup Id="BundlePackages">
+      <MsiPackage Id="PackageA" SourceFile="$(var.PackageA.TargetPath)" />
+    </PackageGroup>
+  </Fragment>
+</Wix>
diff --git a/src/test/burn/TestData/FailureTests/BundleUpdate/BundleUpdate.wixproj b/src/test/burn/TestData/FailureTests/BundleUpdate/BundleUpdate.wixproj
new file mode 100644
index 00000000..34ab1aef
--- /dev/null
+++ b/src/test/burn/TestData/FailureTests/BundleUpdate/BundleUpdate.wixproj
@@ -0,0 +1,16 @@
+<!-- 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>{C60B9483-CE87-4FDA-AE5A-B39A52E956E8}</UpgradeCode>
+  </PropertyGroup>
+  <ItemGroup>
+    <ProjectReference Include="..\PackageA\PackageA.wixproj" />
+    <ProjectReference Include="..\PackageB\PackageB.wixproj" />
+    <ProjectReference Include="..\..\TestBA\TestBAWixlib\testbawixlib.wixproj" />
+  </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.BootstrapperApplications.wixext" />
+    <PackageReference Include="WixToolset.NetFx.wixext" />
+  </ItemGroup>
+</Project>
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/FailureTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/FailureTests.cs
index 4bbfc7e9..7d399df0 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/FailureTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/FailureTests.cs
@@ -30,6 +30,24 @@ namespace WixToolsetTest.BurnE2E
         }
 
         [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/Utilities/TestBAController.cs b/src/test/burn/WixToolsetTest.BurnE2E/Utilities/TestBAController.cs
index ca395f56..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>