scripts: clean up enumeration type am: a9dc1ef0f6

Original change: https://android-review.googlesource.com/c/trusty/vendor/google/aosp/+/2909321

Change-Id: I5d0550ad2e1d3e2aead0dca4b5ee20510cbd65f3
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

2 files changed, 12 insertions, 16 deletions

diff --git a/scripts/run_tests.py b/scripts/run_tests.py
index 332949a..f88caeb 100755
--- a/scripts/run_tests.py
+++ b/scripts/run_tests.py
@@ -38,7 +38,7 @@ from enum import Enum
 from typing import Optional
 from collections import namedtuple
 
-from trusty_build_config import PORT_TYPE, TrustyTest, TrustyCompositeTest
+from trusty_build_config import PortType, TrustyTest, TrustyCompositeTest
 from trusty_build_config import TrustyRebootCommand, TrustyHostTest
 from trusty_build_config import TrustyAndroidTest, TrustyBuildConfig
 
@@ -339,7 +339,7 @@ def run_tests(
             case TrustyTest():
                 # Benchmark runs on QEMU are meaningless and take a lot of CI time
                 # One can still run the bootport test manually if desired
-                if test.port_type == PORT_TYPE.BENCHMARK:
+                if test.port_type == PortType.BENCHMARK:
                     ignore_tests = True
                 else:
                     if isinstance(test, TrustyAndroidTest):
diff --git a/scripts/trusty_build_config.py b/scripts/trusty_build_config.py
index 226210a..2432e36 100755
--- a/scripts/trusty_build_config.py
+++ b/scripts/trusty_build_config.py
@@ -24,16 +24,14 @@
 import argparse
 import os
 import re
-from enum import Enum
+from enum import StrEnum, auto
 from typing import List, Dict, Optional
 
 script_dir = os.path.dirname(os.path.abspath(__file__))
 
-class PORT_TYPE(Enum):
-    TEST = 1
-    BENCHMARK = 2
-    def __str__(self):
-        return str(self.name)
+class PortType(StrEnum):
+    TEST = auto(),
+    BENCHMARK = auto(),
 
 
 class TrustyBuildConfigProject(object):
@@ -88,16 +86,14 @@ class TrustyArchiveBuildFile(object):
 
 class TrustyTest(object):
     """Stores a pair of a test name and a command to run"""
-    def __init__(self, name, command, enabled, port_type = PORT_TYPE.TEST):
+    def __init__(self, name, command, enabled, port_type = PortType.TEST):
         self.name = name
         self.command = command
         self.enabled = enabled
         self.port_type = port_type
 
     def type(self, port_type):
-        if not isinstance(port_type, PORT_TYPE):
-            raise TypeError("Unexpected port_type: " + str(port_type))
-        self.port_type = port_type
+        self.port_type = PortType(port_type)  # ensure we have a valid port type
         return self
 
 class TrustyHostTest(TrustyTest):
@@ -118,7 +114,7 @@ class TrustyHostTest(TrustyTest):
 class TrustyAndroidTest(TrustyTest):
     """Stores a test name and command to run inside Android"""
 
-    def __init__(self, name, command, need = None, port_type = PORT_TYPE.TEST, enabled=True, nameprefix="", runargs=(),
+    def __init__(self, name, command, need = None, port_type = PortType.TEST, enabled=True, nameprefix="", runargs=(),
                  timeout=None):
         nameprefix = nameprefix + "android-test:"
         cmd = ["run", "--headless", "--shell-command", command]
@@ -134,7 +130,7 @@ class TrustyAndroidTest(TrustyTest):
 class TrustyPortTest(TrustyTest):
     """Stores a trusty port name for a test to run."""
 
-    def __init__(self, port, port_type=PORT_TYPE.TEST, enabled=True, timeout=None):
+    def __init__(self, port, port_type=PortType.TEST, enabled=True, timeout=None):
         super().__init__(port, None, enabled, port_type)
         self.port = port
         self.need = TrustyPortTestFlags()
@@ -367,8 +363,8 @@ class TrustyBuildConfig(object):
 
 
         file_format = {
-            "BENCHMARK": PORT_TYPE.BENCHMARK,
-            "TEST": PORT_TYPE.TEST,
+            "BENCHMARK": PortType.BENCHMARK,
+            "TEST": PortType.TEST,
             "include": include,
             "build": build,
             "builddep": builddep,