From a68265bb1350837db05334064fb8a26eda6e35cd Mon Sep 17 00:00:00 2001
From: Marshall Hallenbeck <marshall.hallenbeck@gmail.com>
Date: Fri, 20 Oct 2023 15:02:37 -0400
Subject: [PATCH] feat(tests): allow running specific proto tests via
 --protocols

---
 tests/e2e_test.py | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/tests/e2e_test.py b/tests/e2e_test.py
index f7aabda6..81c44c7a 100644
--- a/tests/e2e_test.py
+++ b/tests/e2e_test.py
@@ -53,15 +53,19 @@ def get_cli_args():
         required=False,
         help="Use poetry to run commands",
     )
+    parser.add_argument(
+        "--protocols",
+        nargs="+",
+        default=[],
+        required=False,
+        help="Protocols to test",
+    )
 
     return parser.parse_args()
 
 
 def generate_commands(args):
     lines = []
-
-    kerberos = "-k " if args.kerberos else ""
-
     file_loc = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
     commands_file = os.path.join(file_loc, "e2e_commands.txt")
 
@@ -70,12 +74,21 @@ def generate_commands(args):
             if line.startswith("#"):
                 continue
             line = line.strip()
-            line = line.replace("TARGET_HOST", args.target).replace("LOGIN_USERNAME", f'"{args.username}"').replace("LOGIN_PASSWORD", f'"{args.password}"').replace("KERBEROS ", kerberos)
-            if args.poetry:
-                line = f"poetry run {line}"
-            lines.append(line)
+            if args.protocols:
+                if line.split()[1] in args.protocols:
+                    lines.append(replace_command(args, line))
+            else:
+                lines.append(replace_command(args, line))
     return lines
 
+def replace_command(args, line):
+    kerberos = "-k " if args.kerberos else ""
+
+    line = line.replace("TARGET_HOST", args.target).replace("LOGIN_USERNAME", f'"{args.username}"').replace("LOGIN_PASSWORD", f'"{args.password}"').replace("KERBEROS ", kerberos)
+    if args.poetry:
+        line = f"poetry run {line}"
+    return line
+
 
 def run_e2e_tests(args):
     console = Console()