Update mock servicer

modal-labs · Oct 6, 2024 · 1c86b0c · 1c86b0c
1 parent 898b711
commit 1c86b0c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/test/conftest.py b/test/conftest.py
@@ -870,8 +870,6 @@ async def FunctionCreate(self, stream):
         else:
             self.n_functions += 1
             function_id = f"fu-{self.n_functions}"
-        if request.schedule:
-            self.function2schedule[function_id] = request.schedule
 
         function: Optional[api_pb2.Function] = None
         function_data: Optional[api_pb2.FunctionData] = None
@@ -890,8 +888,12 @@ async def FunctionCreate(self, stream):
 
         assert (function is None) != (function_data is None)
         function_defn = function or function_data
+        assert function_defn
         self.app_functions[function_id] = function_defn
 
+        if function_defn.schedule:
+            self.function2schedule[function_id] = function_defn.schedule
+
         await stream.send_message(
             api_pb2.FunctionCreateResponse(
                 function_id=function_id,