Merge "Fix built-in controller registration flow."
diff --git a/acts/framework/acts/test_runner.py b/acts/framework/acts/test_runner.py
index a71a9a3..228dae7 100644
--- a/acts/framework/acts/test_runner.py
+++ b/acts/framework/acts/test_runner.py
@@ -139,6 +139,28 @@
                             test_classes[member_name] = test_class
         return test_classes
 
+    def _import_builtin_controllers(self):
+        """Import built-in controller modules.
+
+        Go through the testbed configs, find any built-in controller configs
+        and import the corresponding controller module from acts.controllers
+        package.
+
+        TODO(angli): Remove this when all scripts change to explicitly declare
+                     controller dependency.
+
+        Returns:
+            A list of controller modules.
+        """
+        builtin_controllers = []
+        for ctrl_name in keys.Config.builtin_controller_names.value:
+            if ctrl_name in self.testbed_configs:
+                module_name = keys.get_module_name(ctrl_name)
+                module = importlib.import_module("acts.controllers.%s" %
+                                                 module_name)
+                builtin_controllers.append(module)
+        return builtin_controllers
+
     @staticmethod
     def verify_controller_module(module):
         """Verifies a module object follows the required interface for
@@ -173,7 +195,7 @@
         module will be instantiated with corresponding configs in the test
         config file. The module should be imported first.
 
-        Params:
+        Args:
             module: A module that follows the controller module interface.
 
         Returns:
@@ -251,13 +273,6 @@
         """
         self.test_run_info[
             keys.Config.ikey_testbed_name.value] = self.testbed_name
-        # Instantiate builtin controllers
-        for ctrl_name in keys.Config.builtin_controller_names.value:
-            if ctrl_name in self.testbed_configs:
-                module_name = keys.get_module_name(ctrl_name)
-                module = importlib.import_module("acts.controllers.%s" %
-                                                 module_name)
-                self.register_controller(module)
         # Unpack other params.
         self.test_run_info["register_controller"] = self.register_controller
         self.test_run_info[keys.Config.ikey_logpath.value] = self.log_path
@@ -342,23 +357,26 @@
         t_configs = self.test_configs[keys.Config.key_test_paths.value]
         self.test_classes = self.import_test_modules(t_configs)
         self.log.debug("Executing run list %s.", self.run_list)
-        try:
-            for test_cls_name, test_case_names in self.run_list:
-                if not self.running:
-                    break
-                if test_case_names:
-                    self.log.debug("Executing test cases %s in test class %s.",
-                                   test_case_names, test_cls_name)
-                else:
-                    self.log.debug("Executing test class %s", test_cls_name)
-                try:
-                    self.run_test_class(test_cls_name, test_case_names)
-                except signals.TestAbortAll as e:
-                    self.log.warning(
-                        "Abort all subsequent test classes. Reason: %s", e)
-                    raise
-        finally:
-            self.unregister_controllers()
+        for test_cls_name, test_case_names in self.run_list:
+            if not self.running:
+                break
+            if test_case_names:
+                self.log.debug("Executing test cases %s in test class %s.",
+                               test_case_names, test_cls_name)
+            else:
+                self.log.debug("Executing test class %s", test_cls_name)
+            try:
+                # Import and register the built-in controller modules specified
+                # in testbed config.
+                for module in self._import_builtin_controllers():
+                    self.register_controller(module)
+                self.run_test_class(test_cls_name, test_case_names)
+            except signals.TestAbortAll as e:
+                self.log.warning(
+                    "Abort all subsequent test classes. Reason: %s", e)
+                raise
+            finally:
+                self.unregister_controllers()
 
     def stop(self):
         """Releases resources from test run. Should always be called after
diff --git a/acts/framework/tests/acts_test_runner_test.py b/acts/framework/tests/acts_test_runner_test.py
index f9c7b8e..f6376f3 100755
--- a/acts/framework/tests/acts_test_runner_test.py
+++ b/acts/framework/tests/acts_test_runner_test.py
@@ -15,6 +15,7 @@
 #   limitations under the License.
 
 
+import mock
 import shutil
 import tempfile
 import unittest
@@ -22,6 +23,8 @@
 from acts import keys
 from acts import signals
 from acts import test_runner
+
+import acts_android_device_test
 import mock_controller
 
 
@@ -139,6 +142,38 @@
         self.assertEqual(results["Executed"], 2)
         self.assertEqual(results["Passed"], 2)
 
+    @mock.patch('acts.controllers.adb.AdbProxy',
+                return_value=acts_android_device_test.MockAdbProxy(1))
+    @mock.patch('acts.controllers.android_device.list_adb_devices',
+                return_value=["1"])
+    @mock.patch('acts.controllers.android_device.get_all_instances',
+                return_value=acts_android_device_test.get_mock_ads(1))
+    def test_run_two_test_classes(self, mock_adb, mock_list_adb, mock_get_all):
+        """Verifies that runing more than one test class in one test run works
+        proerly.
+
+        This requires using a built-in controller module. Using AndroidDevice
+        module since it has all the mocks needed already.
+        """
+        mock_test_config = dict(self.base_mock_test_config)
+        tb_key = keys.Config.key_testbed.value
+        mock_ctrlr_config_name = mock_controller.ACTS_CONTROLLER_CONFIG_NAME
+        my_config = [{"serial": "xxxx", "magic": "Magic1"},
+                     {"serial": "xxxx", "magic": "Magic2"}]
+        mock_test_config[tb_key][mock_ctrlr_config_name] = my_config
+        mock_test_config[tb_key]["AndroidDevice"] = [
+            {"serial": "1", "skip_sl4a": True}]
+        tr = test_runner.TestRunner(mock_test_config,
+            [('IntegrationTest', None), ('IntegrationTest', None)])
+        tr.run()
+        tr.stop()
+        self.assertFalse(tr.controller_registry)
+        self.assertFalse(tr.controller_destructors)
+        results = tr.results.summary_dict()
+        self.assertEqual(results["Requested"], 2)
+        self.assertEqual(results["Executed"], 2)
+        self.assertEqual(results["Passed"], 2)
+
     def test_verify_controller_module(self):
         test_runner.TestRunner.verify_controller_module(mock_controller)