From f88de145f4ff05a05bdf4b4b9a9003259cb72d57 Mon Sep 17 00:00:00 2001
From: Matt Curtis <matt@obscu.ro>
Date: Tue, 3 Oct 2023 15:18:04 +0100
Subject: [PATCH] Add balance endpoint to faucet for monitoring

---
 integration/faucet/faucet_test.go    | 43 ++++++++++++++++++++++++++++
 tools/faucet/faucet/faucet.go        |  4 +++
 tools/faucet/webserver/web_server.go | 16 +++++++++++
 3 files changed, 63 insertions(+)

diff --git a/integration/faucet/faucet_test.go b/integration/faucet/faucet_test.go
index 0ddbae6964..3085a5324e 100644
--- a/integration/faucet/faucet_test.go
+++ b/integration/faucet/faucet_test.go
@@ -2,6 +2,7 @@ package faucet
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"math/big"
@@ -60,6 +61,10 @@ func TestFaucet(t *testing.T) {
 	err = faucetContainer.Start()
 	assert.NoError(t, err)
 
+	initialFaucetBal, err := getFaucetBalance(faucetConfig.ServerPort)
+	assert.NoError(t, err)
+	assert.NotZero(t, initialFaucetBal)
+
 	rndWallet := datagenerator.RandomWallet(integration.ObscuroChainID)
 	err = fundWallet(faucetConfig.ServerPort, rndWallet)
 	assert.NoError(t, err)
@@ -73,6 +78,12 @@ func TestFaucet(t *testing.T) {
 	if currentBalance.Cmp(big.NewInt(0)) <= 0 {
 		t.Fatalf("Unexpected balance, got: %d, expected > 0", currentBalance.Int64())
 	}
+
+	endFaucetBal, err := getFaucetBalance(faucetConfig.ServerPort)
+	assert.NoError(t, err)
+	assert.NotZero(t, endFaucetBal)
+	// faucet balance should have decreased
+	assert.Less(t, endFaucetBal.Cmp(initialFaucetBal), 0)
 }
 
 // Creates a single-node Obscuro network for testing.
@@ -124,3 +135,35 @@ func fundWallet(port int, w wallet.Wallet) error {
 	fmt.Println(string(body))
 	return nil
 }
+
+func getFaucetBalance(port int) (*big.Int, error) {
+	url := fmt.Sprintf("http://localhost:%d/balance", port)
+	method := "GET"
+
+	client := &http.Client{}
+	req, err := http.NewRequestWithContext(context.Background(), method, url, nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Add("Content-Type", "application/json")
+
+	res, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer res.Body.Close()
+
+	var resp struct {
+		Balance string `json:"balance"`
+	}
+	err = json.NewDecoder(res.Body).Decode(&resp)
+	if err != nil {
+		return nil, err
+	}
+	bal, success := new(big.Int).SetString(resp.Balance, 10)
+	if !success {
+		return nil, fmt.Errorf("failed to parse balance - %s", resp.Balance)
+	}
+
+	return bal, nil
+}
diff --git a/tools/faucet/faucet/faucet.go b/tools/faucet/faucet/faucet.go
index ff5d5e6a37..a0da725431 100644
--- a/tools/faucet/faucet/faucet.go
+++ b/tools/faucet/faucet/faucet.go
@@ -140,3 +140,7 @@ func (f *Faucet) fundNativeToken(address *common.Address, amount *big.Int) (*typ
 
 	return signedTx, nil
 }
+
+func (f *Faucet) Balance(ctx context.Context) (*big.Int, error) {
+	return f.client.BalanceAt(ctx, nil)
+}
diff --git a/tools/faucet/webserver/web_server.go b/tools/faucet/webserver/web_server.go
index f3bf7967bb..77b944b763 100644
--- a/tools/faucet/webserver/web_server.go
+++ b/tools/faucet/webserver/web_server.go
@@ -35,6 +35,8 @@ func NewWebServer(faucetServer *faucet.Faucet, bindAddress string, jwtSecret []b
 	// todo (@matt) we need to remove this unsecure endpoint before we provide a fully public sepolia faucet
 	r.POST("/fund/:token", fundingHandler(faucetServer, defaultAmount))
 
+	r.GET("/balance", balanceReqHandler(faucetServer))
+
 	return &WebServer{
 		engine:      r,
 		faucet:      faucetServer,
@@ -153,3 +155,17 @@ func fundingHandler(faucetServer *faucet.Faucet, defaultAmount *big.Int) gin.Han
 		c.JSON(http.StatusOK, gin.H{"status": "ok"})
 	}
 }
+
+// returns the remaining native balance of the faucet
+func balanceReqHandler(faucetServer *faucet.Faucet) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		// get the balance
+		balance, err := faucetServer.Balance(c)
+		if err != nil {
+			errorHandler(c, fmt.Errorf("unable to get balance %w", err), faucetServer.Logger)
+			return
+		}
+
+		c.JSON(http.StatusOK, gin.H{"balance": balance.String()})
+	}
+}