SameerMesiah97 commented on code in PR #62626:
URL: https://github.com/apache/airflow/pull/62626#discussion_r2868065466


##########
airflow-core/src/airflow/cli/commands/task_command.py:
##########
@@ -203,7 +204,16 @@ def _get_ti(
                 f"TaskInstance for {dag.dag_id}, {task.task_id}, 
map={map_index} with "
                 f"run_id or logical_date of {logical_date_or_run_id!r} not 
found"
             )
-        # TODO: Validate map_index is in range?
+        if map_index >= 0:
+            try:
+                total = task.get_parse_time_mapped_ti_count()
+                if map_index >= total:
+                    raise ValueError(
+                        f"map_index {map_index} is out of range. "
+                        f"Task '{task.task_id}' has {total} mapped instance(s) 
[0..{total - 1}]."
+                    )
+            except (NotFullyPopulated, NotMapped):
+                pass  # Dynamic mapping — skip validation

Review Comment:
   I believe `NotFullyPopulated` and `NotMapped` should not be grouped 
together. The first does not indicate invalid user input; it arises during 
normal operation i.e. we are waiting for parsing to complete. `NotMapped` means 
the user is treating the task instance as it if is mapped, when it isn't and 
they should be prompted to correct their next input. I would do something like 
this instead:
   
   ```
       except NotFullyPopulated:
           # Dynamic mapping – cannot validate at parse time
           pass
       except NotMapped:
           raise ValueError(
               f"Task '{task.task_id}' is not mapped; map_index must be -1."
           )
   ```



##########
airflow-core/tests/unit/cli/commands/test_task_command.py:
##########
@@ -310,6 +310,72 @@ def test_mapped_task_render(self):
         assert "[3]" not in output
         assert "property: op_args" in output
 
+    @pytest.mark.usefixtures("testing_dag_bundle")
+    def test_mapped_task_render_out_of_range_map_index(self):
+        """Raise ValueError when map_index exceeds the parse-time mapped 
count."""
+        with pytest.raises(ValueError, match=r"map_index 5 is out of range.*3 
mapped instance"):
+            task_command.task_render(
+                self.parser.parse_args(
+                    [
+                        "tasks",
+                        "render",
+                        "test_mapped_classic",
+                        "consumer_literal",
+                        "2022-01-01",
+                        "--map-index",
+                        "5",
+                    ]
+                )
+            )
+
+    @pytest.mark.usefixtures("testing_dag_bundle")
+    def test_mapped_task_render_boundary_map_index(self):
+        """Render should succeed for the last valid map_index (count - 1)."""
+        with redirect_stdout(io.StringIO()) as stdout:
+            task_command.task_render(
+                self.parser.parse_args(
+                    [
+                        "tasks",
+                        "render",
+                        "test_mapped_classic",
+                        "consumer_literal",
+                        "2022-01-01",
+                        "--map-index",
+                        "2",
+                    ]
+                )
+            )
+        output = stdout.getvalue()
+        assert "[3]" in output
+        assert "property: op_args" in output
+
+    @pytest.mark.usefixtures("testing_dag_bundle")
+    def test_mapped_task_render_dynamic_skips_validation(self):
+        """Dynamic (XCom-based) mapping should skip map_index validation."""
+        # consumer depends on XCom from make_arg_lists, so parse-time count
+        # is not available. Validation should be skipped (NotFullyPopulated).
+        # The render may fail later for other reasons, but not with our
+        # "out of range" ValueError.
+        try:
+            task_command.task_render(
+                self.parser.parse_args(
+                    [
+                        "tasks",
+                        "render",
+                        "test_mapped_classic",
+                        "consumer",
+                        "2022-01-01",
+                        "--map-index",
+                        "999",
+                    ]
+                )
+            )
+        except ValueError as e:
+            if "out of range" in str(e):
+                pytest.fail(f"Should not raise 'out of range' for dynamic 
mapping, but got: {e}")
+        except Exception:
+            pass  # Other errors are expected (e.g., missing XCom data)

Review Comment:
   This does not look like a canonical pattern for asserting exceptions in 
tests. Also, this is too permissive. What if an exception besides `ValueError` 
is caused by an accidental refactor e.g. if `ValueError` is changed to say 
`RuntimeError`. I would suggest something like this instead:
   
   ```
   with pytest.raises(Exception) as excinfo:
       task_command.task_render(...)
   
   assert "out of range" not in str(excinfo.value)
   ```
   
   Using this pattern would be more consistent with the existing test suite + 
your new tests.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to