Minor updates to mock server behavior (#3314)

* Minor updates to mock server behavior

* Fix tests that assert on servicer.deployed_apps directly

Author: mwaskom

test/app_test.py

@@ -91,14 +91,16 @@ def test_create_object_invalid_exception(servicer, client):
 def test_deploy_falls_back_to_app_name(servicer, client):
     named_app = App(name="foo_app")
     named_app.deploy(client=client)
-    assert "foo_app" in servicer.deployed_apps
+    app_names = {app_name for (_, app_name) in servicer.deployed_apps}
+    assert "foo_app" in app_names
 
 
 def test_deploy_uses_deployment_name_if_specified(servicer, client):
     named_app = App(name="foo_app")
     named_app.deploy(name="bar_app", client=client)
-    assert "bar_app" in servicer.deployed_apps
-    assert "foo_app" not in servicer.deployed_apps
+    app_names = {app_name for (_, app_name) in servicer.deployed_apps}
+    assert "bar_app" in app_names
+    assert "foo_app" not in app_names
 
 
 def test_run_function_without_app_error():

test/cli_test.py

@@ -76,14 +76,16 @@ def test_app_deploy_success(servicer, mock_dir, set_env_client):
         # Deploy as a script with an absolute path
         _run(["deploy", os.path.abspath("myapp.py")])
 
-    assert "my_app" in servicer.deployed_apps
+    app_names = {app_name for (_, app_name) in servicer.deployed_apps}
+    assert "my_app" in app_names
 
 
 def test_app_deploy_with_name(servicer, mock_dir, set_env_client):
     with mock_dir({"myapp.py": dummy_app_file, "other_module.py": dummy_other_module_file}):
         _run(["deploy", "myapp.py", "--name", "my_app_foo"])
 
-    assert "my_app_foo" in servicer.deployed_apps
+    app_names = {app_name for (_, app_name) in servicer.deployed_apps}
+    assert "my_app_foo" in app_names
 
 
 def test_secret_create_list_delete(servicer, set_env_client):
@@ -987,7 +989,7 @@ def test_app_history(servicer, mock_dir, set_env_client):
     with mock_dir({"myapp.py": dummy_app_file, "other_module.py": dummy_other_module_file}):
         _run(["deploy", "myapp.py", "--name", "my_app_foo"])
 
-    app_id = servicer.deployed_apps.get("my_app_foo")
+    app_id = servicer.deployed_apps.get(("main", "my_app_foo"))
 
     servicer.app_deployment_history[app_id][-1]["commit_info"] = api_pb2.CommitInfo(
         vcs="git", branch="main", commit_hash="abc123"
@@ -1027,11 +1029,11 @@ def test_app_rollback(servicer, mock_dir, set_env_client):
         for _ in range(4):
             _run(["deploy", "myapp.py", "--name", "my_app"])
     _run(["app", "rollback", "my_app"])
-    app_id = servicer.deployed_apps.get("my_app")
+    app_id = servicer.deployed_apps.get(("main", "my_app"))
     assert servicer.app_deployment_history[app_id][-1]["rollback_version"] == 3
 
     _run(["app", "rollback", "my_app", "v2"])
-    app_id = servicer.deployed_apps.get("my_app")
+    app_id = servicer.deployed_apps.get(("main", "my_app"))
     assert servicer.app_deployment_history[app_id][-1]["rollback_version"] == 2
 
     _run(["app", "rollback", "my_app", "2"], expected_exit_code=2)

test/conftest.py

@@ -186,9 +186,8 @@ def __init__(self, blob_host, blobs, blocks, files_sha2data, credentials, port):
         self.fc_data_in = defaultdict(lambda: asyncio.Queue())  # unbounded
         self.fc_data_out = defaultdict(lambda: asyncio.Queue())  # unbounded
         self.queue: dict[bytes, list[bytes]] = {b"": []}
-        self.deployed_apps = {
-            client_mount_name(): "ap-x",
-        }
+        self.deployed_apps: dict[tuple[str, str], str] = {}
+        self.app_environments: dict[str, str] = {}
         self.app_deployment_history: defaultdict[str, list[dict[str, Any]]] = defaultdict(list)
         self.app_deployment_history["ap-x"] = [
             {
@@ -474,6 +473,11 @@ def get_object_metadata(self, object_id) -> api_pb2.Object:
         res.object_id = object_id
         return res
 
+    def get_environment(self, environment_name: Optional[str] = None) -> str:
+        if environment_name is None:
+            return next(iter(self.environments))  # Use first environment as default
+        return environment_name
+
     def mounts_excluding_published_client(self):
         return {
             mount_id: content
@@ -513,6 +517,7 @@ async def AppCreate(self, stream):
         self.n_apps += 1
         app_id = f"ap-{self.n_apps}"
         self.app_state_history[app_id].append(api_pb2.APP_STATE_INITIALIZING)
+        self.app_environments[app_id] = self.get_environment(request.environment_name)
         await stream.send_message(
             api_pb2.AppCreateResponse(app_id=app_id, app_page_url="https://modaltest.com/apps/ap-123")
         )
@@ -582,13 +587,6 @@ async def AppRollback(self, stream):
         self.app_state_history[request.app_id].append(api_pb2.APP_STATE_DEPLOYED)
         await stream.send_message(Empty())
 
-    async def AppDeploy(self, stream):
-        request: api_pb2.AppDeployRequest = await stream.recv_message()
-        self.deployed_apps[request.name] = request.app_id
-        self.app_state_history[request.app_id].append(api_pb2.APP_STATE_DEPLOYED)
-
-        await stream.send_message(api_pb2.AppDeployResponse(url="http://test.modal.com/foo/bar"))
-
     async def AppPublish(self, stream):
         request: api_pb2.AppPublishRequest = await stream.recv_message()
         for key, val in request.definition_ids.items():
@@ -599,7 +597,8 @@ async def AppPublish(self, stream):
         self.app_objects[request.app_id] = {**request.function_ids, **request.class_ids}
         self.app_state_history[request.app_id].append(request.app_state)
         if request.app_state == api_pb2.AppState.APP_STATE_DEPLOYED:
-            self.deployed_apps[request.name] = request.app_id
+            app_key = (self.app_environments[request.app_id], request.name)
+            self.deployed_apps[app_key] = request.app_id
             await stream.send_message(api_pb2.AppPublishResponse(url="http://test.modal.com/foo/bar"))
         else:
             await stream.send_message(api_pb2.AppPublishResponse())
@@ -623,7 +622,9 @@ async def AppPublish(self, stream):
 
     async def AppGetByDeploymentName(self, stream):
         request: api_pb2.AppGetByDeploymentNameRequest = await stream.recv_message()
-        await stream.send_message(api_pb2.AppGetByDeploymentNameResponse(app_id=self.deployed_apps.get(request.name)))
+        app_key = (self.get_environment(request.environment_name), request.name)
+        app_id = self.deployed_apps.get(app_key)
+        await stream.send_message(api_pb2.AppGetByDeploymentNameResponse(app_id=app_id))
 
     async def AppHeartbeat(self, stream):
         request: api_pb2.AppHeartbeatRequest = await stream.recv_message()
@@ -656,9 +657,12 @@ async def AppDeploymentHistory(self, stream):
         )
 
     async def AppList(self, stream):
-        await stream.recv_message()
+        req = await stream.recv_message()
         apps = []
-        for app_name, app_id in self.deployed_apps.items():
+        requested_environment = self.get_environment(req.environment_name)
+        for (environment_name, app_name), app_id in self.deployed_apps.items():
+            if environment_name != requested_environment:
+                continue
             apps.append(
                 api_pb2.AppListResponse.AppListItem(
                     name=app_name,
@@ -757,7 +761,8 @@ async def ClassCreate(self, stream):
 
     async def ClassGet(self, stream):
         request: api_pb2.ClassGetRequest = await stream.recv_message()
-        if not (app_id := self.deployed_apps.get(request.app_name)):
+        app_key = (self.get_environment(request.environment_name), request.app_name)
+        if not (app_id := self.deployed_apps.get(app_key)):
             raise GRPCError(Status.NOT_FOUND, f"can't find app {request.app_name}")
         app_objects = self.app_objects[app_id]
         object_id = app_objects.get(request.object_tag)
@@ -867,7 +872,6 @@ async def DictGetOrCreate(self, stream):
             dict_id = f"di-{len(self.dicts)}"
             self.dicts[dict_id] = {entry.key: entry.value for entry in request.data}
             self.deployed_dicts[k] = dict_id
-            self.deployed_apps[request.deployment_name] = f"ap-{dict_id}"
         elif request.object_creation_type == api_pb2.OBJECT_CREATION_TYPE_EPHEMERAL:
             dict_id = f"di-{len(self.dicts)}"
             self.dicts[dict_id] = {entry.key: entry.value for entry in request.data}
@@ -900,11 +904,7 @@ async def DictLen(self, stream):
         await stream.send_message(api_pb2.DictLenResponse(len=len(self.dicts[request.dict_id])))
 
     async def DictList(self, stream):
-        dicts = [
-            api_pb2.DictListResponse.DictInfo(name=name, created_at=1)
-            for name, _ in self.deployed_dicts
-            if name in self.deployed_apps
-        ]
+        dicts = [api_pb2.DictListResponse.DictInfo(name=name, created_at=1) for name, _ in self.deployed_dicts]
         await stream.send_message(api_pb2.DictListResponse(dicts=dicts))
 
     async def DictUpdate(self, stream):
@@ -1140,7 +1140,8 @@ async def FunctionCreate(self, stream):
 
     async def FunctionGet(self, stream):
         request: api_pb2.FunctionGetRequest = await stream.recv_message()
-        if not (app_id := self.deployed_apps.get(request.app_name)):
+        app_key = (self.get_environment(request.environment_name), request.app_name)
+        if not (app_id := self.deployed_apps.get(app_key)):
             raise GRPCError(Status.NOT_FOUND, f"can't find app {request.app_name}")
 
         app_objects = self.app_objects[app_id]
@@ -1514,7 +1515,6 @@ async def QueueGetOrCreate(self, stream):
             self.n_queues += 1
             queue_id = f"qu-{self.n_queues}"
             self.deployed_queues[k] = queue_id
-            self.deployed_apps[request.deployment_name] = f"ap-{queue_id}"
         elif request.object_creation_type == api_pb2.OBJECT_CREATION_TYPE_EPHEMERAL:
             self.n_queues += 1
             queue_id = f"qu-{self.n_queues}"
@@ -1562,11 +1562,7 @@ async def QueueLen(self, stream):
     async def QueueList(self, stream):
         # TODO Note that the actual self.queue holding the data assumes we have a single queue
         # So there is a mismatch and I am not implementing a mock for the num_partitions / total_size
-        queues = [
-            api_pb2.QueueListResponse.QueueInfo(name=name, created_at=1)
-            for name, _ in self.deployed_queues
-            if name in self.deployed_apps
-        ]
+        queues = [api_pb2.QueueListResponse.QueueInfo(name=name, created_at=1) for name, _ in self.deployed_queues]
         await stream.send_message(api_pb2.QueueListResponse(queues=queues))
 
     async def QueueNextItems(self, stream):