1 files changed, 25 insertions, 11 deletions
diff --git a/testing/multiprocess_func_list.cc b/testing/multiprocess_func_list.cc
index 1e19704..49ae07d 100644
--- a/testing/multiprocess_func_list.cc
+++ b/testing/multiprocess_func_list.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2008 The Chromium Authors. All rights reserved.
+// Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
@@ -12,29 +12,43 @@ namespace multi_process_function_list {
 
 namespace {
 
-typedef std::map<std::string, ChildFunctionPtr> MultiProcessTestMap;
+struct ProcessFunctions {
+  ProcessFunctions() : main(NULL), setup(NULL) {}
+  ProcessFunctions(TestMainFunctionPtr main, SetupFunctionPtr setup)
+      : main(main),
+        setup(setup) {
+  }
+  TestMainFunctionPtr main;
+  SetupFunctionPtr setup;
+};
+
+typedef std::map<std::string, ProcessFunctions> MultiProcessTestMap;
 
 // Retrieve a reference to the global 'func name' -> func ptr map.
-MultiProcessTestMap &GetMultiprocessFuncMap() {
+MultiProcessTestMap& GetMultiprocessFuncMap() {
   static MultiProcessTestMap test_name_to_func_ptr_map;
   return test_name_to_func_ptr_map;
 }
 
 }  // namespace
 
-AppendMultiProcessTest::AppendMultiProcessTest(std::string test_name,
-                                               ChildFunctionPtr func_ptr) {
-  GetMultiprocessFuncMap()[test_name] = func_ptr;
+AppendMultiProcessTest::AppendMultiProcessTest(
+    std::string test_name,
+    TestMainFunctionPtr main_func_ptr,
+    SetupFunctionPtr setup_func_ptr) {
+  GetMultiprocessFuncMap()[test_name] =
+      ProcessFunctions(main_func_ptr, setup_func_ptr);
 }
 
 int InvokeChildProcessTest(std::string test_name) {
-  MultiProcessTestMap &func_lookup_table = GetMultiprocessFuncMap();
+  MultiProcessTestMap& func_lookup_table = GetMultiprocessFuncMap();
   MultiProcessTestMap::iterator it = func_lookup_table.find(test_name);
   if (it != func_lookup_table.end()) {
-    ChildFunctionPtr func = it->second;
-    if (func) {
-      return (*func)();
-    }
+    const ProcessFunctions& process_functions = it->second;
+    if (process_functions.setup)
+      (*process_functions.setup)();
+    if (process_functions.main)
+      return (*process_functions.main)();
   }
 
   return -1;