Fix timeout regex for spark submit command

paasta_tools/spark_tools.py

@@ -184,8 +184,13 @@ def auto_add_timeout_for_spark_job(
     if "spark-submit" not in cmd:
         return cmd
     try:
+        # This is not an exhaustive regex, matches the invalid ones also, where as the invalid
+        # timeout command will fail during execution
+        options_regex = r"(--?[a-z][a-z-]*((\s+|=)[\w\d-]+)?\s+)*"
+        duration_regex = r"\d+\.?\d*[smhd]?"
+
         timeout_present = re.match(
-            r"^.*timeout[\s]+[\d]+[\.]?[\d]*[m|h][\s]+spark-submit .*$", cmd
+            rf"^.*timeout\s+{options_regex}{duration_regex}\s+spark-submit .*$", cmd
         )
         if not timeout_present:
             split_cmd = cmd.split("spark-submit")

tests/test_spark_tools.py

@@ -4,6 +4,7 @@
 import pytest
 
 from paasta_tools import spark_tools
+from paasta_tools.spark_tools import auto_add_timeout_for_spark_job
 
 
 def test_get_webui_url():
@@ -133,3 +134,48 @@ def test_get_volumes_from_spark_k8s_configs(mock_sys, spark_conf, expected):
 def test_get_spark_driver_monitoring_annotations(spark_config, expected):
     result = spark_tools.get_spark_driver_monitoring_annotations(spark_config)
     assert result == expected
+
+
+@pytest.mark.parametrize(
+    argnames=[
+        "cmd",
+        "timeout_duration",
+        "expected",
+    ],
+    argvalues=[
+        pytest.param(
+            "spark-submit abc.py",
+            "4h",
+            "timeout 4h spark-submit abc.py",
+            id="No timeout",
+        ),
+        pytest.param(
+            "timeout 2h spark-submit abc.py",
+            "12h",
+            "timeout 2h spark-submit abc.py",
+            id="Timeout without options",
+        ),
+        pytest.param(
+            "timeout -v 2h spark-submit abc.py",
+            "12h",
+            "timeout -v 2h spark-submit abc.py",
+            id="Timeout with options",
+        ),
+        pytest.param(
+            "timeout -v -s 1 2h spark-submit abc.py",
+            "12h",
+            "timeout -v -s 1 2h spark-submit abc.py",
+            id="Timeout with multiple options",
+        ),
+        pytest.param(
+            "timeout -k 10m --signal=SIGKILL 2h spark-submit abc.py",
+            "12h",
+            "timeout -k 10m --signal=SIGKILL 2h spark-submit abc.py",
+            id="Timeout with double dash option",
+        ),
+    ],
+)
+def test_auto_add_timeout_for_spark_job(cmd, timeout_duration, expected):
+    result = auto_add_timeout_for_spark_job(cmd, timeout_duration)
+
+    assert result == expected