feat(faucet): ensure successful transfer before responding

starship/faucet/handler.go

@@ -2,7 +2,11 @@ package main
 
 import (
 	"context"
+	"fmt"
+	"math/big"
+	"time"
 
+	"go.uber.org/zap"
 	"google.golang.org/protobuf/types/known/emptypb"
 
 	pb "github.com/cosmology-tech/starship/faucet/faucet"
@@ -37,11 +41,64 @@ func (a *AppServer) Status(ctx context.Context, _ *emptypb.Empty) (*pb.State, er
 	return state, nil
 }
 
+func (a *AppServer) getBalance(address, denom string) (*big.Int, error) {
+  account := &Account{config: a.config, logger: a.logger, Address: address}
+  coin, err := account.GetBalanceByDenom(denom)
+  if err != nil {
+    // Log the error, but don't return it
+    a.logger.Debug("Error getting balance, assuming new account", zap.Error(err))
+    return new(big.Int), nil // Return 0 balance
+  }
+  balance, ok := new(big.Int).SetString(coin.Amount, 10)
+  if !ok {
+    return nil, fmt.Errorf("failed to parse balance")
+  }
+  return balance, nil
+}
+
 func (a *AppServer) Credit(ctx context.Context, requestCredit *pb.RequestCredit) (*pb.ResponseCredit, error) {
-	err := a.distributor.SendTokens(requestCredit.GetAddress(), requestCredit.GetDenom())
-	if err != nil {
-		return nil, err
-	}
+  // Get initial balance before sending tokens
+  initialBalance, err := a.getBalance(requestCredit.GetAddress(), requestCredit.GetDenom())
+  if err != nil {
+    return nil, fmt.Errorf("failed to get initial balance: %v", err)
+  }
+
+  err = a.distributor.SendTokens(requestCredit.GetAddress(), requestCredit.GetDenom())
+  if err != nil {
+    return nil, err
+  }
+
+  // Check balance after transfer
+  confirmed, err := a.confirmBalanceUpdate(requestCredit.GetAddress(), requestCredit.GetDenom(), initialBalance)
+  if err != nil {
+    return &pb.ResponseCredit{Status: fmt.Sprintf("error: %v", err)}, err
+  }
+  if !confirmed {
+    return &pb.ResponseCredit{Status: "error: failed to confirm balance update (timeout)"}, nil
+  }
+
+  return &pb.ResponseCredit{Status: "ok"}, nil
+}
+
+func (a *AppServer) confirmBalanceUpdate(address, denom string, initialBalance *big.Int) (bool, error) {
+  expectedIncrease, ok := new(big.Int).SetString(a.distributor.CreditCoins.GetDenomAmount(denom), 10)
+  if !ok {
+    return false, fmt.Errorf("failed to parse expected amount")
+  }
+
+  expectedFinalBalance := new(big.Int).Add(initialBalance, expectedIncrease)
 
-	return &pb.ResponseCredit{Status: "ok"}, nil
+  for i := 0; i < 3; i++ { // Try 3 times with 5-second intervals
+    currentBalance, err := a.getBalance(address, denom)
+    if err != nil {
+      return false, err
+    }
+    if currentBalance.Cmp(expectedFinalBalance) >= 0 {
+      return true, nil
+    }
+    if i < 2 {
+      time.Sleep(5 * time.Second)
+    }
+  }
+  return false, nil
 }

starship/tests/e2e/faucet_test.go

@@ -4,11 +4,11 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	pb "github.com/cosmology-tech/starship/registry/registry"
 	"net/http"
 	urlpkg "net/url"
 	"strconv"
-	"time"
+
+	pb "github.com/cosmology-tech/starship/registry/registry"
 )
 
 func (s *TestSuite) MakeFaucetRequest(chain *Chain, req *http.Request, unmarshal map[string]interface{}) {
@@ -115,6 +115,28 @@ func (s *TestSuite) getAccountBalance(chain *Chain, address string, denom string
 	return float64(0)
 }
 
+func (s *TestSuite) creditAccount(chain *Chain, addr, denom string) error {
+	body := map[string]string{
+		"denom":   denom,
+		"address": addr,
+	}
+	postBody, err := json.Marshal(body)
+	if err != nil {
+		return err
+	}
+	resp, err := http.Post(
+		fmt.Sprintf("http://0.0.0.0:%d/credit", chain.Ports.Faucet),
+		"application/json",
+		bytes.NewBuffer(postBody))
+	if err != nil {
+		return err
+	}
+	if resp.StatusCode != 200 {
+		return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
+	}
+	return nil
+}
+
 func (s *TestSuite) TestFaucet_Credit() {
 	s.T().Log("running test for /credit endpoint for faucet")
 
@@ -132,20 +154,9 @@ func (s *TestSuite) TestFaucet_Credit() {
 			addr := getAddressFromType(chain.Name)
 			beforeBalance := s.getAccountBalance(chain, addr, denom)
 
-			body := map[string]string{
-				"denom":   denom,
-				"address": addr,
-			}
-			postBody, err := json.Marshal(body)
-			s.Require().NoError(err)
-			resp, err := http.Post(
-				fmt.Sprintf("http://0.0.0.0:%d/credit", chain.Ports.Faucet),
-				"application/json",
-				bytes.NewBuffer(postBody))
+			err := s.creditAccount(chain, addr, denom)
 			s.Require().NoError(err)
-			s.Require().Equal(200, resp.StatusCode)
 
-			time.Sleep(4 * time.Second)
 			afterBalance := s.getAccountBalance(chain, addr, denom)
 			s.T().Log("address:", addr, "after balance: ", afterBalance, "before balance:", beforeBalance)
 			// note sometimes expected difference is 9x expected value (bug due to using holder address for test)
@@ -154,3 +165,40 @@ func (s *TestSuite) TestFaucet_Credit() {
 		})
 	}
 }
+
+func (s *TestSuite) TestFaucet_Credit_MultipleRequests() {
+	s.T().Log("running test for multiple requests to /credit endpoint for faucet")
+
+	// expected amount to be credited via faucet
+	expCreditedAmt := float64(10000000000)
+
+	for _, chain := range s.config.Chains {
+		s.Run(fmt.Sprintf("multiple faucet requests test for: %s", chain.ID), func() {
+			if chain.Ports.Faucet == 0 {
+				s.T().Skip("faucet not exposed via ports")
+			}
+
+			// fetch denom and address from an account on chain
+			denom := s.getChainDenoms(chain)
+			addr := getAddressFromType(chain.Name)
+			beforeBalance := s.getAccountBalance(chain, addr, denom)
+
+			// Send multiple requests
+			numRequests := 3
+			for i := 0; i < numRequests; i++ {
+				err := s.creditAccount(chain, addr, denom)
+				s.Require().NoError(err)
+			}
+
+			afterBalance := s.getAccountBalance(chain, addr, denom)
+			s.T().Log("address:", addr, "after balance: ", afterBalance, "before balance:", beforeBalance)
+
+			// Check that the balance has increased by at least the expected amount times the number of requests
+			expectedIncrease := expCreditedAmt * float64(numRequests)
+			actualIncrease := afterBalance - beforeBalance
+			s.Require().GreaterOrEqual(actualIncrease, expectedIncrease, 
+				"Balance didn't increase as expected. Actual increase: %f, Expected increase: %f", 
+				actualIncrease, expectedIncrease)
+		})
+	}
+}