fix: adapt to pydantic 2.x.x

Signed-off-by: Guilhem Barthés <[email protected]>
Substra · Oct 6, 2023 · bd03de1 · bd03de1
1 parent 0bee8a1
commit bd03de1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/backend/builder/tests/test_task_build_image.py b/backend/builder/tests/test_task_build_image.py
@@ -37,8 +37,12 @@ def test_order_building_success(celery_app, celery_worker, mocker, execution_num
     mocker.patch("builder.tasks.task.get_orchestrator_client")
     mocker.patch("builder.tasks.tasks_build_image.build_image_if_missing", side_effect=lambda x, y: time.sleep(0.5))
 
-    result_1 = build_image.apply_async(kwargs={"function_serialized": function_1.json(), "channel_name": CHANNEL})
-    result_2 = build_image.apply_async(kwargs={"function_serialized": function_2.json(), "channel_name": CHANNEL})
+    result_1 = build_image.apply_async(
+        kwargs={"function_serialized": function_1.model_dump_json(), "channel_name": CHANNEL}
+    )
+    result_2 = build_image.apply_async(
+        kwargs={"function_serialized": function_2.model_dump_json(), "channel_name": CHANNEL}
+    )
     # get waits for the completion
     result_1.get()
 
@@ -70,10 +74,10 @@ def side_effect(*args, **kwargs):
     mocker.patch("builder.tasks.tasks_build_image.build_image_if_missing", side_effect=side_effect_creator())
 
     result_retry = build_image.apply_async(
-        kwargs={"function_serialized": function_retry.json(), "channel_name": CHANNEL}
+        kwargs={"function_serialized": function_retry.model_dump_json(), "channel_name": CHANNEL}
     )
     result_other = build_image.apply_async(
-        kwargs={"function_serialized": function_other.json(), "channel_name": CHANNEL}
+        kwargs={"function_serialized": function_other.model_dump_json(), "channel_name": CHANNEL}
     )
 
     result_retry.get()

diff --git a/backend/substrapp/events/reactor.py b/backend/substrapp/events/reactor.py
@@ -99,10 +99,9 @@ def on_function_event(payload):
 
             building_params = {
                 "channel_name": channel_name,
-                "function_serialized": orc_function.json(),
+                "function_serialized": orc_function.model_dump_json(),
             }
             (
-                # TODO switch to function.model_dump_json() as soon as pydantic is updated to > 2.0
                 build_image.si(**building_params).set(queue=builder_queue)
                 | save_image_task.si(**building_params).set(queue=WORKER_QUEUE)
             ).apply_async()