Make the estimated size in ARP a little more accurate.

Fixes 4039

1 files changed, 17 insertions, 0 deletions

diff --git a/src/test/burn/TestBA/TestBA.cs b/src/test/burn/TestBA/TestBA.cs
index fdbcc6d4..b4d74341 100644
--- a/src/test/burn/TestBA/TestBA.cs
+++ b/src/test/burn/TestBA/TestBA.cs
@@ -27,6 +27,7 @@ namespace WixToolset.Test.BA
 
         private string updateBundlePath;
 
+        private bool forceKeepRegistration;
         private bool immediatelyQuit;
         private bool quitAfterDetect;
         private bool explicitlyElevateAndPlanFromOnElevateBegin;
@@ -131,6 +132,12 @@ namespace WixToolset.Test.BA
                 this.explicitlyElevateAndPlanFromOnElevateBegin = false;
             }
 
+            string forceKeepRegistration = this.ReadPackageAction(null, "ForceKeepRegistration");
+            if (String.IsNullOrEmpty(forceKeepRegistration) || !Boolean.TryParse(forceKeepRegistration, out this.forceKeepRegistration))
+            {
+                this.forceKeepRegistration = false;
+            }
+
             string quitAfterDetect = this.ReadPackageAction(null, "QuitAfterDetect");
             if (String.IsNullOrEmpty(quitAfterDetect) || !Boolean.TryParse(quitAfterDetect, out this.quitAfterDetect))
             {
@@ -533,6 +540,16 @@ namespace WixToolset.Test.BA
             this.ShutdownUiThread();
         }
 
+        protected override void OnUnregisterBegin(UnregisterBeginEventArgs args)
+        {
+            if (this.forceKeepRegistration && args.RegistrationType == RegistrationType.None)
+            {
+                args.RegistrationType = RegistrationType.InProgress;
+            }
+
+            this.Log("OnUnregisterBegin, default: {0}, requested: {1}", args.RecommendedRegistrationType, args.RegistrationType);
+        }
+
         private void TestVariables()
         {
             // First make sure we can check and get standard variables of each type.