Improve balance api (#879)

* Get balance from the last address * Fix some unit tests
archethic-foundation · Feb 1, 2023 · f8dfd04 · f8dfd04
1 parent 8563823
commit f8dfd04
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 28 deletions.
diff --git 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
@@ -34,22 +34,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
 

diff --git a/test/archethic_web/graphql_schema_test.exs b/test/archethic_web/graphql_schema_test.exs
@@ -476,8 +476,12 @@ defmodule ArchethicWeb.GraphQLSchemaTest do
       addr = <<0::8, 0::8, :crypto.strong_rand_bytes(32)::binary>>
 
       MockClient
-      |> stub(:send_message, fn _, %GetBalance{}, _ ->
-        {:ok, %Balance{uco: 218_000_000}}
+      |> stub(:send_message, fn
+        _, %GetBalance{}, _ ->
+          {:ok, %Balance{uco: 218_000_000}}
+
+        _, %GetLastTransactionAddress{address: address}, _ ->
+          {:ok, %LastTransactionAddress{address: address}}
       end)
 
       conn =
@@ -492,15 +496,19 @@ defmodule ArchethicWeb.GraphQLSchemaTest do
       addr = <<0::8, 0::8, :crypto.strong_rand_bytes(32)::binary>>
 
       MockClient
-      |> stub(:send_message, fn _, %GetBalance{}, _ ->
-        {:ok,
-         %Balance{
-           token: %{
-             {"@Token1", 0} => 200_000_000,
-             {"@Token2", 0} => 500_000_000,
-             {"@Token3", 0} => 1_000_000_000
-           }
-         }}
+      |> stub(:send_message, fn
+        _, %GetBalance{}, _ ->
+          {:ok,
+           %Balance{
+             token: %{
+               {"@Token1", 0} => 200_000_000,
+               {"@Token2", 0} => 500_000_000,
+               {"@Token3", 0} => 1_000_000_000
+             }
+           }}
+
+        _, %GetLastTransactionAddress{address: address}, _ ->
+          {:ok, %LastTransactionAddress{address: address}}
       end)
 
       conn =