From ac39ba4f53579ebd1128053f2588c8630439124d Mon Sep 17 00:00:00 2001
From: jx2lee <dev.jaejun.lee.1991@gmail.com>
Date: Thu, 19 Dec 2024 10:29:12 +0900
Subject: [PATCH] test

---
 .../routes/test_task_instances.py             | 83 +++++++++++++++++++
 .../execution_api/routes/test_variables.py    |  2 +-
 2 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/tests/api_fastapi/execution_api/routes/test_task_instances.py b/tests/api_fastapi/execution_api/routes/test_task_instances.py
index 4ed5f8f1598f3..cd7a49971cdb6 100644
--- a/tests/api_fastapi/execution_api/routes/test_task_instances.py
+++ b/tests/api_fastapi/execution_api/routes/test_task_instances.py
@@ -136,6 +136,39 @@ def test_ti_run_state_conflict_if_not_queued(
 
         assert session.scalar(select(TaskInstance.state).where(TaskInstance.id == ti.id)) == initial_ti_state
 
+    def test_ti_run_failed_with_extra(self, client, session, create_task_instance, time_machine):
+        """
+        Test that a 422 error is returned when extra fields are included in the payload.
+        """
+        instant_str = "2024-12-19T00:00:00Z"
+        instant = timezone.parse(instant_str)
+        time_machine.move_to(instant, tick=False)
+
+        ti = create_task_instance(
+            task_id="test_ti_run_failed_with_extra",
+            state=State.QUEUED,
+            session=session,
+            start_date=instant,
+        )
+
+        session.commit()
+
+        response = client.patch(
+            f"/execution/task-instances/{ti.id}/run",
+            json={
+                "state": "running",
+                "hostname": "random-hostname",
+                "unixname": "random-unixname",
+                "pid": 100,
+                "start_date": instant_str,
+                "foo": "bar",
+            },
+        )
+
+        assert response.status_code == 422
+        assert response.json()["detail"][0]["type"] == "extra_forbidden"
+        assert response.json()["detail"][0]["msg"] == "Extra inputs are not permitted"
+
 
 class TestTIUpdateState:
     def setup_method(self):
@@ -340,6 +373,27 @@ def test_ti_update_state_to_reschedule(self, client, session, create_task_instan
         assert trs[0].map_index == -1
         assert trs[0].duration == 129600
 
+    def test_ti_update_state_failed_with_extra(self, client, session, create_task_instance, time_machine):
+        """
+        Test that a 422 error is returned when extra fields are included in the payload.
+        """
+        ti = create_task_instance(
+            task_id="test_ti_update_state_failed_with_extra",
+            state=State.RUNNING,
+            session=session,
+            start_date=DEFAULT_START_DATE,
+        )
+
+        session.commit()
+
+        response = client.patch(
+            f"/execution/task-instances/{ti.id}/state", json={"state": "scheduled", "foo": "bar"}
+        )
+
+        assert response.status_code == 422
+        assert response.json()["detail"][0]["type"] == "extra_forbidden"
+        assert response.json()["detail"][0]["msg"] == "Extra inputs are not permitted"
+
 
 class TestTIHealthEndpoint:
     def setup_method(self):
@@ -536,6 +590,35 @@ def test_ti_update_state_to_failed_table_check(self, client, session, create_tas
         assert ti.next_kwargs is None
         assert ti.duration == 3600.00
 
+    def test_ti_heartbeat_with_extra(
+        self,
+        client,
+        session,
+        create_task_instance,
+        time_machine,
+    ):
+        """
+        Test that a 422 error is returned when extra fields are included in the payload.
+        """
+        ti = create_task_instance(
+            task_id="test_ti_heartbeat_when_task_not_running",
+            state=State.RUNNING,
+            hostname="random-hostname",
+            pid=1547,
+            session=session,
+        )
+        session.commit()
+        task_instance_id = ti.id
+
+        response = client.put(
+            f"/execution/task-instances/{task_instance_id}/heartbeat",
+            json={"hostname": "random-hostname", "pid": 1547, "foo": "bar"},
+        )
+
+        assert response.status_code == 422
+        assert response.json()["detail"][0]["type"] == "extra_forbidden"
+        assert response.json()["detail"][0]["msg"] == "Extra inputs are not permitted"
+
 
 class TestTIPutRTIF:
     def setup_method(self):
diff --git a/tests/api_fastapi/execution_api/routes/test_variables.py b/tests/api_fastapi/execution_api/routes/test_variables.py
index 45868e2a6092e..20a9b43c07ace 100644
--- a/tests/api_fastapi/execution_api/routes/test_variables.py
+++ b/tests/api_fastapi/execution_api/routes/test_variables.py
@@ -54,7 +54,7 @@ def test_variable_get_from_db(self, client, session):
         {"AIRFLOW_VAR_KEY1": "VALUE"},
     )
     def test_variable_get_from_env_var(self, client, session):
-        response = client.get("/execution/variables/key1", params={"foo": "bar"})
+        response = client.get("/execution/variables/key1")
 
         assert response.status_code == 200
         assert response.json() == {"key": "key1", "value": "VALUE"}