From e0743fc5591d9f5d7e6a0c6462deff5791a7bb89 Mon Sep 17 00:00:00 2001
From: Matthieu Moquet <matthieu@moquet.net>
Date: Mon, 9 Oct 2023 15:12:19 +0200
Subject: [PATCH] feat: add token-exponent option

---
 pkg/app/flags.go               |  5 +++++
 pkg/app/run.go                 |  5 ++++-
 pkg/watcher/validators.go      | 17 ++++++++++++++---
 pkg/watcher/validators_test.go |  3 +++
 4 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/pkg/app/flags.go b/pkg/app/flags.go
index 9f64856..e279cef 100644
--- a/pkg/app/flags.go
+++ b/pkg/app/flags.go
@@ -39,6 +39,11 @@ var Flags = []cli.Flag{
 		Name:  "no-staking",
 		Usage: "disable calls to staking module (useful for consumer chains)",
 	},
+	&cli.UintFlag{
+		Name:  "token-exponent",
+		Usage: "token exponent (ie. 6 for uatom)",
+		Value: 6,
+	},
 	&cli.StringSliceFlag{
 		Name:  "validator",
 		Usage: "validator address(es) to track (use :my-label to add a custom label in metrics & ouput)",
diff --git a/pkg/app/run.go b/pkg/app/run.go
index 2b39556..348fbc7 100644
--- a/pkg/app/run.go
+++ b/pkg/app/run.go
@@ -37,6 +37,7 @@ func RunFunc(cCtx *cli.Context) error {
 		nodes      = cCtx.StringSlice("node")
 		noGov      = cCtx.Bool("no-gov")
 		noStaking  = cCtx.Bool("no-staking")
+		tokenExpon = cCtx.Uint("token-exponent")
 		validators = cCtx.StringSlice("validator")
 		xGov       = cCtx.String("x-gov")
 	)
@@ -95,7 +96,9 @@ func RunFunc(cCtx *cli.Context) error {
 	// Pool watchers
 	//
 	if !noStaking {
-		validatorsWatcher := watcher.NewValidatorsWatcher(trackedValidators, metrics, pool)
+		validatorsWatcher := watcher.NewValidatorsWatcher(trackedValidators, metrics, pool, watcher.ValidatorsWatcherOptions{
+			TokenExponent: tokenExpon,
+		})
 		errg.Go(func() error {
 			return validatorsWatcher.Start(ctx)
 		})
diff --git a/pkg/watcher/validators.go b/pkg/watcher/validators.go
index ca53f8b..7845ca4 100644
--- a/pkg/watcher/validators.go
+++ b/pkg/watcher/validators.go
@@ -20,13 +20,19 @@ type ValidatorsWatcher struct {
 	metrics    *metrics.Metrics
 	validators []TrackedValidator
 	pool       *rpc.Pool
+	opts       ValidatorsWatcherOptions
 }
 
-func NewValidatorsWatcher(validators []TrackedValidator, metrics *metrics.Metrics, pool *rpc.Pool) *ValidatorsWatcher {
+type ValidatorsWatcherOptions struct {
+	TokenExponent uint
+}
+
+func NewValidatorsWatcher(validators []TrackedValidator, metrics *metrics.Metrics, pool *rpc.Pool, opts ValidatorsWatcherOptions) *ValidatorsWatcher {
 	return &ValidatorsWatcher{
 		metrics:    metrics,
 		validators: validators,
 		pool:       pool,
+		opts:       opts,
 	}
 }
 
@@ -71,9 +77,14 @@ func (w *ValidatorsWatcher) handleValidators(chainID string, validators []stakin
 	// Sort validators by tokens & status (bonded, unbonded, jailed)
 	sort.Sort(RankedValidators(validators))
 
+	tokenExponent := w.opts.TokenExponent
+	if tokenExponent == 0 {
+		tokenExponent = 1
+	}
+
 	seatPrice := decimal.Zero
 	for _, val := range validators {
-		tokens := decimal.NewFromBigInt(val.Tokens.BigInt(), -6)
+		tokens := decimal.NewFromBigInt(val.Tokens.BigInt(), -int32(tokenExponent))
 		if val.Status == staking.Bonded && (seatPrice.IsZero() || seatPrice.GreaterThan(tokens)) {
 			seatPrice = tokens
 		}
@@ -92,7 +103,7 @@ func (w *ValidatorsWatcher) handleValidators(chainID string, validators []stakin
 					rank     = i + 1
 					isBonded = val.Status == staking.Bonded
 					isJailed = val.Jailed
-					tokens   = decimal.NewFromBigInt(val.Tokens.BigInt(), -6)
+					tokens   = decimal.NewFromBigInt(val.Tokens.BigInt(), -int32(tokenExponent))
 				)
 
 				w.metrics.Rank.WithLabelValues(chainID, address, name).Set(float64(rank))
diff --git a/pkg/watcher/validators_test.go b/pkg/watcher/validators_test.go
index ceb2b26..ef1376b 100644
--- a/pkg/watcher/validators_test.go
+++ b/pkg/watcher/validators_test.go
@@ -29,6 +29,9 @@ func TestValidatorsWatcher(t *testing.T) {
 		},
 		metrics.New("cosmos_validator_watcher"),
 		nil,
+		ValidatorsWatcherOptions{
+			TokenExponent: 6,
+		},
 	)
 
 	t.Run("Handle Validators", func(t *testing.T) {