Fix command add_args

src/deployments/core/configs/command.py

@@ -66,20 +66,16 @@ def _add_arg_ignore(self, arg, existing_args_map: dict):
             self.args.append(arg)
 
     def _add_arg_replace(self, arg, existing_args_map: dict):
-        """
-        Helper for add_args with on_duplicate == "replace"
-
-        Note: This may modify `existing_args_map` for the caller.
-        """
         try:
             flag, _value = arg
         except ValueError:
             flag = arg
-        if flag in existing_args_map:
-            self.args = [item for item in self.args if item != arg]
-            # Remove from map so it won't be replaced again.
-            del existing_args_map[flag]
-            self.args.insert(existing_args_map[flag], arg)
+
+        if flag in existing_args_map and existing_args_map[flag]:
+            index = existing_args_map[flag].pop(0)
+            self.args[index] = arg
+            if not existing_args_map[flag]:
+                del existing_args_map[flag]
         else:
             self.args.append(arg)
 
@@ -89,28 +85,16 @@ def add_args(
         *,
         on_duplicate: Literal["error", "ignore", "replace"] = "error",
     ):
-        """Add args to command.
-
-        `on_duplicate` determines the behavior when an argument with the same key already exists.
-
-        `"error"`: raises ValueError if flag existed in self.args when this function was called
-
-        Note: This allows you to have duplicate flags in `args` param and add them to self.args.
-
-        `"ignore"`: does not add the argument
-
-        `"replace"`: replaces all instances of the argument in self.args with all new arguments with the same flag
-
-        Note: This allows you to pass in a new set of flags/value pairs to replace an existing set.
-
-        For example: --name=Alice --name=Bob can be replaced with --name=Alfred --name=Ben --name=Carl
-        """
         if isinstance(args, dict):
             args = [(key, value) for key, value in args.items()]
 
         existing_args_map = defaultdict(list)
         for index, arg in enumerate(self.args):
-            existing_args_map[arg].append(index)
+            try:
+                flag, _value = arg
+            except ValueError:
+                flag = arg
+            existing_args_map[flag].append(index)
 
         for arg in args:
             if on_duplicate == "error":

src/deployments/core/configs/tests/test_command.py

@@ -91,3 +91,27 @@ def test_build_command_script_mode_returns_shell_wrapper():
 def test_build_container_command_uses_custom_prefix():
     result = build_container_command(["echo hello"], prefix=["bash", "-lc"])
     assert result == ["bash", "-lc", "echo hello\n"]
+
+
+def test_add_args_replace_replaces_existing_flag_value_pair():
+    command = Command(command="app", args=[("--max-connections", "1000")])
+    command.add_args([("--max-connections", "200")], on_duplicate="replace")
+    assert command.args == [("--max-connections", "200")]
+
+
+def test_add_args_replace_consumes_duplicate_existing_args_in_order():
+    command = Command(command="app", args=[("--name", "alice"), ("--name", "bob")])
+    command.add_args(
+        [("--name", "carl"), ("--name", "dave")],
+        on_duplicate="replace",
+    )
+    assert command.args == [("--name", "carl"), ("--name", "dave")]
+
+
+def test_add_args_replace_appends_new_values_after_replacing_existing_duplicates():
+    command = Command(command="app", args=[("--name", "alice"), ("--name", "bob")])
+    command.add_args(
+        [("--name", "carl"), ("--name", "dave"), ("--name", "eric")],
+        on_duplicate="replace",
+    )
+    assert command.args == [("--name", "carl"), ("--name", "dave"), ("--name", "eric")]