1 files changed, 64 insertions, 0 deletions

diff --git a/src/burn/test/BurnUnitTest/PlanTest.cpp b/src/burn/test/BurnUnitTest/PlanTest.cpp
index 99644115..81484234 100644
--- a/src/burn/test/BurnUnitTest/PlanTest.cpp
+++ b/src/burn/test/BurnUnitTest/PlanTest.cpp
@@ -170,6 +170,10 @@ namespace Bootstrapper
             Assert::Equal(7ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            ValidateRestoreRelatedBundle(pPlan, dwIndex++, L"{FD9920AD-DBCA-4C6C-8CD5-B47431CE8D21}", BOOTSTRAPPER_ACTION_STATE_INSTALL, NULL);
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -277,6 +281,9 @@ namespace Bootstrapper
             Assert::Equal(3ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageC");
             ValidateCleanAction(pPlan, dwIndex++, L"PackageB");
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
@@ -350,6 +357,9 @@ namespace Bootstrapper
             Assert::Equal(1ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             ValidateCleanCompatibleAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
@@ -438,6 +448,9 @@ namespace Bootstrapper
             Assert::Equal(2ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -508,6 +521,9 @@ namespace Bootstrapper
             Assert::Equal(1ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -579,6 +595,9 @@ namespace Bootstrapper
             Assert::Equal(2ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
@@ -658,6 +677,10 @@ namespace Bootstrapper
             Assert::Equal(3ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            ValidateRestoreRelatedBundle(pPlan, dwIndex++, L"{FD9920AD-DBCA-4C6C-8CD5-B47431CE8D21}", BOOTSTRAPPER_ACTION_STATE_INSTALL, NULL);
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -744,6 +767,10 @@ namespace Bootstrapper
             Assert::Equal(3ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            ValidateRestoreRelatedBundle(pPlan, dwIndex++, L"{FD9920AD-DBCA-4C6C-8CD5-B47431CE8D21}", BOOTSTRAPPER_ACTION_STATE_INSTALL, NULL);
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -805,6 +832,9 @@ namespace Bootstrapper
             Assert::Equal(0ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
@@ -879,6 +909,9 @@ namespace Bootstrapper
             Assert::Equal(1ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
@@ -946,6 +979,9 @@ namespace Bootstrapper
             Assert::Equal(0ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -1057,6 +1093,9 @@ namespace Bootstrapper
             Assert::Equal(5ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -1147,6 +1186,9 @@ namespace Bootstrapper
             Assert::Equal(2ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PatchA");
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
@@ -1829,6 +1871,28 @@ namespace Bootstrapper
             NativeAssert::StringEqual(wzName, pProvider->sczName);
         }
 
+        void ValidateRestoreRelatedBundle(
+            __in BURN_PLAN* pPlan,
+            __in DWORD dwIndex,
+            __in LPCWSTR wzPackageId,
+            __in BOOTSTRAPPER_ACTION_STATE action,
+            __in LPCWSTR wzIgnoreDependencies
+            )
+        {
+            BURN_EXECUTE_ACTION* pAction = ValidateRestoreRelatedBundleActionExists(pPlan, dwIndex);
+            Assert::Equal<DWORD>(BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE, pAction->type);
+            NativeAssert::StringEqual(wzPackageId, pAction->relatedBundle.pRelatedBundle->package.sczId);
+            Assert::Equal<DWORD>(action, pAction->relatedBundle.action);
+            NativeAssert::StringEqual(wzIgnoreDependencies, pAction->relatedBundle.sczIgnoreDependencies);
+            Assert::Equal<BOOL>(FALSE, pAction->fDeleted);
+        }
+
+        BURN_EXECUTE_ACTION* ValidateRestoreRelatedBundleActionExists(BURN_PLAN* pPlan, DWORD dwIndex)
+        {
+            Assert::InRange(dwIndex + 1ul, 1ul, pPlan->cRestoreRelatedBundleActions);
+            return pPlan->rgRestoreRelatedBundleActions + dwIndex;
+        }
+
         void ValidateUninstallMsiCompatiblePackage(
             __in BURN_PLAN* pPlan,
             __in BOOL fRollback,