From 74b91c1f84b50c53f70e87007035d886cc7552cd Mon Sep 17 00:00:00 2001
From: stefan-mysten <135084671+stefan-mysten@users.noreply.github.com>
Date: Tue, 1 Oct 2024 06:40:15 -0700
Subject: [PATCH] sui-graphql-client: add method for querying total supply of a
 coin (#17)

---
 crates/sui-graphql-client/src/lib.rs | 29 ++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/crates/sui-graphql-client/src/lib.rs b/crates/sui-graphql-client/src/lib.rs
index 00bde3522..b1c3eb9b5 100644
--- a/crates/sui-graphql-client/src/lib.rs
+++ b/crates/sui-graphql-client/src/lib.rs
@@ -316,6 +316,7 @@ impl Client {
         })
     }
 
+    /// Get the coin metadata for the coin type.
     pub async fn coin_metadata(&self, coin_type: &str) -> Result<Option<CoinMetadata>, Error> {
         let operation = CoinMetadataQuery::build(CoinMetadataArgs { coin_type });
         let response = self.run_query(&operation).await?;
@@ -327,6 +328,16 @@ impl Client {
         Ok(response.data.and_then(|x| x.coin_metadata))
     }
 
+    /// Get total supply for the coin type.
+    pub async fn total_supply(&self, coin_type: &str) -> Result<Option<u64>, Error> {
+        let coin_metadata = self.coin_metadata(coin_type).await?;
+
+        coin_metadata
+            .and_then(|c| c.supply)
+            .map(|c| c.try_into())
+            .transpose()
+    }
+
     // ===========================================================================
     // Checkpoints API
     // ===========================================================================
@@ -908,4 +919,22 @@ mod tests {
         }
         assert!(num_coins > 0);
     }
+
+    #[tokio::test]
+    async fn test_total_supply() {
+        for (n, _) in NETWORKS {
+            let client = Client::new(n).unwrap();
+            let ts = client.total_supply("0x2::sui::SUI").await;
+            assert!(
+                ts.is_ok(),
+                "Total supply query failed for network: {n}. Error: {}",
+                ts.unwrap_err()
+            );
+            assert_eq!(
+                ts.unwrap().unwrap(),
+                10_000_000_000,
+                "Total supply mismatch for network: {n}"
+            );
+        }
+    }
 }