use enum for deprecation reasons

greenbone · Jul 4, 2024 · 139fb8e · 139fb8e
1 parent 42dceb0
commit 139fb8e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 16 deletions.
diff --git a/tests/standalone_plugins/test_deprecate_vts.py b/tests/standalone_plugins/test_deprecate_vts.py
@@ -22,7 +22,7 @@ class ParseArgsTestCase(unittest.TestCase):
     def test_parse_args(self):
         testfile = "testfile.nasl"
         output_path = "attic/"
-        reason = "notus"
+        reason = "NOTUS"
 
         args = parse_args(
             [
@@ -38,11 +38,11 @@ def test_parse_args(self):
         self.assertEqual(args.output_path, Path(output_path))
         self.assertEqual(args.deprecation_reason, reason)
 
-    def test_parse_args_invalid_reason(self):
+    def test_mandatory_arg_group_both(self):
         testfile = "testfile.nasl"
         output_path = "attic/"
         input_path = "nasl/common"
-        reason = "notus"
+        reason = "NOTUS"
 
         with self.assertRaises(SystemExit):
             parse_args(
@@ -58,7 +58,7 @@ def test_parse_args_invalid_reason(self):
                 ]
             )
 
-    def test_mandatory_arg_group_both(self):
+    def test_invalid_reason(self):
         output_path = "attic/"
         input_path = "nasl/common"
         reason = "foo"
@@ -76,7 +76,7 @@ def test_mandatory_arg_group_both(self):
 
     def test_mandatory_arg_group_neither(self):
         output_path = "attic/"
-        reason = "notus"
+        reason = "NOTUS"
         with self.assertRaises(SystemExit):
             parse_args(
                 [
@@ -124,7 +124,7 @@ def test_deprecate(self):
                     content=NASL_CONTENT,
                 )
             ]
-            deprecate(out_dir, to_deprecate, "notus")
+            deprecate(out_dir, to_deprecate, "NOTUS")
 
             result = testfile2.read_text(encoding="utf8")
             self.assertNotIn(result, "script_mandatory_keys")
@@ -147,7 +147,7 @@ def test_deprecate_kb_item(self):
                     content=NASL_CONTENT_KB,
                 )
             ]
-            deprecate(out_dir, to_deprecate, "notus")
+            deprecate(out_dir, to_deprecate, "NOTUS")
             self.assertLogs(
                 "Unable to deprecate testfile1.nasl. There are still KB keys "
                 "remaining."
@@ -177,7 +177,7 @@ def test_update_summary_no_oid_match(self):
             full_path=Path("dir/testfile.nasl"),
             content=NASL_CONTENT,
         )
-        result = update_summary(file, "notus")
+        result = update_summary(file, "NOTUS")
         self.assertIn("This VT has been deprecated", result)
 
     def test_get_files_from_path(self):

diff --git a/troubadix/standalone_plugins/deprecate_vts.py b/troubadix/standalone_plugins/deprecate_vts.py
@@ -4,6 +4,7 @@
 import re
 from argparse import ArgumentParser, Namespace
 from dataclasses import dataclass
+from enum import Enum
 from pathlib import Path
 from typing import Iterable, Optional
 
@@ -15,12 +16,12 @@
     SpecialScriptTag,
 )
 
-DEPRECATIONS = {
-    "notus": "and replaced by a Notus scanner based one.",
-    "merged": "because it has been merged into a different VT.",
-    "defunct": "and is therefore no longer functional.",
-    "duplicate": "as a duplicate.",
-}
+
+class Deprecations(Enum):
+    NOTUS = "and replaced by a Notus scanner based one."
+    MERGED = "because it has been merged into a different VT."
+    DEFUNCT = "and is therefore no longer functional."
+    DUPLICATE = "as a duplicate."
 
 
 @dataclass
@@ -52,7 +53,8 @@ def update_summary(file: DeprecatedFile, deprecation_reason: str) -> str:
         return file.content
 
     deprecate_text = (
-        f"Note: This VT has been deprecated {DEPRECATIONS[deprecation_reason]}"
+        f"Note: This VT has been deprecated "
+        f"{Deprecations[deprecation_reason].value}"
     )
 
     new_summary = old_summary + "\n" + deprecate_text
@@ -205,7 +207,7 @@ def parse_args(args: Iterable[str] = None) -> Namespace:
         "-r",
         "--deprecation-reason",
         metavar="<deprecation_reason>",
-        choices=DEPRECATIONS.keys(),
+        choices=[reason.name for reason in Deprecations],
         type=str,
         help="The reason the VT is being deprecated. Options are 'notus':"
         "The VT has been replaced by a new Notus VT. 'Merged': the VT has"