From 97189d78e8407ace9e411795440546582d7f84bd Mon Sep 17 00:00:00 2001
From: pmantica11 <151664502+pmantica11@users.noreply.github.com>
Date: Tue, 19 Nov 2024 20:33:02 -0600
Subject: [PATCH] Fix pagination for token balances (#251)

* Fix pagination for token balances

* Fix pagination for token balances
---
 Cargo.lock                                    |  2 +-
 Cargo.toml                                    |  2 +-
 .../get_compressed_token_balances_by_owner.rs |  2 +-
 tests/integration_tests/mock_tests.rs         | 39 +++++++++++++++++++
 4 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 76392f3..8836f86 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4591,7 +4591,7 @@ dependencies = [
 
 [[package]]
 name = "photon-indexer"
-version = "0.49.0"
+version = "0.50.0"
 dependencies = [
  "anchor-lang",
  "anyhow",
diff --git a/Cargo.toml b/Cargo.toml
index 7f0da29..c5422c9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,7 +7,7 @@ name = "photon-indexer"
 publish = true
 readme = "README.md"
 repository = "https://github.com/helius-labs/photon"
-version = "0.49.0"
+version = "0.50.0"
 
 [[bin]]
 name = "photon"
diff --git a/src/api/method/get_compressed_token_balances_by_owner.rs b/src/api/method/get_compressed_token_balances_by_owner.rs
index 24eacb9..732a32c 100644
--- a/src/api/method/get_compressed_token_balances_by_owner.rs
+++ b/src/api/method/get_compressed_token_balances_by_owner.rs
@@ -66,7 +66,7 @@ pub async fn get_compressed_token_balances_by_owner(
                 bytes.len()
             )));
         };
-        filter = filter.and(token_owner_balances::Column::Mint.gte::<Vec<u8>>(mint.into()));
+        filter = filter.and(token_owner_balances::Column::Mint.gt::<Vec<u8>>(mint.into()));
     }
     let limit = limit.map(|l| l.value()).unwrap_or(PAGE_LIMIT);
 
diff --git a/tests/integration_tests/mock_tests.rs b/tests/integration_tests/mock_tests.rs
index fbc1a6a..32c67b0 100644
--- a/tests/integration_tests/mock_tests.rs
+++ b/tests/integration_tests/mock_tests.rs
@@ -657,6 +657,45 @@ async fn test_persist_token_data(
             assert_eq!(item.balance.0, *owner_to_balance.get(&item.owner).unwrap());
         }
     }
+
+    let mut owner_to_balances = HashMap::new();
+    for (mint, balances) in mint_to_owner_to_balance.into_iter() {
+        for (owner, balance) in balances.into_iter() {
+            owner_to_balances
+                .entry(owner)
+                .or_insert(HashMap::new())
+                .insert(mint, balance);
+        }
+    }
+    for owner in owner_to_balances.keys() {
+        let mut cursor = None;
+        let mut mint_to_balance = HashMap::new();
+
+        loop {
+            let request = GetCompressedTokenBalancesByOwnerRequest {
+                owner: owner.clone(),
+                cursor,
+                limit: Some(photon_indexer::api::method::utils::Limit::new(1).unwrap()),
+                ..Default::default()
+            };
+            let res = setup
+                .api
+                .get_compressed_token_balances_by_owner(request)
+                .await
+                .unwrap()
+                .value;
+            let token_balances = res.token_balances;
+            for token_balance in token_balances.iter() {
+                let balance = mint_to_balance.entry(token_balance.mint).or_insert(0);
+                *balance += token_balance.balance.0;
+            }
+            cursor = res.cursor;
+            if cursor.is_none() {
+                break;
+            }
+        }
+        assert_eq!(mint_to_balance, *owner_to_balances.get(owner).unwrap());
+    }
 }
 
 #[tokio::test]