From 695de9e4f734a68468bdb977f127b386d7c99d06 Mon Sep 17 00:00:00 2001 From: Daniel Shaar Date: Mon, 9 Dec 2024 20:33:35 +0000 Subject: [PATCH] Fix experimental spawn test. --- test/conftest.py | 7 +++++++ test/function_test.py | 7 ++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/test/conftest.py b/test/conftest.py index 3dee4c7cf2..5a297b3ae3 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -819,6 +819,13 @@ async def EnvironmentGetOrCreate(self, stream): ### Function + async def FunctionAsyncInvoke(self, stream): + self.fcidx += 1 + request: api_pb2.FunctionAsyncInvokeRequest = await stream.recv_message() + function_call_id = f"fc-{self.fcidx}" + self.function_id_for_function_call[function_call_id] = request.function_id + await stream.send_message(api_pb2.FunctionAsyncInvokeResponse(function_call_id=function_call_id)) + async def FunctionBindParams(self, stream): from modal._serialization import deserialize diff --git a/test/function_test.py b/test/function_test.py index 28787f5d28..2ab15d1bf6 100644 --- a/test/function_test.py +++ b/test/function_test.py @@ -11,6 +11,7 @@ import modal from modal import App, Image, Mount, NetworkFileSystem, Proxy, asgi_app, batched, web_endpoint +from modal._serialization import deserialize from modal._utils.async_utils import synchronize_api from modal._vendor import cloudpickle from modal.exception import ExecutionError, InvalidError @@ -1012,9 +1013,9 @@ def test_experimental_spawn(client, servicer): with app.run(client=client): dummy_modal._experimental_spawn(1, 2) - # Verify the correct invocation type is set - function_map = ctx.pop_request("FunctionMap") - assert function_map.function_call_invocation_type == api_pb2.FUNCTION_CALL_INVOCATION_TYPE_ASYNC + # Verify the correct input was passed to the function. + request = ctx.pop_request("FunctionAsyncInvoke") + assert deserialize(request.input.args, client) == ((1, 2), {}) def test_from_name_web_url(servicer, set_env_client):