bisect-kit: clean up experiment version

BUG=b:328452015
TEST=./diagnose_cros_tast.py {init,run} && ./run_formatter.sh

Change-Id: I288d3fd014777794966a33f6ac49bc1631576645
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/bisect-kit/+/5348017
Reviewed-by: MD Nayeem Jahan Rafi <njrafi@chromium.org>
Commit-Queue: Cheng-Yu Lee <cylee@google.com>
Tested-by: Cheng-Yu Lee <cylee@google.com>
diff --git a/experiment.py b/experiment.py
index f284a00..05a777e 100644
--- a/experiment.py
+++ b/experiment.py
@@ -13,7 +13,6 @@
 """
 
 import argparse
-import dataclasses
 from enum import StrEnum
 import logging
 
@@ -28,86 +27,29 @@
     VM = 'vm'
 
 
+# TODO b/328452015: either record the version when an experimental bisection
+# runs, or remove it if the version doesn't matter.
 _VERSION_MAP = {
     ID.STATELESS: '20230901',
     ID.VM: '20240207',
 }
 
 
-@dataclasses.dataclass
-class Experiment:
-    """Defines an experiment.
-
-    name (ID) is mandatory while version is optional
-    """
-
-    name: ID
-    version: str | None = None
-
-
-def experiment_type(s: str):
-    """Parse experiments from the command line.
-
-    The command line expects the value of --experiments to be a list of
-      - <exp_name> or
-      - <exp_name:version>
-    """
-    name = None
-    version = None
-
-    parts = s.split(':')
-    if len(parts) == 1:
-        (name,) = parts
-    elif len(parts) == 2:
-        name, version = parts
-    else:
-        raise argparse.ArgumentTypeError('Invalid experiment format %s' % s)
-
-    try:
-        name = ID(name)
-    except ValueError as e:
-        raise argparse.ArgumentTypeError(e)
-
-    return Experiment(name, version)
-
-
-def is_in_experiment(experiments: list[Experiment], exp_in_query: ID) -> bool:
+def is_in_experiment(experiments: list[ID], exp_in_query: ID) -> bool:
     """Returns whether a bisect is in a specific exepriment.
 
-    For each experiment this bisect is in:
-    If version is missing, we don't check experiment versions.
-    If version is given,  we check whether current status of the experiment
-    matches the specified version.
-
-    For example, if the bisect is run with "--experiments exp1 exp2:beta exp3:alpha"
-    and current experiments version is
-    {
-      exp1: "dev",
-      exp2: "beta",
-      exp3: "beta",
-    }
-    Then the bisect would be run with "exp1", "exp2" but not "exp3".
-
     Args:
-      experiments: A list of Experiment object.
+      experiments: A list of experiment IDs.
       exp_in_query: The Experiment ID in query.
     """
-    for exp in experiments:
-        if exp.name == exp_in_query:
-            # If the version is not specified, current version doesn't matter.
-            # If the version is specified, the experiment must specify its
-            # current version in _VERSION_MAP and the two must match.
-            if not exp.version or exp.version == _VERSION_MAP.get(exp.name):
-                return True
-
-    return False
+    return exp_in_query in experiments
 
 
 def common_flags() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(add_help=False)
     parser.add_argument(
         '--experiments',
-        type=experiment_type,
+        type=str,
         action="extend",
         nargs="+",
         metavar='EXPERIMENTS',
diff --git a/experiment_test.py b/experiment_test.py
index 1df95dc..86722c2 100644
--- a/experiment_test.py
+++ b/experiment_test.py
@@ -4,7 +4,6 @@
 
 """Test experiment module."""
 
-import argparse
 from enum import StrEnum
 import unittest
 from unittest import mock
@@ -29,43 +28,6 @@
 }
 
 
-class ParseExperimentTest(unittest.TestCase):
-    """Test parsing Experiment from string."""
-
-    def setUp(self):
-        super().setUp()
-        self._mock_id_class = self.enterContext(
-            mock.patch.object(experiment, 'ID', FakeID)
-        )
-
-        self._mock_version_map = self.enterContext(
-            mock.patch.object(experiment, '_VERSION_MAP', FAKE_VERSION_MAP)
-        )
-
-    def test_parse_exp_name(self):
-        self.assertEqual(
-            experiment.experiment_type('exp1'),
-            experiment.Experiment(experiment.ID.EXP1),
-        )
-
-    def test_parse_exp_name_and_version(self):
-        self.assertEqual(
-            experiment.experiment_type('exp1:dev'),
-            experiment.Experiment(experiment.ID.EXP1, 'dev'),
-        )
-
-    def test_parse_exp_format_error(self):
-        with self.assertRaises(argparse.ArgumentTypeError):
-            experiment.experiment_type('a:b:c')
-
-    def test_parse_exp_invalid_id(self):
-        with self.assertRaises(argparse.ArgumentTypeError):
-            experiment.experiment_type('exp0')
-
-        with self.assertRaises(argparse.ArgumentTypeError):
-            experiment.experiment_type('exp0:dev')
-
-
 class IsInExperimentTest(unittest.TestCase):
     """Test is_in_experiment."""
 
@@ -82,50 +44,15 @@
     def test_experiments_empty_list(self):
         self.assertFalse(experiment.is_in_experiment([], experiment.ID.EXP1))
 
-    def test_experiments_single_experiment_without_version(self):
+    def test_experiments_single_experiment(self):
         self.assertTrue(
             experiment.is_in_experiment(
-                [experiment.Experiment(experiment.ID.EXP1)], experiment.ID.EXP1
+                [experiment.ID.EXP1], experiment.ID.EXP1
             )
         )
         self.assertFalse(
             experiment.is_in_experiment(
-                [experiment.Experiment(experiment.ID.EXP1)], experiment.ID.EXP2
-            )
-        )
-
-        # EXP4 has no version defined.
-        self.assertTrue(
-            experiment.is_in_experiment(
-                [experiment.Experiment(experiment.ID.EXP4)], experiment.ID.EXP4
-            )
-        )
-
-    def test_experiments_single_experiment_with_version(self):
-        self.assertTrue(
-            experiment.is_in_experiment(
-                [experiment.Experiment(experiment.ID.EXP1, 'dev')],
-                experiment.ID.EXP1,
-            )
-        )
-        self.assertFalse(
-            experiment.is_in_experiment(
-                [experiment.Experiment(experiment.ID.EXP1, 'alpha')],
-                experiment.ID.EXP1,
-            )
-        )
-        self.assertFalse(
-            experiment.is_in_experiment(
-                [experiment.Experiment(experiment.ID.EXP1, 'dev')],
-                experiment.ID.EXP2,
-            )
-        )
-
-        # EXP4 has no version defined.
-        self.assertFalse(
-            experiment.is_in_experiment(
-                [experiment.Experiment(experiment.ID.EXP4, 'dev')],
-                experiment.ID.EXP4,
+                [experiment.ID.EXP1], experiment.ID.EXP2
             )
         )
 
@@ -133,21 +60,10 @@
         self.assertTrue(
             experiment.is_in_experiment(
                 [
-                    experiment.Experiment(experiment.ID.EXP1, 'dev'),
-                    experiment.Experiment(experiment.ID.EXP2),
-                    experiment.Experiment(experiment.ID.EXP3, 'alpha'),
-                    experiment.Experiment(experiment.ID.EXP4),
-                ],
-                experiment.ID.EXP1,
-            )
-        )
-        self.assertTrue(
-            experiment.is_in_experiment(
-                [
-                    experiment.Experiment(experiment.ID.EXP1, 'dev'),
-                    experiment.Experiment(experiment.ID.EXP2),
-                    experiment.Experiment(experiment.ID.EXP3, 'alpha'),
-                    experiment.Experiment(experiment.ID.EXP4),
+                    experiment.ID.EXP1,
+                    experiment.ID.EXP2,
+                    experiment.ID.EXP3,
+                    experiment.ID.EXP4,
                 ],
                 experiment.ID.EXP2,
             )
@@ -155,34 +71,9 @@
         self.assertFalse(
             experiment.is_in_experiment(
                 [
-                    experiment.Experiment(experiment.ID.EXP1, 'dev'),
-                    experiment.Experiment(experiment.ID.EXP2),
-                    experiment.Experiment(experiment.ID.EXP3, 'alpha'),
-                    experiment.Experiment(experiment.ID.EXP4),
-                ],
-                experiment.ID.EXP3,
-            )
-        )
-        self.assertTrue(
-            experiment.is_in_experiment(
-                [
-                    experiment.Experiment(experiment.ID.EXP1, 'dev'),
-                    experiment.Experiment(experiment.ID.EXP2),
-                    experiment.Experiment(experiment.ID.EXP3, 'alpha'),
-                    experiment.Experiment(experiment.ID.EXP4),
-                ],
-                experiment.ID.EXP4,
-            )
-        )
-
-        # EXP4 has no version defined.
-        self.assertFalse(
-            experiment.is_in_experiment(
-                [
-                    experiment.Experiment(experiment.ID.EXP1, 'dev'),
-                    experiment.Experiment(experiment.ID.EXP2),
-                    experiment.Experiment(experiment.ID.EXP3, 'alpha'),
-                    experiment.Experiment(experiment.ID.EXP4, 'beta'),
+                    experiment.ID.EXP1,
+                    experiment.ID.EXP2,
+                    experiment.ID.EXP3,
                 ],
                 experiment.ID.EXP4,
             )