Repository: apache/airflow
File: airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py (Python)

Code before:
            mock_register_asset_changes_in_db.return_value = None
            response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
            assert response.status_code == 500
            assert response.json()["detail"] == "Database error occurred"

    def test_ti_update_state_to_deferred(self, client, session, create_task_instance, time_machine):
        """
        Test that tests if the transition to deferred state is handled correctly.
        """
        ti = create_task_instance(
            task_id="test_ti_update_state_to_deferred",
            state=State.RUNNING,
            session=session,
        )
        session.commit()

        instant = timezone.datetime(2024, 11, 22)
        time_machine.move_to(instant, tick=False)

        from airflow.sdk.serde import serialize

        payload = {
            "state": "deferred",
            "trigger_kwargs": serialize(
                {"key": "value", "moment": datetime(2024, 12, 18, 0, 0, 1, tzinfo=timezone.utc)}
            ),
            "trigger_timeout": "P1D",  # 1 day
            "classpath": "my-classpath",
            "next_method": "execute_callback",
            "next_kwargs": serialize(
                {"foo": datetime(2024, 12, 18, 0, 0, 0, tzinfo=timezone.utc), "bar": "abc"}
            ),
        }

        response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)

        assert response.status_code == 204
        assert response.text == ""

        session.expire_all()

        tis = session.scalars(select(TaskInstance)).all()
        assert len(tis) == 1

        assert tis[0].state == TaskInstanceState.DEFERRED
        assert tis[0].next_method == "execute_callback"
        from airflow.sdk.serde import serialize

        assert tis[0].next_kwargs == serialize(
            {
                "bar": "abc",

Review comment (suggestion):
let's put the raw value -- i.e. output of `serialize(
                {"key": "value", "moment": datetime(2024, 12, 18, 0, 0, 1, tzinfo=timezone.utc)}`

Code after:
        )
        session.commit()
        payload = {
            "state": "success",
            "end_date": "2024-10-31T12:00:00Z",
        }

        with (
            mock.patch(
                "airflow.api_fastapi.common.db.common.Session.execute",
                side_effect=[
                    mock.Mock(one=lambda: ("running", 1, 0, "dag")),  # First call returns "queued"
                    mock.Mock(one=lambda: ("running", 1, 0, "dag")),  # Second call returns "queued"
                    SQLAlchemyError("Database error"),  # Last call raises an error
                ],
            ),
            mock.patch(
                "airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db",
            ) as mock_register_asset_changes_in_db,
        ):
            mock_register_asset_changes_in_db.return_value = None
            response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
            assert response.status_code == 500
            assert response.json()["detail"] == "Database error occurred"

    def test_ti_update_state_to_deferred(self, client, session, create_task_instance, time_machine):
        """
        Test that tests if the transition to deferred state is handled correctly.
        """
        ti = create_task_instance(
            task_id="test_ti_update_state_to_deferred",
            state=State.RUNNING,
            session=session,
        )
        session.commit()

        instant = timezone.datetime(2024, 11, 22)
        time_machine.move_to(instant, tick=False)

        payload = {
            "state": "deferred",
            # expected format is now in serde serialized format
            "trigger_kwargs": {
                "key": "value",
                "moment": {
                    "__classname__": "datetime.datetime",
                    "__version__": 2,
                    "__data__": {
                        "timestamp": 1734480001.0,
                        "tz": {
                            "__classname__": "builtins.tuple",