fix tests

fastpay/src/network.rs

@@ -221,9 +221,41 @@ impl AuthorityClient for Client {
 
     fn handle_info_request(
         &self,
-        request: AccountInfoRequest,
-    ) -> AsyncResult<'_, AccountInfoResponse, FastPayError> {
-        Box::pin(async move { self.send_recv_bytes(serialize_info_request(&request)).await })
+        request: InfoRequest,
+    ) -> AsyncResult<'_, InfoResponse, FastPayError> {
+        Box::pin(async move {
+            match request.kind {
+                InfoRequestKind::AccountInfoRequest(ref account_request) => {
+                    let mut all_object_ids = Vec::new();
+
+                    for shard in 0..self.num_shards {
+                        let result: Result<InfoResponse, FastPayError> = self
+                            .send_recv_bytes(shard, serialize_info_request(&request))
+                            .await;
+
+                        match result {
+                            Ok(InfoResponse {
+                                kind: InfoResponseKind::AccountInfoResponse(response),
+                            }) => all_object_ids.extend(response.object_ids),
+                            Err(e) => return Err(e),
+                            _ => return Err(FastPayError::UnexpectedMessage),
+                        }
+                    }
+                    Ok(InfoResponse::new(InfoResponseKind::AccountInfoResponse(
+                        AccountInfoResponse {
+                            object_ids: all_object_ids,
+                            owner: account_request.account,
+                        },
+                    )))
+                }
+                InfoRequestKind::ObjectInfoRequest(ref object_request) => {
+                    let shard =
+                        AuthorityState::get_shard(self.num_shards, &object_request.object_id);
+                    self.send_recv_bytes(shard, serialize_info_request(&request))
+                        .await
+                }
+            }
+        })
     }
 }
 

fastx_types/src/messages.rs

@@ -169,31 +169,6 @@ impl InfoResponse {
         InfoResponse { kind }
     }
 
-    pub fn new_account_info_resp(object_ids: Vec<ObjectRef>, owner: FastPayAddress) -> Self {
-        Self::new(InfoResponseKind::AccountInfoResponse(AccountInfoResponse {
-            object_ids,
-            owner,
-        }))
-    }
-
-    pub fn new_object_info_resp(
-        object_id: ObjectID,
-        owner: FastPayAddress,
-        next_sequence_number: SequenceNumber,
-        requested_certificate: Option<CertifiedOrder>,
-        pending_confirmation: Option<SignedOrder>,
-        requested_received_transfers: Vec<CertifiedOrder>,
-    ) -> Self {
-        Self::new(InfoResponseKind::ObjectInfoResponse(ObjectInfoResponse {
-            object_id,
-            owner,
-            next_sequence_number,
-            requested_certificate,
-            pending_confirmation,
-            requested_received_transfers,
-        }))
-    }
-
     pub fn to_account_info_response(&self) -> Option<AccountInfoResponse> {
         match self.kind.clone() {
             InfoResponseKind::AccountInfoResponse(response) => Some(response),

fastx_types/src/unit_tests/serialize_tests.rs

@@ -21,16 +21,16 @@ fn test_error() {
 
 #[test]
 fn test_info_request() {
-    let req1 = ObjectInfoRequest {
+    let req1 = InfoRequest::new(InfoRequestKind::ObjectInfoRequest(ObjectInfoRequest {
         object_id: dbg_object_id(0x20),
         request_sequence_number: None,
         request_received_transfers_excluding_first_nth: None,
-    };
-    let req2 = ObjectInfoRequest {
+    }));
+    let req2 = InfoRequest::new(InfoRequestKind::ObjectInfoRequest(ObjectInfoRequest {
         object_id: dbg_object_id(0x20),
         request_sequence_number: Some(SequenceNumber::from(129)),
         request_received_transfers_excluding_first_nth: None,
-    };
+    }));
 
     let buf1 = serialize_info_request(&req1);
     let buf2 = serialize_info_request(&req2);
@@ -40,13 +40,13 @@ fn test_info_request() {
     assert!(result1.is_ok());
     assert!(result2.is_ok());
 
-    if let SerializedMessage::ObjectInfoReq(o) = result1.unwrap() {
-        assert!(*o == req1);
+    if let SerializedMessage::InfoReq(o) = result1.unwrap() {
+        assert_eq!(*o, req1);
     } else {
         panic!()
     }
-    if let SerializedMessage::ObjectInfoReq(o) = result2.unwrap() {
-        assert!(*o == req2);
+    if let SerializedMessage::InfoReq(o) = result2.unwrap() {
+        assert_eq!(*o, req2);
     } else {
         panic!()
     }
@@ -219,45 +219,45 @@ fn test_info_response() {
         cert.signatures.push((authority_name, sig));
     }
 
-    let resp1 = AccountInfoResponse {
+    let resp1 = InfoResponse::new(InfoResponseKind::ObjectInfoResponse(ObjectInfoResponse {
         object_id: dbg_object_id(0x20),
         owner: dbg_addr(0x20),
         next_sequence_number: SequenceNumber::new(),
         pending_confirmation: None,
         requested_certificate: None,
         requested_received_transfers: Vec::new(),
-    };
-    let resp2 = AccountInfoResponse {
+    }));
+    let resp2 = InfoResponse::new(InfoResponseKind::ObjectInfoResponse(ObjectInfoResponse {
         object_id: dbg_object_id(0x20),
         owner: dbg_addr(0x20),
         next_sequence_number: SequenceNumber::new(),
         pending_confirmation: Some(vote.clone()),
         requested_certificate: None,
         requested_received_transfers: Vec::new(),
-    };
-    let resp3 = AccountInfoResponse {
+    }));
+    let resp3 = InfoResponse::new(InfoResponseKind::ObjectInfoResponse(ObjectInfoResponse {
         object_id: dbg_object_id(0x20),
         owner: dbg_addr(0x20),
         next_sequence_number: SequenceNumber::new(),
         pending_confirmation: None,
         requested_certificate: Some(cert.clone()),
         requested_received_transfers: Vec::new(),
-    };
-    let resp4 = AccountInfoResponse {
+    }));
+    let resp4 = InfoResponse::new(InfoResponseKind::ObjectInfoResponse(ObjectInfoResponse {
         object_id: dbg_object_id(0x20),
         owner: dbg_addr(0x20),
         next_sequence_number: SequenceNumber::new(),
         pending_confirmation: Some(vote),
         requested_certificate: Some(cert),
         requested_received_transfers: Vec::new(),
-    };
+    }));
 
     for resp in [resp1, resp2, resp3, resp4].iter() {
         let buf = serialize_info_response(resp);
         let result = deserialize_message(buf.as_slice());
         assert!(result.is_ok());
-        if let SerializedMessage::ObjectInfoResp(o) = result.unwrap() {
-            assert!(*o == *resp);
+        if let SerializedMessage::InfoResp(o) = result.unwrap() {
+            assert_eq!(*o, *resp);
         } else {
             panic!()
         }