Adding more tests for pod merging labels process

tests/bin/action_runner_test.py

@@ -156,6 +156,24 @@ def test_build_labels(self):
             TRON_RUN_NUM="10",
         )
 
+    def test_build_labels_with_merging(self):
+        current_labels = {"LABEL1": "value_1"}
+        labels = action_runner.build_labels("MASTER.foo.10.bar", current_labels)
+
+        assert labels == dict(
+            TRON_RUN_NUM="10",
+            LABEL1="value_1",
+        )
+
+    def test_build_labels_with_merging_on_unknown(self):
+        current_labels = {"LABEL1": "value_1"}
+        labels = action_runner.build_labels("asdf", current_labels)
+
+        assert labels == dict(
+            TRON_RUN_NUM="UNKNOWN",
+            LABEL1="value_1",
+        )
+
     def test_build_labels_invalid_run_id(self):
         labels = action_runner.build_labels("asdf")
 

tron/bin/action_runner.py

@@ -10,6 +10,8 @@
 import sys
 import threading
 import time
+from typing import Dict
+from typing import Optional
 
 from tron import yaml
 
@@ -93,21 +95,21 @@ def build_environment(run_id, original_env=None):
     return new_env
 
 
-def build_labels(run_id, original_labels=None):
+def build_labels(run_id: str, original_labels: Optional[Dict[str, str]] = None) -> Dict[str, str]:
     if original_labels is None:
         original_labels = dict()
 
     try:
-        run_num = run_id.split(".", maxsplit=3)[2]
-    except IndexError:
+        # reminder: the format here is "namespace.job.run_num.action"
+        _, _, run_num, _ = run_id.split(".", maxsplit=3)
+    except ValueError:
         # if we can't parse the run_id, we don't want to abort, so just
         # set these semi-arbitrarily
         run_num = "UNKNOWN"
 
     new_labels = dict(original_labels)
     new_labels["TRON_RUN_NUM"] = run_num
 
-    logging.debug(new_labels)
     return new_labels