1 files changed, 50 insertions, 0 deletions

diff --git a/src/test/burn/TestBA/TestBA.cs b/src/test/burn/TestBA/TestBA.cs
index 1556acd0..51fc1cbd 100644
--- a/src/test/burn/TestBA/TestBA.cs
+++ b/src/test/burn/TestBA/TestBA.cs
@@ -41,6 +41,8 @@ namespace WixToolset.Test.BA
         private int cancelOnProgressAtProgress;
         private int retryExecuteFilesInUse;
         private bool rollingBack;
+        private string forceDownloadSource;
+        private string forceUpdateSource;
 
         private IBootstrapperCommand Command { get; set; }
 
@@ -216,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;
@@ -375,6 +396,12 @@ namespace WixToolset.Test.BA
             {
                 this.Log("    CancelCacheAtProgress: {0}", this.cancelCacheAtProgress);
             }
+
+            this.forceDownloadSource = this.ReadPackageAction(args.PackageId, "ForceDownloadSource");
+            if (!String.IsNullOrEmpty(this.forceDownloadSource))
+            {
+                this.Log("    ForceDownloadSource: {0}", this.forceDownloadSource);
+            }
         }
 
         protected override void OnCachePackageNonVitalValidationFailure(CachePackageNonVitalValidationFailureEventArgs args)
@@ -387,6 +414,29 @@ namespace WixToolset.Test.BA
             this.Log("OnCachePackageNonVitalValidationFailure() - id: {0}, default: {1}, requested: {2}", args.PackageId, args.Recommendation, args.Action);
         }
 
+        protected override void OnCacheAcquireResolving(CacheAcquireResolvingEventArgs args)
+        {
+            if (!String.IsNullOrEmpty(this.forceDownloadSource))
+            {
+                args.Action = CacheResolveOperation.Download;
+                var url = String.Format(this.forceDownloadSource, args.PayloadId);
+
+                this.Log("OnCacheAcquireResolving: {0} => {1}", this.forceDownloadSource, url);
+
+                this.engine.SetDownloadSource(
+                    String.Empty,
+                    args.PayloadId,
+                    url,
+                    String.Empty,
+                    String.Empty,
+                    String.Empty);
+            }
+            else
+            {
+                this.Log("OnCacheAcquireResolving not forcing download");
+            }
+        }
+
         protected override void OnCacheAcquireProgress(CacheAcquireProgressEventArgs args)
         {
             this.Log("OnCacheAcquireProgress() - container/package: {0}, payload: {1}, progress: {2}, total: {3}, overall progress: {4}%", args.PackageOrContainerId, args.PayloadId, args.Progress, args.Total, args.OverallPercentage);