diff --git a/python/test.go b/python/test.go
index 7eb9136..826f353 100644
--- a/python/test.go
+++ b/python/test.go
@@ -158,35 +158,25 @@
 	}
 
 	runner := proptools.StringDefault(p.testProperties.Test_options.Runner, "tradefed")
-	if runner == "tradefed" {
-		p.testConfig = tradefed.AutoGenTestConfig(ctx, tradefed.AutoGenTestConfigOptions{
-			TestConfigProp:          p.testProperties.Test_config,
-			TestConfigTemplateProp:  p.testProperties.Test_config_template,
-			TestSuites:              p.binaryProperties.Test_suites,
-			OptionsForAutogenerated: configs,
-			AutoGenConfig:           p.binaryProperties.Auto_gen_config,
-			DeviceTemplate:          "${PythonBinaryHostTestConfigTemplate}",
-			HostTemplate:            "${PythonBinaryHostTestConfigTemplate}",
-		})
-	} else if runner == "mobly" {
-		if p.testProperties.Test_config != nil || p.testProperties.Test_config_template != nil || p.binaryProperties.Auto_gen_config != nil {
-			panic(fmt.Errorf("cannot set test_config, test_config_template or auto_gen_config for mobly test"))
+	template := "${PythonBinaryHostTestConfigTemplate}"
+	if runner == "mobly" {
+		// Add tag to enable Atest mobly runner
+		if !android.InList("mobly", p.testProperties.Test_options.Tags) {
+			p.testProperties.Test_options.Tags = append(p.testProperties.Test_options.Tags, "mobly")
 		}
-
-		for _, testSuite := range p.binaryProperties.Test_suites {
-			if testSuite == "cts" {
-				configs = append(configs, tradefed.Option{Name: "test-suite-tag", Value: "cts"})
-				break
-			}
-		}
-		p.testConfig = tradefed.AutoGenTestConfig(ctx, tradefed.AutoGenTestConfigOptions{
-			OptionsForAutogenerated: configs,
-			DeviceTemplate:          "${PythonBinaryHostMoblyTestConfigTemplate}",
-			HostTemplate:            "${PythonBinaryHostMoblyTestConfigTemplate}",
-		})
-	} else {
+		template = "${PythonBinaryHostMoblyTestConfigTemplate}"
+	} else if runner != "tradefed" {
 		panic(fmt.Errorf("unknown python test runner '%s', should be 'tradefed' or 'mobly'", runner))
 	}
+	p.testConfig = tradefed.AutoGenTestConfig(ctx, tradefed.AutoGenTestConfigOptions{
+		TestConfigProp:          p.testProperties.Test_config,
+		TestConfigTemplateProp:  p.testProperties.Test_config_template,
+		TestSuites:              p.binaryProperties.Test_suites,
+		OptionsForAutogenerated: configs,
+		AutoGenConfig:           p.binaryProperties.Auto_gen_config,
+		DeviceTemplate:          template,
+		HostTemplate:            template,
+	})
 
 	for _, dataSrcPath := range android.PathsForModuleSrc(ctx, p.testProperties.Data) {
 		p.data = append(p.data, android.DataPath{SrcPath: dataSrcPath})
@@ -228,6 +218,12 @@
 				entries.SetString("LOCAL_FULL_TEST_CONFIG", p.testConfig.String())
 			}
 
+			// ATS 2.0 is the test harness for mobly tests and the test config is for ATS 2.0.
+			// Add "v2" suffix to test config name to distinguish it from the config for TF.
+			if proptools.String(p.testProperties.Test_options.Runner) == "mobly" {
+				entries.SetString("LOCAL_TEST_CONFIG_SUFFIX", "v2")
+			}
+
 			entries.SetBoolIfTrue("LOCAL_DISABLE_AUTO_GENERATE_TEST_CONFIG", !BoolDefault(p.binaryProperties.Auto_gen_config, true))
 			android.SetAconfigFileMkEntries(&p.ModuleBase, entries, p.mergedAconfigFiles)