Make the estimated size in ARP a little more accurate.

Fixes 4039

7 files changed, 63 insertions, 11 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.
diff --git a/src/test/burn/WixTestTools/BundleVerifier.cs b/src/test/burn/WixTestTools/BundleVerifier.cs
index 56044c5f..594b19aa 100644
--- a/src/test/burn/WixTestTools/BundleVerifier.cs
+++ b/src/test/burn/WixTestTools/BundleVerifier.cs
@@ -87,7 +87,7 @@ namespace WixTestTools
             }
         }
 
-        public string VerifyRegisteredAndInPackageCache(int? expectedSystemComponent = null)
+        public BundleRegistration VerifyRegisteredAndInPackageCache(int? expectedSystemComponent = null)
         {
             Assert.True(this.TryGetRegistration(out var registration));
 
@@ -99,7 +99,7 @@ namespace WixTestTools
             var expectedCachePath = this.GetExpectedCachedBundlePath();
             WixAssert.StringEqual(expectedCachePath, registration.CachePath, true);
 
-            return registration.CachePath;
+            return registration;
         }
 
         public void VerifyUnregisteredAndRemovedFromPackageCache()
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/BasicFunctionalityTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/BasicFunctionalityTests.cs
index 89c5be9b..efe4f05b 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/BasicFunctionalityTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/BasicFunctionalityTests.cs
@@ -82,7 +82,8 @@ namespace WixToolsetTest.BurnE2E
 
             bundle.Install();
 
-            var cachedBundlePath = bundle.VerifyRegisteredAndInPackageCache();
+            var registration = bundle.VerifyRegisteredAndInPackageCache();
+            var cachedBundlePath = registration.CachePath;
 
             // Source file should be installed
             Assert.True(File.Exists(packageSourceCodeInstalled), $"Should have found {packageName} payload installed at: {packageSourceCodeInstalled}");
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/CacheTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/CacheTests.cs
index 943939e3..4d5fb811 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/CacheTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/CacheTests.cs
@@ -162,14 +162,14 @@ namespace WixToolsetTest.BurnE2E
                 bundleA.Install(bundleACopiedPath);
             }
 
-            var bundlePackageCachePath = bundleA.VerifyRegisteredAndInPackageCache();
+            var bundlePackageRegistration = bundleA.VerifyRegisteredAndInPackageCache();
 
             packageA.VerifyInstalled(true);
             packageB.VerifyInstalled(false);
 
             testBAController.SetPackageRequestedState("PackageB", RequestState.Present);
 
-            var modifyLogPath = bundleA.Modify(bundlePackageCachePath);
+            var modifyLogPath = bundleA.Modify(bundlePackageRegistration.CachePath);
             bundleA.VerifyRegisteredAndInPackageCache();
 
             packageA.VerifyInstalled(true);
@@ -204,14 +204,14 @@ namespace WixToolsetTest.BurnE2E
 
                 bundleB.Install(bundleBCopiedPath);
 
-                var bundlePackageCachePath = bundleB.VerifyRegisteredAndInPackageCache();
+                var bundlePackageRegistration = bundleB.VerifyRegisteredAndInPackageCache();
 
                 packageA.VerifyInstalled(true);
                 packageB.VerifyInstalled(false);
 
                 testBAController.SetPackageRequestedState("PackageB", RequestState.Present);
 
-                bundleB.Modify(bundlePackageCachePath);
+                bundleB.Modify(bundlePackageRegistration.CachePath);
                 bundleB.VerifyRegisteredAndInPackageCache();
 
                 packageA.VerifyInstalled(true);
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/MsiTransactionTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/MsiTransactionTests.cs
index 3d9748bb..cbfee806 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/MsiTransactionTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/MsiTransactionTests.cs
@@ -42,7 +42,7 @@ namespace WixToolsetTest.BurnE2E
 
             bundleAv1.Install();
 
-            var bundleAv1CachedPath = bundleAv1.VerifyRegisteredAndInPackageCache();
+            var bundleAv1Registration = bundleAv1.VerifyRegisteredAndInPackageCache();
 
             // Source file should be installed
             Assert.True(File.Exists(packageASourceCodeInstalled), String.Concat("Should have found Package A payload installed at: ", packageASourceCodeInstalled));
@@ -51,7 +51,7 @@ namespace WixToolsetTest.BurnE2E
 
             bundleAv2.Install();
 
-            var bundleAv2CachedPath = bundleAv2.VerifyRegisteredAndInPackageCache();
+            var bundleAv2Registration = bundleAv2.VerifyRegisteredAndInPackageCache();
 
             // Source file should be upgraded
             Assert.True(File.Exists(packageDSourceCodeInstalled), String.Concat("Should have found Package D payload installed at: ", packageDSourceCodeInstalled));
@@ -61,7 +61,7 @@ namespace WixToolsetTest.BurnE2E
             Assert.False(File.Exists(packageBv1SourceCodeInstalled), String.Concat("Package Bv1 payload should have been removed by upgrade uninstall from: ", packageBv1SourceCodeInstalled));
             Assert.False(File.Exists(packageASourceCodeInstalled), String.Concat("Package A payload should have been removed by upgrade uninstall from: ", packageASourceCodeInstalled));
 
-            bundleAv1.VerifyUnregisteredAndRemovedFromPackageCache(bundleAv1CachedPath);
+            bundleAv1.VerifyUnregisteredAndRemovedFromPackageCache(bundleAv1Registration.CachePath);
 
             // Uninstall everything.
             bundleAv2.Uninstall();
@@ -71,7 +71,7 @@ namespace WixToolsetTest.BurnE2E
             Assert.False(File.Exists(packageBv2SourceCodeInstalled), String.Concat("Package Bv2 payload should have been removed by uninstall from: ", packageBv2SourceCodeInstalled));
             Assert.False(File.Exists(packageCv2SourceCodeInstalled), String.Concat("Package Cv2 payload should have been removed by uninstall from: ", packageCv2SourceCodeInstalled));
 
-            bundleAv2.VerifyUnregisteredAndRemovedFromPackageCache(bundleAv2CachedPath);
+            bundleAv2.VerifyUnregisteredAndRemovedFromPackageCache(bundleAv2Registration.CachePath);
         }
 
         /// <summary>
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/RegistrationTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/RegistrationTests.cs
index 51122c28..01ffa942 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/RegistrationTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/RegistrationTests.cs
@@ -3,6 +3,7 @@
 namespace WixToolsetTest.BurnE2E
 {
     using System;
+    using WixToolset.Mba.Core;
     using Xunit;
     using Xunit.Abstractions;
 
@@ -11,8 +12,35 @@ namespace WixToolsetTest.BurnE2E
         public RegistrationTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
 
         [Fact]
+        public void AllowsBAToKeepRegistration()
+        {
+            this.CreatePackageInstaller("PackageA");
+            var bundleA = this.CreateBundleInstaller("BundleA");
+            var testBAController = this.CreateTestBAController();
+
+            testBAController.SetPackageRequestedState("PackageA", RequestState.Absent);
+            testBAController.SetForceKeepRegistration();
+
+            bundleA.Install();
+            var initialRegistration = bundleA.VerifyRegisteredAndInPackageCache();
+
+            Assert.NotNull(initialRegistration.EstimatedSize);
+
+            testBAController.SetForceKeepRegistration(null);
+            testBAController.ResetPackageStates("PackageA");
+
+            bundleA.Install();
+            var finalRegistration = bundleA.VerifyRegisteredAndInPackageCache();
+
+            // Verifies https://github.com/wixtoolset/issues/issues/4039
+            Assert.NotNull(finalRegistration.EstimatedSize);
+            Assert.InRange(finalRegistration.EstimatedSize.Value, initialRegistration.EstimatedSize.Value + 1, Int32.MaxValue);
+        }
+
+        [Fact]
         public void AutomaticallyUncachesBundleWhenNotInstalled()
         {
+            this.CreatePackageInstaller("PackageA");
             var bundleA = this.CreateBundleInstaller("BundleA");
             var testBAController = this.CreateTestBAController();
 
@@ -40,6 +68,7 @@ namespace WixToolsetTest.BurnE2E
         [Fact]
         public void RegistersInARPIfPrecached()
         {
+            this.CreatePackageInstaller("PackageA");
             var bundleA = this.CreateBundleInstaller("BundleA");
 
             bundleA.ManuallyCache();
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/TestBAController.cs b/src/test/burn/WixToolsetTest.BurnE2E/TestBAController.cs
index fa553919..8e6611a2 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/TestBAController.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/TestBAController.cs
@@ -47,6 +47,11 @@ namespace WixToolsetTest.BurnE2E
             this.SetBurnTestValue("ExplicitlyElevateAndPlanFromOnElevateBegin", value);
         }
 
+        public void SetForceKeepRegistration(string value = "true")
+        {
+            this.SetBurnTestValue("ForceKeepRegistration", value);
+        }
+
         public void SetImmediatelyQuit(string value = "true")
         {
             this.SetBurnTestValue("ImmediatelyQuit", value);