read_topic API

xmtp_mls/src/api_client_wrapper.rs

@@ -1,21 +1,27 @@
 use std::collections::HashMap;
 
 use xmtp_proto::{
-    api_client::{Error as ApiError, ErrorKind, XmtpApiClient, XmtpMlsClient},
+    api_client::{
+        Envelope, Error as ApiError, ErrorKind, PagingInfo, QueryRequest, XmtpApiClient,
+        XmtpMlsClient,
+    },
     xmtp::{
-        message_api::v3::GetIdentityUpdatesRequest,
+        message_api::{v1::Cursor, v3::GetIdentityUpdatesRequest},
         mls::message_contents::{
             group_message::{Version as GroupMessageVersion, V1 as GroupMessageV1},
             welcome_message::{Version as WelcomeMessageVersion, V1 as WelcomeMessageV1},
             WelcomeMessage as WelcomeMessageProto,
         },
     },
     xmtp::{
-        message_api::v3::{
-            get_identity_updates_response::update::Kind as UpdateKind,
-            publish_welcomes_request::WelcomeMessageRequest, ConsumeKeyPackagesRequest,
-            KeyPackageUpload, PublishToGroupRequest, PublishWelcomesRequest,
-            RegisterInstallationRequest, UploadKeyPackagesRequest,
+        message_api::{
+            v1::SortDirection,
+            v3::{
+                get_identity_updates_response::update::Kind as UpdateKind,
+                publish_welcomes_request::WelcomeMessageRequest, ConsumeKeyPackagesRequest,
+                KeyPackageUpload, PublishToGroupRequest, PublishWelcomesRequest,
+                RegisterInstallationRequest, UploadKeyPackagesRequest,
+            },
         },
         mls::message_contents::GroupMessage,
     },
@@ -34,6 +40,46 @@ where
         Self { api_client }
     }
 
+    pub async fn read_topic(
+        &self,
+        topic: &str,
+        start_time_ns: u64,
+    ) -> Result<Vec<Envelope>, ApiError> {
+        let mut cursor: Option<Cursor> = None;
+        let mut out: Vec<Envelope> = vec![];
+        let page_size = 100;
+        loop {
+            let mut result = self
+                .api_client
+                .query(QueryRequest {
+                    content_topics: vec![topic.to_string()],
+                    start_time_ns,
+                    end_time_ns: 0,
+                    paging_info: Some(PagingInfo {
+                        cursor,
+                        limit: page_size,
+                        direction: SortDirection::Ascending as i32,
+                    }),
+                })
+                .await?;
+
+            let num_envelopes = result.envelopes.len();
+            out.append(&mut result.envelopes);
+
+            if num_envelopes < page_size as usize || result.paging_info.is_none() {
+                break;
+            }
+
+            cursor = result.paging_info.expect("Empty paging info").cursor;
+
+            if cursor.is_none() {
+                break;
+            }
+        }
+
+        Ok(out)
+    }
+
     pub async fn register_installation(
         &self,
         last_resort_key_package: Vec<u8>,
@@ -228,7 +274,10 @@ type IdentityUpdatesMap = HashMap<String, Vec<IdentityUpdate>>;
 mod tests {
     use super::ApiClientWrapper;
     use mockall::mock;
-    use xmtp_proto::api_client::{Error, XmtpApiClient, XmtpApiSubscription, XmtpMlsClient};
+    use xmtp_proto::api_client::{
+        Error, PagingInfo, XmtpApiClient, XmtpApiSubscription, XmtpMlsClient,
+    };
+    use xmtp_proto::xmtp::message_api::v1::IndexCursor;
     use xmtp_proto::xmtp::message_api::v3::consume_key_packages_response::KeyPackage;
     use xmtp_proto::xmtp::message_api::v3::get_identity_updates_response::update::Kind as UpdateKind;
     use xmtp_proto::xmtp::message_api::v3::get_identity_updates_response::{
@@ -241,12 +290,24 @@ mod tests {
     };
 
     use xmtp_proto::xmtp::message_api::v1::{
-        BatchQueryRequest, BatchQueryResponse, Envelope, PublishRequest, PublishResponse,
-        QueryRequest, QueryResponse, SubscribeRequest,
+        cursor::Cursor as InnerCursor, BatchQueryRequest, BatchQueryResponse, Cursor, Envelope,
+        PublishRequest, PublishResponse, QueryRequest, QueryResponse, SubscribeRequest,
     };
 
     use async_trait::async_trait;
 
+    fn build_envelopes(num_envelopes: usize, topic: &str) -> Vec<Envelope> {
+        let mut out: Vec<Envelope> = vec![];
+        for i in 0..num_envelopes {
+            out.push(Envelope {
+                content_topic: topic.to_string(),
+                message: vec![i as u8],
+                timestamp_ns: i as u64,
+            })
+        }
+        out
+    }
+
     mock! {
         pub Subscription {}
 
@@ -442,4 +503,115 @@ mod tests {
             }
         }
     }
+
+    #[tokio::test]
+    async fn test_read_topic_single_page() {
+        let mut mock_api = MockApiClient::new();
+        let topic = "topic";
+        let start_time_ns = 10;
+        // Set expectation for first request with no cursor
+        mock_api.expect_query().returning(move |req| {
+            assert_eq!(req.content_topics[0], topic);
+
+            Ok(QueryResponse {
+                paging_info: Some(PagingInfo {
+                    cursor: None,
+                    limit: 100,
+                    direction: 0,
+                }),
+                envelopes: build_envelopes(10, topic),
+            })
+        });
+
+        let wrapper = ApiClientWrapper::new(mock_api);
+
+        let result = wrapper.read_topic(topic, start_time_ns).await.unwrap();
+        assert_eq!(result.len(), 10);
+    }
+
+    #[tokio::test]
+    async fn test_read_topic_single_page_exactly_100_results() {
+        let mut mock_api = MockApiClient::new();
+        let topic = "topic";
+        let start_time_ns = 10;
+        // Set expectation for first request with no cursor
+        mock_api.expect_query().times(1).returning(move |req| {
+            assert_eq!(req.content_topics[0], topic);
+
+            Ok(QueryResponse {
+                paging_info: Some(PagingInfo {
+                    cursor: None,
+                    limit: 100,
+                    direction: 0,
+                }),
+                envelopes: build_envelopes(100, topic),
+            })
+        });
+
+        let wrapper = ApiClientWrapper::new(mock_api);
+
+        let result = wrapper.read_topic(topic, start_time_ns).await.unwrap();
+        assert_eq!(result.len(), 100);
+    }
+
+    #[tokio::test]
+    async fn test_read_topic_multi_page() {
+        let mut mock_api = MockApiClient::new();
+        let topic = "topic";
+        let start_time_ns = 10;
+        // Set expectation for first request with no cursor
+        mock_api
+            .expect_query()
+            .withf(move |req| match req.paging_info.clone() {
+                Some(paging_info) => match paging_info.cursor {
+                    Some(_) => false,
+                    None => true,
+                },
+                None => true,
+            } && req.start_time_ns == 10)
+            .returning(move |req| {
+                assert_eq!(req.content_topics[0], topic);
+
+                Ok(QueryResponse {
+                    paging_info: Some(PagingInfo {
+                        cursor: Some(Cursor {
+                            cursor: Some(InnerCursor::Index(IndexCursor {
+                                digest: vec![],
+                                sender_time_ns: 0,
+                            })),
+                        }),
+                        limit: 100,
+                        direction: 0,
+                    }),
+                    envelopes: build_envelopes(100, topic),
+                })
+            });
+        // Set expectation for requests with a cursor
+        mock_api
+            .expect_query()
+            .withf(|req| match req.paging_info.clone() {
+                Some(paging_info) => match paging_info.cursor {
+                    Some(_) => true,
+                    None => false,
+                },
+                None => false,
+            })
+            .returning(move |req| {
+                assert_eq!(req.content_topics[0], topic);
+
+                Ok(QueryResponse {
+                    paging_info: Some(PagingInfo {
+                        cursor: None,
+                        limit: 100,
+                        direction: 0,
+                    }),
+                    envelopes: build_envelopes(100, topic),
+                })
+            });
+
+        let wrapper = ApiClientWrapper::new(mock_api);
+
+        let result = wrapper.read_topic(topic, start_time_ns).await.unwrap();
+        assert_eq!(result.len(), 200);
+    }
 }

xmtp_mls/src/client.rs

@@ -211,7 +211,7 @@ mod tests {
         let wallet = generate_local_wallet();
         let wallet_address = wallet.get_address();
         let client = ClientBuilder::new_test_client(wallet.clone().into()).await;
-
+        
         client.register_identity().await.unwrap();
         client.top_up_key_packages().await.unwrap();