7 files changed, 32 insertions, 59 deletions

diff --git a/src/api/burn/test/BalUtilUnitTest/BAFunctionsTests.cpp b/src/api/burn/test/BalUtilUnitTest/BAFunctionsTests.cpp
index 2ae539ee..ba04c1bb 100644
--- a/src/api/burn/test/BalUtilUnitTest/BAFunctionsTests.cpp
+++ b/src/api/burn/test/BalUtilUnitTest/BAFunctionsTests.cpp
@@ -12,32 +12,30 @@ namespace BalUtilTests
     public ref class BAFunctions
     {
     public:
-        [Fact]
+        [Fact(Skip = "Need a mock implementation of IBootstrapperEngine to test BAFunctions.")]
         void CanCreateTestBAFunctions()
         {
             HRESULT hr = S_OK;
-            BOOTSTRAPPER_CREATE_ARGS bootstrapperArgs = { };
-            BOOTSTRAPPER_COMMAND bootstrapperCommand = { };
             BA_FUNCTIONS_CREATE_ARGS args = { };
             BA_FUNCTIONS_CREATE_RESULTS results = { };
             IBootstrapperEngine* pEngine = NULL;
+            BOOTSTRAPPER_COMMAND command = { };
             IBAFunctions* pBAFunctions = NULL;
 
-            bootstrapperArgs.cbSize = sizeof(bootstrapperArgs);
-            bootstrapperArgs.pCommand = &bootstrapperCommand;
-
             args.cbSize = sizeof(args);
-            args.pBootstrapperCreateArgs = &bootstrapperArgs;
+            args.pEngine = pEngine;
+            args.pCommand = &command;
 
             results.cbSize = sizeof(results);
 
             try
             {
-                hr = BalInitializeFromCreateArgs(&bootstrapperArgs, &pEngine);
-                NativeAssert::Succeeded(hr, "Failed to create engine.");
+                BalInitialize(pEngine);
 
-                hr = CreateBAFunctions(NULL, pEngine, &args, &results, &pBAFunctions);
+                hr = CreateBAFunctions(NULL, &args, &results);
                 NativeAssert::Succeeded(hr, "Failed to create BAFunctions.");
+
+                pBAFunctions = reinterpret_cast<IBAFunctions*>(results.pvBAFunctionsProcContext);
             }
             finally
             {
diff --git a/src/api/burn/test/BalUtilUnitTest/BootstrapperApplicationTests.cpp b/src/api/burn/test/BalUtilUnitTest/BootstrapperApplicationTests.cpp
index 70b01a42..9a68c738 100644
--- a/src/api/burn/test/BalUtilUnitTest/BootstrapperApplicationTests.cpp
+++ b/src/api/burn/test/BalUtilUnitTest/BootstrapperApplicationTests.cpp
@@ -12,28 +12,21 @@ namespace BalUtilTests
     public ref class BootstrapperApplication
     {
     public:
-        [Fact]
+        [Fact(Skip = "Need a mock implementation of IBootstrapperEngine to test BootstrapperApplication.")]
         void CanCreateTestBootstrapperApplication()
         {
             HRESULT hr = S_OK;
-            BOOTSTRAPPER_CREATE_ARGS args = { };
-            BOOTSTRAPPER_COMMAND command = { };
-            BOOTSTRAPPER_CREATE_RESULTS results = { };
-            IBootstrapperEngine* pEngine = NULL;
             IBootstrapperApplication* pApplication = NULL;
-
-            args.cbSize = sizeof(args);
-            args.pCommand = &command;
-
-            results.cbSize = sizeof(results);
+            IBootstrapperEngine* pEngine = NULL;
+            BOOTSTRAPPER_COMMAND command = { };
 
             try
             {
-                hr = BalInitializeFromCreateArgs(&args, &pEngine);
-                NativeAssert::Succeeded(hr, "Failed to create engine.");
-
-                hr = CreateBootstrapperApplication(pEngine, &args, &results, &pApplication);
+                hr = CreateBootstrapperApplication(&pApplication);
                 NativeAssert::Succeeded(hr, "Failed to create BootstrapperApplication.");
+
+                hr = pApplication->OnCreate(pEngine, &command);
+                NativeAssert::Succeeded(hr, "Failed to initialize BootstrapperApplication.");
             }
             finally
             {
diff --git a/src/api/burn/test/BalUtilUnitTest/TestBAFunctions.cpp b/src/api/burn/test/BalUtilUnitTest/TestBAFunctions.cpp
index 927a8d10..3e850442 100644
--- a/src/api/burn/test/BalUtilUnitTest/TestBAFunctions.cpp
+++ b/src/api/burn/test/BalUtilUnitTest/TestBAFunctions.cpp
@@ -8,34 +8,32 @@ class CTestBAFunctions : public CBalBaseBAFunctions
 {
 public:
     CTestBAFunctions(
-        __in HMODULE hModule,
-        __in IBootstrapperEngine* pEngine,
-        __in const BA_FUNCTIONS_CREATE_ARGS* pArgs
-        ) : CBalBaseBAFunctions(hModule, pEngine, pArgs)
+        __in HMODULE hModule
+        ) : CBalBaseBAFunctions(hModule)
     {
     }
 };
 
 HRESULT CreateBAFunctions(
     __in HMODULE hModule,
-    __in IBootstrapperEngine* pEngine,
     __in const BA_FUNCTIONS_CREATE_ARGS* pArgs,
-    __in BA_FUNCTIONS_CREATE_RESULTS* pResults,
-    __out IBAFunctions** ppApplication
+    __inout BA_FUNCTIONS_CREATE_RESULTS* pResults
     )
 {
     HRESULT hr = S_OK;
-    CTestBAFunctions* pApplication = NULL;
+    CTestBAFunctions* pFunction = NULL;
 
-    pApplication = new CTestBAFunctions(hModule, pEngine, pArgs);
-    ExitOnNull(pApplication, hr, E_OUTOFMEMORY, "Failed to create new test bafunctions object.");
+    pFunction = new CTestBAFunctions(hModule);
+    ExitOnNull(pFunction, hr, E_OUTOFMEMORY, "Failed to create new test bafunctions object.");
+
+    hr = pFunction->OnCreate(pArgs->pEngine, pArgs->pCommand);
+    ExitOnFailure(hr, "Failed to initialize new test bafunctions.");
 
     pResults->pfnBAFunctionsProc = BalBaseBAFunctionsProc;
-    pResults->pvBAFunctionsProcContext = pApplication;
-    *ppApplication = pApplication;
-    pApplication = NULL;
+    pResults->pvBAFunctionsProcContext = pFunction;
+    pFunction = NULL;
 
 LExit:
-    ReleaseObject(pApplication);
+    ReleaseObject(pFunction);
     return hr;
 }
diff --git a/src/api/burn/test/BalUtilUnitTest/TestBAFunctions.h b/src/api/burn/test/BalUtilUnitTest/TestBAFunctions.h
index 41f36df8..e25e40c3 100644
--- a/src/api/burn/test/BalUtilUnitTest/TestBAFunctions.h
+++ b/src/api/burn/test/BalUtilUnitTest/TestBAFunctions.h
@@ -3,8 +3,6 @@
 
 HRESULT CreateBAFunctions(
     __in HMODULE hModule,
-    __in IBootstrapperEngine* pEngine,
     __in const BA_FUNCTIONS_CREATE_ARGS* pArgs,
-    __in BA_FUNCTIONS_CREATE_RESULTS* pResults,
-    __out IBAFunctions** ppApplication
+    __inout BA_FUNCTIONS_CREATE_RESULTS* pResults
     );
diff --git a/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.cpp b/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.cpp
index daa1d690..b345ab9f 100644
--- a/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.cpp
+++ b/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.cpp
@@ -2,36 +2,25 @@
 
 #include "precomp.h"
 #include "BalBaseBootstrapperApplication.h"
-#include "BalBaseBootstrapperApplicationProc.h"
 
 class CTestBootstrapperApplication : public CBalBaseBootstrapperApplication
 {
 public:
-    CTestBootstrapperApplication(
-        __in IBootstrapperEngine* pEngine
-        ) : CBalBaseBootstrapperApplication(pEngine)
+    CTestBootstrapperApplication() : CBalBaseBootstrapperApplication()
     {
     }
 };
 
 HRESULT CreateBootstrapperApplication(
-    __in IBootstrapperEngine* pEngine,
-    __in const BOOTSTRAPPER_CREATE_ARGS* pArgs,
-    __inout BOOTSTRAPPER_CREATE_RESULTS* pResults,
     __out IBootstrapperApplication** ppApplication
     )
 {
     HRESULT hr = S_OK;
     CTestBootstrapperApplication* pApplication = NULL;
 
-    pApplication = new CTestBootstrapperApplication(pEngine);
+    pApplication = new CTestBootstrapperApplication();
     ExitOnNull(pApplication, hr, E_OUTOFMEMORY, "Failed to create new test bootstrapper application object.");
 
-    hr = pApplication->Initialize(pArgs);
-    ExitOnFailure(hr, "CTestBootstrapperApplication initialization failed.");
-
-    pResults->pfnBootstrapperApplicationProc = BalBaseBootstrapperApplicationProc;
-    pResults->pvBootstrapperApplicationProcContext = pApplication;
     *ppApplication = pApplication;
     pApplication = NULL;
 
diff --git a/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.h b/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.h
index c173e9ee..313bfede 100644
--- a/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.h
+++ b/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.h
@@ -2,8 +2,5 @@
 // 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.
 
 HRESULT CreateBootstrapperApplication(
-    __in IBootstrapperEngine* pEngine,
-    __in const BOOTSTRAPPER_CREATE_ARGS* pArgs,
-    __inout BOOTSTRAPPER_CREATE_RESULTS* pResults,
     __out IBootstrapperApplication** ppApplication
     );
diff --git a/src/api/burn/test/BalUtilUnitTest/precomp.h b/src/api/burn/test/BalUtilUnitTest/precomp.h
index 218cab68..57487bc5 100644
--- a/src/api/burn/test/BalUtilUnitTest/precomp.h
+++ b/src/api/burn/test/BalUtilUnitTest/precomp.h
@@ -15,8 +15,8 @@
 #include <dutil.h>
 #include <dictutil.h>
 
-#include <BootstrapperEngine.h>
-#include <BootstrapperApplication.h>
+#include <baenginetypes.h>
+#include <batypes.h>
 
 #include <BAFunctions.h>
 #include <IBootstrapperEngine.h>