Tools: autotest: use class map for getting test class for full test name

Tools/autotest/autotest.py

@@ -295,22 +295,23 @@ def find_specific_test_to_run(step):
     (testname, test) = t
     return "%s.%s" % (testname, test)
 
+tester_class_map = {
+    "fly.ArduCopter": arducopter.AutoTestCopter,
+    "fly.ArduPlane": arduplane.AutoTestPlane,
+    "fly.QuadPlane": quadplane.AutoTestQuadPlane,
+    "drive.APMrover2": apmrover2.AutoTestRover,
+    "drive.balancebot": balancebot.AutoTestBalanceBot,
+    "fly.CopterAVC": arducopter.AutoTestHeli,
+    "dive.ArduSub": ardusub.AutoTestSub,
+    "test.AntennaTracker": antennatracker.AutoTestTracker,
+}
+
 def run_specific_test(step, *args, **kwargs):
     t = split_specific_test_step(step)
     if t is None:
         return []
     (testname, test) = t
 
-    tester_class_map = {
-        "fly.ArduCopter": arducopter.AutoTestCopter,
-        "fly.ArduPlane": arduplane.AutoTestPlane,
-        "fly.QuadPlane": quadplane.AutoTestQuadPlane,
-        "drive.APMrover2": apmrover2.AutoTestRover,
-        "drive.balancebot": balancebot.AutoTestBalanceBot,
-        "fly.CopterAVC": arducopter.AutoTestHeli,
-        "dive.ArduSub": ardusub.AutoTestSub,
-        "test.AntennaTracker": antennatracker.AutoTestTracker,
-    }
     tester_class = tester_class_map[testname]
     tester = tester_class(*args, **kwargs)
 
@@ -380,37 +381,9 @@ def run_step(step):
     if opts.speedup is not None:
         fly_opts["speedup"] = opts.speedup
 
-    if step == 'fly.ArduCopter':
-        tester = arducopter.AutoTestCopter(binary, **fly_opts)
-        return tester.autotest()
-
-    if step == 'fly.CopterAVC':
-        tester = arducopter.AutoTestHeli(binary, **fly_opts)
-        return tester.autotest()
-
-    if step == 'fly.ArduPlane':
-        tester = arduplane.AutoTestPlane(binary, **fly_opts)
-        return tester.autotest()
-
-    if step == 'fly.QuadPlane':
-        tester = quadplane.AutoTestQuadPlane(binary, **fly_opts)
-        return tester.autotest()
-
-    if step == 'drive.APMrover2':
-        tester = apmrover2.AutoTestRover(binary, **fly_opts)
-        return tester.autotest()
-
-    if step == 'drive.balancebot':
-        tester = balancebot.AutoTestBalanceBot(binary, **fly_opts)
-        return tester.autotest()
-
-    if step == 'dive.ArduSub':
-        tester = ardusub.AutoTestSub(binary, **fly_opts)
-        return tester.autotest()
-
-    if step == 'test.AntennaTracker':
-        tester = antennatracker.AutoTestTracker(binary, **fly_opts)
-        return tester.autotest()
+    # handle "fly.ArduCopter" etc:
+    if step in tester_class_map:
+        return tester_class_map[step](binary, **fly_opts).autotest()
 
     specific_test_to_run = find_specific_test_to_run(step)
     if specific_test_to_run is not None: