From b1715a740104cbdb50ee6c826fa5955a101cc1a7 Mon Sep 17 00:00:00 2001
From: Gustavo Chain <me@qustavo.cc>
Date: Thu, 5 May 2022 20:34:00 +0200
Subject: [PATCH] allow user to specify a custom endpoint

---
 client/client.go       | 36 +++++++++++++++++++++++-------------
 main.go                |  6 +++++-
 ui/fee_distribution.go |  9 +++++----
 ui/ui.go               | 13 +++++++------
 4 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/client/client.go b/client/client.go
index d76046c..ad4d34f 100644
--- a/client/client.go
+++ b/client/client.go
@@ -5,14 +5,11 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"strings"
 
 	"github.com/gorilla/websocket"
 )
 
-const (
-	API_URL = "https://mempool.space/api/v1/ws"
-)
-
 type MempoolInfo struct {
 	Size  int `json:"size"`
 	Bytes int `json:"bytes"`
@@ -69,12 +66,22 @@ type Response struct {
 }
 
 type Client struct {
-	conn *websocket.Conn
+	conn     *websocket.Conn
+	endpoint string
 }
 
 func New() (*Client, error) {
+	return NewWithEndpoint("mempool.space")
+}
+
+func NewWithEndpoint(endpoint string) (*Client, error) {
+	if !strings.HasSuffix(endpoint, "/") {
+		endpoint = endpoint + "/"
+	}
+	endpoint = endpoint + "api/v1/ws"
+
 	dialer := websocket.Dialer{}
-	conn, _, err := dialer.Dial("wss://mempool.space/ws", nil)
+	conn, _, err := dialer.Dial("wss://"+endpoint, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -83,7 +90,10 @@ func New() (*Client, error) {
 		`{"action": "init"}`,
 	))
 
-	return &Client{conn: conn}, nil
+	return &Client{
+		conn:     conn,
+		endpoint: endpoint,
+	}, nil
 }
 
 func (c *Client) Read() (*Response, error) {
@@ -114,8 +124,8 @@ func (f Fees) Len() int           { return len(f) }
 func (f Fees) Less(i, j int) bool { return f[i].FPV < f[j].FPV }
 func (f Fees) Swap(i, j int)      { f[i], f[j] = f[j], f[i] }
 
-func Get(ctx context.Context, path string, v interface{}) error {
-	req, err := http.NewRequest("GET", API_URL+path, nil)
+func (c *Client) Get(ctx context.Context, path string, v interface{}) error {
+	req, err := http.NewRequest("GET", "https://"+c.endpoint+path, nil)
 	if err != nil {
 		return err
 	}
@@ -135,17 +145,17 @@ func Get(ctx context.Context, path string, v interface{}) error {
 	return json.NewDecoder(r.Body).Decode(v)
 }
 
-func GetMempoolFee(ctx context.Context, n int) (Fees, error) {
+func (c *Client) GetMempoolFee(ctx context.Context, n int) (Fees, error) {
 	var fees Fees
-	if err := Get(ctx, fmt.Sprintf("transactions/mempool/%d", n), &fees); err != nil {
+	if err := c.Get(ctx, fmt.Sprintf("/transactions/mempool/%d", n), &fees); err != nil {
 		return nil, err
 	}
 	return fees, nil
 }
 
-func GetBlockFee(ctx context.Context, n int) (Fees, error) {
+func (c *Client) GetBlockFee(ctx context.Context, n int) (Fees, error) {
 	var fees Fees
-	if err := Get(ctx, fmt.Sprintf("transactions/height/%d", n), &fees); err != nil {
+	if err := c.Get(ctx, fmt.Sprintf("/transactions/height/%d", n), &fees); err != nil {
 		return nil, err
 	}
 	return fees, nil
diff --git a/main.go b/main.go
index b0dfb05..6060403 100644
--- a/main.go
+++ b/main.go
@@ -1,13 +1,17 @@
 package main
 
 import (
+	"flag"
 	"log"
 
 	"github.com/mempool/mempool-cli/ui"
 )
 
 func main() {
-	gui, err := ui.New()
+	endpoint := flag.String("endpoint", "mempool.space", "The API endpoint")
+	flag.Parse()
+
+	gui, err := ui.New(*endpoint)
 	if err != nil {
 		log.Fatal(err)
 	}
diff --git a/ui/fee_distribution.go b/ui/fee_distribution.go
index f7c6c69..705e505 100644
--- a/ui/fee_distribution.go
+++ b/ui/fee_distribution.go
@@ -21,10 +21,11 @@ type FeeDistribution struct {
 	cancelFn    context.CancelFunc
 	fees        client.Fees
 	title       string
+	client      *client.Client
 }
 
-func NewFeeDistribution(g *gocui.Gui) *FeeDistribution {
-	return &FeeDistribution{gui: g}
+func NewFeeDistribution(g *gocui.Gui, c *client.Client) *FeeDistribution {
+	return &FeeDistribution{gui: g, client: c}
 }
 
 func (fd *FeeDistribution) newCtx() context.Context {
@@ -35,14 +36,14 @@ func (fd *FeeDistribution) newCtx() context.Context {
 
 func (fd *FeeDistribution) FetchProjection(n int) error {
 	fn := func(ctx context.Context) (client.Fees, error) {
-		return client.GetMempoolFee(ctx, n)
+		return fd.client.GetMempoolFee(ctx, n)
 	}
 	return fd.fetch(fn)
 }
 
 func (fd *FeeDistribution) FetchBlock(n int) error {
 	fn := func(ctx context.Context) (client.Fees, error) {
-		return client.GetBlockFee(ctx, n)
+		return fd.client.GetBlockFee(ctx, n)
 	}
 	return fd.fetch(fn)
 }
diff --git a/ui/ui.go b/ui/ui.go
index f451f95..4111f43 100644
--- a/ui/ui.go
+++ b/ui/ui.go
@@ -34,14 +34,19 @@ type UI struct {
 	state  state
 }
 
-func New() (*UI, error) {
+func New(endpoint string) (*UI, error) {
+	c, err := client.NewWithEndpoint(endpoint)
+	if err != nil {
+		log.Fatal(err)
+	}
+
 	gui, err := gocui.NewGui(gocui.Output256)
 	if err != nil {
 		return nil, err
 	}
 
 	ui := &UI{gui: gui}
-	ui.fd = NewFeeDistribution(gui)
+	ui.fd = NewFeeDistribution(gui, c)
 	ui.ts = NewTXSearch(gui)
 	gui.SetManager(ui, ui.fd, ui.ts)
 
@@ -54,10 +59,6 @@ func New() (*UI, error) {
 	gui.SelFgColor = gocui.ColorWhite
 
 	go func() {
-		c, err := client.New()
-		if err != nil {
-			log.Fatal(err)
-		}
 		if err := c.Want(); err != nil {
 			log.Fatal(err)
 		}