Refactoring: move verify_casting out so it can be re-used by other tests (#1303)

Summary:
Pull Request resolved: #1303

Two refactorings to make the next diff D82537988 looks less busy:

1. extract `cast_v0`;
2. extract `verify_casting`.

Reviewed By: mariusae

Differential Revision: D83001963

fbshipit-source-id: 33cb027b73e26d8b1606ff4eb5baa13454d2de6a

Author: pzhan9

hyperactor_mesh/src/v1/actor_mesh.rs

@@ -33,6 +33,7 @@ use serde::Deserializer;
 use serde::Serialize;
 use serde::Serializer;
 
+use crate::CommActor;
 use crate::actor_mesh as v0_actor_mesh;
 use crate::comm::multicast;
 use crate::proc_mesh::mesh_agent::ActorState;
@@ -136,33 +137,7 @@ impl<A: Actor + RemoteActor> ActorMeshRef<A> {
         M: Castable + RemoteMessage + Clone, // Clone is required until we are fully onto comm actor
     {
         if let Some(root_comm_actor) = self.proc_mesh.root_comm_actor() {
-            let cast_mesh_shape = view::Ranked::region(self).into();
-            let actor_mesh_id = ActorMeshId::V1(self.name.clone());
-            match &self.proc_mesh.root_region {
-                Some(root_region) => {
-                    let root_mesh_shape = root_region.into();
-                    v0_actor_mesh::cast_to_sliced_mesh::<A, M>(
-                        cx,
-                        actor_mesh_id,
-                        root_comm_actor,
-                        &sel!(*),
-                        message,
-                        &cast_mesh_shape,
-                        &root_mesh_shape,
-                    )
-                    .map_err(|e| Error::CastingError(self.name.clone(), e.into()))
-                }
-                None => v0_actor_mesh::actor_mesh_cast::<A, M>(
-                    cx,
-                    actor_mesh_id,
-                    root_comm_actor,
-                    sel!(*),
-                    &cast_mesh_shape,
-                    &cast_mesh_shape,
-                    message,
-                )
-                .map_err(|e| Error::CastingError(self.name.clone(), e.into())),
-            }
+            self.cast_v0(cx, message, root_comm_actor)
         } else {
             for (point, actor) in self.iter() {
                 let mut headers = Attrs::new();
@@ -180,6 +155,45 @@ impl<A: Actor + RemoteActor> ActorMeshRef<A> {
         }
     }
 
+    fn cast_v0<M>(
+        &self,
+        cx: &impl context::Actor,
+        message: M,
+        root_comm_actor: &ActorRef<CommActor>,
+    ) -> v1::Result<()>
+    where
+        A: RemoteHandles<M> + RemoteHandles<IndexedErasedUnbound<M>>,
+        M: Castable + RemoteMessage + Clone, // Clone is required until we are fully onto comm actor
+    {
+        let cast_mesh_shape = view::Ranked::region(self).into();
+        let actor_mesh_id = ActorMeshId::V1(self.name.clone());
+        match &self.proc_mesh.root_region {
+            Some(root_region) => {
+                let root_mesh_shape = root_region.into();
+                v0_actor_mesh::cast_to_sliced_mesh::<A, M>(
+                    cx,
+                    actor_mesh_id,
+                    root_comm_actor,
+                    &sel!(*),
+                    message,
+                    &cast_mesh_shape,
+                    &root_mesh_shape,
+                )
+                .map_err(|e| Error::CastingError(self.name.clone(), e.into()))
+            }
+            None => v0_actor_mesh::actor_mesh_cast::<A, M>(
+                cx,
+                actor_mesh_id,
+                root_comm_actor,
+                sel!(*),
+                &cast_mesh_shape,
+                &cast_mesh_shape,
+                message,
+            )
+            .map_err(|e| Error::CastingError(self.name.clone(), e.into())),
+        }
+    }
+
     pub async fn actor_states(
         &self,
         cx: &impl context::Actor,

hyperactor_mesh/src/v1/testactor.rs

@@ -38,14 +38,16 @@ use hyperactor::mailbox;
 use hyperactor::supervision::ActorSupervisionEvent;
 use ndslice::Point;
 #[cfg(test)]
-use ndslice::ViewExt;
+use ndslice::ViewExt as _;
 use serde::Deserialize;
 use serde::Serialize;
 
 use crate::comm::multicast::CastInfo;
 #[cfg(test)]
 use crate::v1::ActorMesh;
 #[cfg(test)]
+use crate::v1::ActorMeshRef;
+#[cfg(test)]
 use crate::v1::testing;
 
 /// A simple test actor used by various unit tests.
@@ -218,28 +220,7 @@ impl Handler<GetCastInfo> for TestActor {
 pub async fn assert_mesh_shape(actor_mesh: ActorMesh<TestActor>) {
     let instance = testing::instance().await;
     // Verify casting to the root actor mesh
-    {
-        let (port, mut rx) = mailbox::open_port(&instance);
-        actor_mesh.cast(instance, GetActorId(port.bind())).unwrap();
-
-        let mut expected_actor_ids: HashSet<_> = actor_mesh
-            .values()
-            .map(|actor_ref| actor_ref.actor_id().clone())
-            .collect();
-
-        while !expected_actor_ids.is_empty() {
-            let actor_id = rx.recv().await.unwrap();
-            assert!(
-                expected_actor_ids.remove(&actor_id),
-                "got {actor_id}, expect {expected_actor_ids:?}"
-            );
-        }
-
-        // No more messages
-        RealClock.sleep(Duration::from_secs(1)).await;
-        let result = rx.try_recv();
-        assert!(result.as_ref().unwrap().is_none(), "got {result:?}");
-    }
+    assert_casting_correctness(&actor_mesh, instance).await;
 
     // Just pick the first dimension. Slice half of it off.
     // actor_mesh.extent().
@@ -248,28 +229,33 @@ pub async fn assert_mesh_shape(actor_mesh: ActorMesh<TestActor>) {
 
     // Verify casting to the sliced actor mesh
     let sliced_actor_mesh = actor_mesh.range(&label, 0..size).unwrap();
-    {
-        let (port, mut rx) = mailbox::open_port(instance);
-        sliced_actor_mesh
-            .cast(instance, GetActorId(port.bind()))
-            .unwrap();
+    assert_casting_correctness(&sliced_actor_mesh, instance).await;
+}
 
-        let mut expected_actor_ids: HashSet<_> = sliced_actor_mesh
-            .values()
-            .map(|actor_ref| actor_ref.actor_id().clone())
-            .collect();
+#[cfg(test)]
+/// Cast to the actor mesh, and verify that all actors are reached.
+pub async fn assert_casting_correctness(
+    actor_mesh: &ActorMeshRef<TestActor>,
+    instance: &Instance<()>,
+) {
+    let (port, mut rx) = mailbox::open_port(instance);
+    actor_mesh.cast(instance, GetActorId(port.bind())).unwrap();
 
-        while !expected_actor_ids.is_empty() {
-            let actor_id = rx.recv().await.unwrap();
-            assert!(
-                expected_actor_ids.remove(&actor_id),
-                "got {actor_id}, expect {expected_actor_ids:?}"
-            );
-        }
+    let mut expected_actor_ids: HashSet<_> = actor_mesh
+        .values()
+        .map(|actor_ref| actor_ref.actor_id().clone())
+        .collect();
 
-        // No more messages
-        RealClock.sleep(Duration::from_secs(1)).await;
-        let result = rx.try_recv();
-        assert!(result.as_ref().unwrap().is_none(), "got {result:?}");
+    while !expected_actor_ids.is_empty() {
+        let actor_id = rx.recv().await.unwrap();
+        assert!(
+            expected_actor_ids.remove(&actor_id),
+            "got {actor_id}, expect {expected_actor_ids:?}"
+        );
     }
+
+    // No more messages
+    RealClock.sleep(Duration::from_secs(1)).await;
+    let result = rx.try_recv();
+    assert!(result.as_ref().unwrap().is_none(), "got {result:?}");
 }