3 files changed, 23 insertions, 30 deletions

diff --git a/src/test/burn/TestData/BAFunctionsTests/BafRelatedBundleVariableTesting/BafRelatedBundleVariableTesting.cpp b/src/test/burn/TestData/BAFunctionsTests/BafRelatedBundleVariableTesting/BafRelatedBundleVariableTesting.cpp
index 01750b4b..ba559b05 100644
--- a/src/test/burn/TestData/BAFunctionsTests/BafRelatedBundleVariableTesting/BafRelatedBundleVariableTesting.cpp
+++ b/src/test/burn/TestData/BAFunctionsTests/BafRelatedBundleVariableTesting/BafRelatedBundleVariableTesting.cpp
@@ -29,37 +29,28 @@ public: //IBootstrapperApplication
 
         hr = BalGetRelatedBundleVariable(wzBundleId, STRING_VARIABLE, &wzValue);
 
-        ExitOnFailure(hr, "Failed to get related bundle string variable.");
-
-        if (wzValue)
-        {
-            BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "Retrieved related bundle variable with BAFunctions: AString = %ws", wzValue);
-        }
+        BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "Retrieved related bundle variable with BAFunctions: AString = %ws, Error: 0x%x", wzValue, hr);
 
         hr = BalGetRelatedBundleVariable(wzBundleId, NUMBER_VARIABLE, &wzValue);
 
-        if (wzValue)
-        {
-            BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "Retrieved related bundle variable with BAFunctions: ANumber = %ws", wzValue);
-        }
+        BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "Retrieved related bundle variable with BAFunctions: ANumber = %ws, Error: 0x%x", wzValue, hr);
 
         hr = __super::OnDetectRelatedBundle(wzBundleId, relationType, wzBundleTag, fPerMachine, wzVersion, fMissingFromCache, pfCancel);
-    LExit:
+
         ReleaseStr(wzValue);
-        return hr;        
+        return hr;
     }
+
 private:
- 
+
 
 public:
     //
     // Constructor - initialize member variables.
     //
     CBafRelatedBundleVariableTesting(
-        __in HMODULE hModule,
-        __in IBootstrapperEngine* pEngine,
-        __in const BA_FUNCTIONS_CREATE_ARGS* pArgs
-        ) : CBalBaseBAFunctions(hModule, pEngine, pArgs)
+        __in HMODULE hModule
+        ) : CBalBaseBAFunctions(hModule)
     {
     }
 
@@ -82,21 +73,21 @@ HRESULT WINAPI CreateBAFunctions(
 {
     HRESULT hr = S_OK;
     CBafRelatedBundleVariableTesting* pBAFunctions = NULL;
-    IBootstrapperEngine* pEngine = NULL;
 
-    hr = BalInitializeFromCreateArgs(pArgs->pBootstrapperCreateArgs, &pEngine);
-    ExitOnFailure(hr, "Failed to initialize Bal.");
+    BalInitialize(pArgs->pEngine);
 
-    pBAFunctions = new CBafRelatedBundleVariableTesting(hModule, pEngine, pArgs);
+    pBAFunctions = new CBafRelatedBundleVariableTesting(hModule);
     ExitOnNull(pBAFunctions, hr, E_OUTOFMEMORY, "Failed to create new CBafRelatedBundleVariableTesting object.");
 
+    hr = pBAFunctions->OnCreate(pArgs->pEngine, pArgs->pCommand);
+    ExitOnFailure(hr, "Failed to create BA function");
+
     pResults->pfnBAFunctionsProc = BalBaseBAFunctionsProc;
     pResults->pvBAFunctionsProcContext = pBAFunctions;
     pBAFunctions = NULL;
 
 LExit:
     ReleaseObject(pBAFunctions);
-    ReleaseObject(pEngine);
 
     return hr;
 }
diff --git a/src/test/burn/TestData/BAFunctionsTests/BundleAv1/BundleAv1.wxs b/src/test/burn/TestData/BAFunctionsTests/BundleAv1/BundleAv1.wxs
index 4aba7a2f..431ef3d1 100644
--- a/src/test/burn/TestData/BAFunctionsTests/BundleAv1/BundleAv1.wxs
+++ b/src/test/burn/TestData/BAFunctionsTests/BundleAv1/BundleAv1.wxs
@@ -1,15 +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. -->
 
-
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
   <Fragment>
-    <BootstrapperApplication>
+    <PayloadGroup Id="override ExtraPayloads">
       <Payload SourceFile="$(var.BafRelatedBundleVariableTesting.TargetPath)" bal:BAFunctions="yes" />
-    </BootstrapperApplication>
+    </PayloadGroup>
+
     <PackageGroup Id="BundlePackages">
       <MsiPackage Id="PackageA" SourceFile="$(var.PackageAv1.TargetPath)" />
     </PackageGroup>
+
     <Variable Name="ANumber" bal:Overridable="yes" Value="42" Persisted="yes" />
-    <Variable Name="AString" bal:Overridable="yes" Value="This is a test" Persisted="yes" /> 
+    <Variable Name="AString" bal:Overridable="yes" Value="This is a test" Persisted="yes" />
   </Fragment>
 </Wix>
diff --git a/src/test/burn/TestData/BAFunctionsTests/BundleAv2/BundleAv2.wxs b/src/test/burn/TestData/BAFunctionsTests/BundleAv2/BundleAv2.wxs
index 8608512c..9c512899 100644
--- a/src/test/burn/TestData/BAFunctionsTests/BundleAv2/BundleAv2.wxs
+++ b/src/test/burn/TestData/BAFunctionsTests/BundleAv2/BundleAv2.wxs
@@ -1,15 +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. -->
 
-
 <Wix xmlns="http://wixtoolset.org/schemas/v4/wxs" xmlns:bal="http://wixtoolset.org/schemas/v4/wxs/bal">
   <Fragment>
-    <BootstrapperApplication>
+    <PayloadGroup Id="override ExtraPayloads">
       <Payload SourceFile="$(var.BafRelatedBundleVariableTesting.TargetPath)" bal:BAFunctions="yes" />
-    </BootstrapperApplication>
+    </PayloadGroup>
+
     <PackageGroup Id="BundlePackages">
       <MsiPackage Id="PackageA" SourceFile="$(var.PackageAv2.TargetPath)" />
     </PackageGroup>
+
     <Variable Name="ANumber" bal:Overridable="yes" Value="4242" Persisted="yes" />
-    <Variable Name="AString" bal:Overridable="yes" Value="This is a test v2" Persisted="yes" /> 
+    <Variable Name="AString" bal:Overridable="yes" Value="This is a test v2" Persisted="yes" />
   </Fragment>
 </Wix>