From 8f0d58aee9f66b356a3191c0cc5b6a38247d8943 Mon Sep 17 00:00:00 2001
From: c9s <yoanlin93@gmail.com>
Date: Fri, 4 Oct 2024 23:32:11 +0800
Subject: [PATCH] add new test helper to create balance map objects

Signed-off-by: c9s <yoanlin93@gmail.com>
---
 pkg/bbgo/account_value_calc_test.go | 24 ++++++++--------
 pkg/testing/testhelper/balance.go   | 43 +++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 12 deletions(-)
 create mode 100644 pkg/testing/testhelper/balance.go

diff --git a/pkg/bbgo/account_value_calc_test.go b/pkg/bbgo/account_value_calc_test.go
index a2a4a6689..6036f18b3 100644
--- a/pkg/bbgo/account_value_calc_test.go
+++ b/pkg/bbgo/account_value_calc_test.go
@@ -215,19 +215,19 @@ func Test_usdFiatBalances(t *testing.T) {
 	}{
 		{
 			args: args{
-				balances: types.BalanceMap{
-					"USDC": types.Balance{Currency: "USDC", Available: Number(70.0 + 80.0)},
-					"USDT": types.Balance{Currency: "USDT", Available: Number(100.0)},
-					"BTC":  types.Balance{Currency: "BTC", Available: Number(0.01)},
-				},
-			},
-			wantFiats: types.BalanceMap{
-				"USDC": types.Balance{Currency: "USDC", Available: Number(70.0 + 80.0)},
-				"USDT": types.Balance{Currency: "USDT", Available: Number(100.0)},
-			},
-			wantRest: types.BalanceMap{
-				"BTC": types.Balance{Currency: "BTC", Available: Number(0.01)},
+				balances: BalancesFromText(`
+					USDC, 150.0
+					USDT, 100.0
+					BTC, 0.01
+				`),
 			},
+			wantFiats: BalancesFromText(`
+				USDC, 150.0
+				USDT, 100.0
+				`),
+			wantRest: BalancesFromText(`
+				BTC, 0.01
+				`),
 		},
 	}
 	for _, tt := range tests {
diff --git a/pkg/testing/testhelper/balance.go b/pkg/testing/testhelper/balance.go
new file mode 100644
index 000000000..4480ecd80
--- /dev/null
+++ b/pkg/testing/testhelper/balance.go
@@ -0,0 +1,43 @@
+package testhelper
+
+import (
+	"strings"
+
+	"github.com/c9s/bbgo/pkg/fixedpoint"
+	"github.com/c9s/bbgo/pkg/types"
+)
+
+func BalancesFromText(str string) types.BalanceMap {
+	balances := make(types.BalanceMap)
+	lines := strings.Split(str, "\n")
+	for _, line := range lines {
+		line = strings.TrimSpace(line)
+		if len(line) == 0 {
+			continue
+		}
+
+		cols := strings.SplitN(line, ",", 2)
+		if len(cols) < 2 {
+			panic("column length should be 2")
+		}
+
+		currency := strings.TrimSpace(cols[0])
+		available := fixedpoint.MustNewFromString(strings.TrimSpace(cols[1]))
+		balances[currency] = Balance(currency, available)
+	}
+
+	return balances
+}
+
+// Balance returns a balance object with the given currency and available amount
+func Balance(currency string, available fixedpoint.Value) types.Balance {
+	return types.Balance{
+		Currency:          currency,
+		Available:         available,
+		Locked:            fixedpoint.Zero,
+		Borrowed:          fixedpoint.Zero,
+		Interest:          fixedpoint.Zero,
+		NetAsset:          fixedpoint.Zero,
+		MaxWithdrawAmount: fixedpoint.Zero,
+	}
+}