improve test

modal/object.py

@@ -44,7 +44,7 @@ class _Object:
     _object_id: str
     _client: _Client
     _is_hydrated: bool
-    _is_rehydrating: bool
+    _is_rehydrated: bool
 
     @classmethod
     def __init_subclass__(cls, type_prefix: Optional[str] = None):
@@ -222,6 +222,7 @@ async def resolve(self):
                 self._is_hydrated = False  # un-hydrate and re-resolve
                 resolver = Resolver(await _Client.from_env())
                 await resolver.load(self)
+                self._is_rehydrated = True
             return
         elif not self._hydrate_lazily:
             self._validate_is_hydrated()

test/container_app_test.py

@@ -110,6 +110,10 @@ async def test_container_snapshot_reference_capture(container_client, tmpdir, se
     assert new_app_id != app_id
     await f.remote.aio()
     assert f.object_id == "fu-2"
+    # Purposefully break FunctionGet to check the hydration is cached.
+    del servicer.app_objects[new_app_id]
+    await f.remote.aio()  # remote call succeeds because it didn't re-hydrate Function
+    assert f.object_id == "fu-2"
     channel.close()