From dc34e82101258ff76acd88e6d5d0d50f55d44762 Mon Sep 17 00:00:00 2001
From: Wilmer Bandres <wilmerrafael@yelp.com>
Date: Tue, 13 Feb 2024 05:32:17 -0800
Subject: [PATCH] Adding more tests for pod merging labels process

---
 tests/bin/action_runner_test.py | 18 ++++++++++++++++++
 tron/bin/action_runner.py       | 10 ++++++----
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/tests/bin/action_runner_test.py b/tests/bin/action_runner_test.py
index d12e49892..1ecc206e6 100644
--- a/tests/bin/action_runner_test.py
+++ b/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")
 
diff --git a/tron/bin/action_runner.py b/tron/bin/action_runner.py
index 5907ae854..177f356d9 100755
--- a/tron/bin/action_runner.py
+++ b/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,13 +95,14 @@ 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"
@@ -107,7 +110,6 @@ def build_labels(run_id, original_labels=None):
     new_labels = dict(original_labels)
     new_labels["TRON_RUN_NUM"] = run_num
 
-    logging.debug(new_labels)
     return new_labels