fix balance map add

This commit is contained in:
c9s 2022-05-04 17:39:35 +08:00
parent 75adb8f3c3
commit 4404098bf9
No known key found for this signature in database
GPG Key ID: 7385E7E464CB0A54
2 changed files with 30 additions and 11 deletions

View File

@ -19,6 +19,23 @@ func TestBalanceMap_Add(t *testing.T) {
}, },
}) })
assert.Len(t, bm2, 1) assert.Len(t, bm2, 1)
var bm3 = bm2.Add(BalanceMap{
"BTC": Balance{
Currency: "BTC",
Available: fixedpoint.MustNewFromString("1.0"),
Locked: fixedpoint.MustNewFromString("0"),
NetAsset: fixedpoint.MustNewFromString("1.0"),
},
"LTC": Balance{
Currency: "LTC",
Available: fixedpoint.MustNewFromString("20.0"),
Locked: fixedpoint.MustNewFromString("0"),
NetAsset: fixedpoint.MustNewFromString("20.0"),
},
})
assert.Len(t, bm3, 2)
assert.Equal(t, fixedpoint.MustNewFromString("11.0"), bm3["BTC"].Available)
} }
func TestAccountLockAndUnlock(t *testing.T) { func TestAccountLockAndUnlock(t *testing.T) {

View File

@ -41,7 +41,6 @@ func (b Balance) String() (o string) {
return o return o
} }
type BalanceMap map[string]Balance type BalanceMap map[string]Balance
func (m BalanceMap) Currencies() (currencies []string) { func (m BalanceMap) Currencies() (currencies []string) {
@ -52,14 +51,18 @@ func (m BalanceMap) Currencies() (currencies []string) {
} }
func (m BalanceMap) Add(bm BalanceMap) BalanceMap { func (m BalanceMap) Add(bm BalanceMap) BalanceMap {
var total = BalanceMap{} var total = m.Copy()
for _, b := range bm { for _, b := range bm {
tb := total[b.Currency] tb, ok := total[b.Currency]
tb.Available = tb.Available.Add(b.Available) if ok {
tb.Locked = tb.Locked.Add(b.Locked) tb.Available = tb.Available.Add(b.Available)
tb.Borrowed = tb.Borrowed.Add(b.Borrowed) tb.Locked = tb.Locked.Add(b.Locked)
tb.NetAsset = tb.NetAsset.Add(b.NetAsset) tb.Borrowed = tb.Borrowed.Add(b.Borrowed)
tb.Interest = tb.Interest.Add(b.Interest) tb.NetAsset = tb.NetAsset.Add(b.NetAsset)
tb.Interest = tb.Interest.Add(b.Interest)
} else {
tb = b
}
total[b.Currency] = tb total[b.Currency] = tb
} }
return total return total
@ -107,14 +110,14 @@ func (m BalanceMap) Assets(prices map[string]fixedpoint.Value, priceTime time.Ti
NetAsset: netAsset, NetAsset: netAsset,
} }
if strings.HasPrefix(currency,"USD") { // for usd if strings.HasPrefix(currency, "USD") { // for usd
asset.InUSD = netAsset asset.InUSD = netAsset
asset.PriceInUSD = fixedpoint.One asset.PriceInUSD = fixedpoint.One
if hasBtcPrice && !asset.InUSD.IsZero() { if hasBtcPrice && !asset.InUSD.IsZero() {
asset.InBTC = asset.InUSD.Div(btcusdt) asset.InBTC = asset.InUSD.Div(btcusdt)
} }
} else { // for crypto } else { // for crypto
if market, usdPrice, ok := findUSDMarketPrice(currency, prices) ; ok { if market, usdPrice, ok := findUSDMarketPrice(currency, prices); ok {
// this includes USDT, USD, USDC and so on // this includes USDT, USD, USDC and so on
if strings.HasPrefix(market, "USD") { if strings.HasPrefix(market, "USD") {
if !asset.Total.IsZero() { if !asset.Total.IsZero() {
@ -154,7 +157,6 @@ func (m BalanceMap) Print() {
} }
} }
func findUSDMarketPrice(currency string, prices map[string]fixedpoint.Value) (string, fixedpoint.Value, bool) { func findUSDMarketPrice(currency string, prices map[string]fixedpoint.Value) (string, fixedpoint.Value, bool) {
usdMarkets := []string{currency + "USDT", currency + "USDC", currency + "USD", "USDT" + currency} usdMarkets := []string{currency + "USDT", currency + "USDC", currency + "USD", "USDT" + currency}
for _, market := range usdMarkets { for _, market := range usdMarkets {