From 74cf7ebd0816ca995e54e768985398e956aed177 Mon Sep 17 00:00:00 2001
From: Samuel Manzanera <samuelmanzanera@protonmail.com>
Date: Tue, 31 Jan 2023 20:50:08 +0100
Subject: [PATCH] Get balance from the last address

---
 lib/archethic_web/graphql_schema.ex          |  2 +-
 lib/archethic_web/graphql_schema/resolver.ex | 29 +++++++++-----------
 2 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/lib/archethic_web/graphql_schema.ex b/lib/archethic_web/graphql_schema.ex
index 5713a3423c..71759217ba 100644
--- a/lib/archethic_web/graphql_schema.ex
+++ b/lib/archethic_web/graphql_schema.ex
@@ -101,7 +101,7 @@ defmodule ArchethicWeb.GraphQLSchema do
     end
 
     @desc """
-    Query the network to find a balance from an address
+    Query the network to find a balance from an address coming from the latest transaction on the chain
     """
     field :balance, :balance do
       arg(:address, non_null(:address))
diff --git a/lib/archethic_web/graphql_schema/resolver.ex b/lib/archethic_web/graphql_schema/resolver.ex
index 2e376eb222..485c97d660 100644
--- a/lib/archethic_web/graphql_schema/resolver.ex
+++ b/lib/archethic_web/graphql_schema/resolver.ex
@@ -33,22 +33,19 @@ defmodule ArchethicWeb.GraphQLSchema.Resolver do
   end
 
   def get_balance(address) do
-    case Archethic.get_balance(address) do
-      {:ok, %{uco: uco, token: token_balances}} ->
-        balance = %{
-          uco: uco,
-          token:
-            token_balances
-            |> Enum.map(fn {{address, token_id}, amount} ->
-              %{address: address, amount: amount, token_id: token_id}
-            end)
-            |> Enum.sort_by(& &1.amount)
-        }
-
-        {:ok, balance}
-
-      {:error, :network_issue} = e ->
-        e
+    with {:ok, last_address} <- Archethic.get_last_transaction_address(address),
+         {:ok, %{uco: uco, token: token_balances}} <- Archethic.get_balance(last_address) do
+      balance = %{
+        uco: uco,
+        token:
+          token_balances
+          |> Enum.map(fn {{address, token_id}, amount} ->
+            %{address: address, amount: amount, token_id: token_id}
+          end)
+          |> Enum.sort_by(& &1.amount)
+      }
+
+      {:ok, balance}
     end
   end