Raises error when unsupported mode is passed

Author: ffmmjj

piranha_python/codemods.py

@@ -46,6 +46,9 @@ def add_args(arg_parser):  # pragma: no cover
 
     def __init__(self, context, flag_name, flag_resolution_methods, ignored_module_check_fn_path=None, mode="treated"):
         super().__init__(context)
+        if mode not in ("treated", "control"):
+            raise ValueError("mode parameter must be 'treated' or 'control' - '%s' was passed" % mode)
+
         self.flag_name = flag_name
         self.is_in_feature_flag_block = False
         self.found_return_stmt_in_ff_block = False

test/test_codemods.py

@@ -19,6 +19,18 @@ def test_lack_of_custom_method_name_raises_exception(self):
 
         self.assertIn("flag_resolution_methods", thrown_exception.exception.args[0])
 
+    def test_unsupported_mode_name_raises_exception(self):
+        with self.assertRaises(ValueError) as thrown_exception:
+            self.assertCodemod(
+                "print('This is not related to the feature flag value at all')",
+                "print('This is not related to the feature flag value at all')",
+                flag_name=FEATURE_FLAG_NAME,
+                flag_resolution_methods="is_flag_active",
+                mode="not-a-supported-mode",
+            )
+
+        self.assertIn("mode", thrown_exception.exception.args[0])
+
 
 class PiranhaTreatmentFlagTest(CodemodTest):
     TRANSFORM = PiranhaCommand