1 files changed, 27 insertions, 10 deletions

diff --git a/src/test/burn/TestData/PrereqBaTests/PrereqBaf/PrereqBaf.cpp b/src/test/burn/TestData/PrereqBaTests/PrereqBaf/PrereqBaf.cpp
index 35949eb9..c93f9d3f 100644
--- a/src/test/burn/TestData/PrereqBaTests/PrereqBaf/PrereqBaf.cpp
+++ b/src/test/burn/TestData/PrereqBaTests/PrereqBaf/PrereqBaf.cpp
@@ -9,6 +9,22 @@ class CPrereqBaf : public CBalBaseBAFunctions
 public: // IBAFunctions
 
 public: //IBootstrapperApplication
+    STDMETHODIMP OnCreate(
+        __in IBootstrapperEngine* pEngine,
+        __in BOOTSTRAPPER_COMMAND* pCommand
+    )
+    {
+        HRESULT hr = S_OK;
+
+        hr = __super::OnCreate(pEngine, pCommand);
+        ExitOnFailure(hr, "CBalBaseBootstrapperApplication initialization failed.");
+
+        hr = StrAllocString(&m_sczBARuntimeDirectory, pCommand->wzBootstrapperWorkingFolder, 0);
+        ExitOnFailure(hr, "Failed to copy working folder");
+
+    LExit:
+        return hr;
+    }
 
     virtual STDMETHODIMP OnDetectBegin(
         __in BOOL /*fCached*/,
@@ -19,7 +35,7 @@ public: //IBootstrapperApplication
     {
         HRESULT hr = S_OK;
 
-        hr = m_pEngine->SetVariableString(L"BARuntimeDirectory", m_command.wzBootstrapperWorkingFolder, FALSE);
+        hr = m_pEngine->SetVariableString(L"BARuntimeDirectory", m_sczBARuntimeDirectory, FALSE);
         ExitOnFailure(hr, "Failed to set BARuntimeDirectory");
 
     LExit:
@@ -33,11 +49,10 @@ public:
     // Constructor - initialize member variables.
     //
     CPrereqBaf(
-        __in HMODULE hModule,
-        __in IBootstrapperEngine* pEngine,
-        __in const BA_FUNCTIONS_CREATE_ARGS* pArgs
-        ) : CBalBaseBAFunctions(hModule, pEngine, pArgs)
+        __in HMODULE hModule
+        ) : CBalBaseBAFunctions(hModule)
     {
+        m_sczBARuntimeDirectory = NULL;
     }
 
     //
@@ -45,9 +60,11 @@ public:
     //
     ~CPrereqBaf()
     {
+        ReleaseNullStr(m_sczBARuntimeDirectory);
     }
 
 private:
+    LPWSTR m_sczBARuntimeDirectory;
 };
 
 
@@ -59,21 +76,21 @@ HRESULT WINAPI CreateBAFunctions(
 {
     HRESULT hr = S_OK;
     CPrereqBaf* pBAFunctions = NULL;
-    IBootstrapperEngine* pEngine = NULL;
 
-    hr = BalInitializeFromCreateArgs(pArgs->pBootstrapperCreateArgs, &pEngine);
-    ExitOnFailure(hr, "Failed to initialize Bal.");
+    BalInitialize(pArgs->pEngine);
 
-    pBAFunctions = new CPrereqBaf(hModule, pEngine, pArgs);
+    pBAFunctions = new CPrereqBaf(hModule);
     ExitOnNull(pBAFunctions, hr, E_OUTOFMEMORY, "Failed to create new CPrereqBaf object.");
 
+    hr = pBAFunctions->OnCreate(pArgs->pEngine, pArgs->pCommand);
+    ExitOnFailure(hr, "Failed to call OnCreate CPrereqBaf.");
+
     pResults->pfnBAFunctionsProc = BalBaseBAFunctionsProc;
     pResults->pvBAFunctionsProcContext = pBAFunctions;
     pBAFunctions = NULL;
 
 LExit:
     ReleaseObject(pBAFunctions);
-    ReleaseObject(pEngine);
 
     return hr;
 }