fix(paperwallet): add logic to release locked funds when we cancel or…

…der (#339)
rodrigo-brito · Jul 13, 2024 · cad0986 · cad0986
1 parent 0474603
commit cad0986
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/exchange/paperwallet.go b/exchange/paperwallet.go
@@ -668,6 +668,21 @@ func (p *PaperWallet) Cancel(order model.Order) error {
 	for i, o := range p.orders {
 		if o.ExchangeID == order.ExchangeID {
 			p.orders[i].Status = model.OrderStatusTypeCanceled
+
+			// unlock funds
+			assset, quote := SplitAssetQuote(o.Pair)
+			// we have open long position
+			if p.assets[assset].Lock > 0 && o.Side == model.SideTypeSell {
+				p.assets[assset].Free += o.Quantity
+				p.assets[assset].Lock -= o.Quantity
+			} else {
+				// we don't have open long position
+				if p.assets[assset].Lock == 0 {
+					amount := order.Price * order.Quantity
+					p.assets[quote].Free += amount
+					p.assets[quote].Lock -= amount
+				}
+			}
 		}
 	}
 	return nil

diff --git a/exchange/paperwallet_test.go b/exchange/paperwallet_test.go
@@ -189,6 +189,52 @@ func TestPaperWallet_OrderLimit(t *testing.T) {
 		require.Equal(t, 0.0, wallet.assets["USDT"].Lock)
 		require.Equal(t, 10.0, wallet.avgLongPrice["BTCUSDT"])
 	})
+
+	t.Run("cancel buy order before executing", func(t *testing.T) {
+		wallet := NewPaperWallet(context.Background(), "USDT", WithPaperAsset("USDT", 100))
+		order, err := wallet.CreateOrderLimit(model.SideTypeBuy, "BTCUSDT", 1, 100)
+		require.NoError(t, err)
+
+		// create order and lock values
+		require.Len(t, wallet.orders, 1)
+		require.Equal(t, 1.0, order.Quantity)
+		require.Equal(t, 100.0, order.Price)
+		require.Equal(t, 0.0, wallet.assets["USDT"].Free)
+		require.Equal(t, 100.0, wallet.assets["USDT"].Lock)
+
+		// cancel limit order and it should unlock funds
+		err = wallet.Cancel(order)
+		require.NoError(t, err)
+
+		require.Equal(t, model.OrderStatusTypeCanceled, wallet.orders[0].Status)
+		require.Equal(t, 100.0, wallet.assets["USDT"].Free)
+		require.Equal(t, 0.0, wallet.assets["USDT"].Lock)
+		require.Equal(t, 0.0, wallet.assets["BTC"].Free)
+		require.Equal(t, 0.0, wallet.assets["BTC"].Lock)
+	})
+
+	t.Run("cancel sell order before executing", func(t *testing.T) {
+		wallet := NewPaperWallet(context.Background(), "USDT", WithPaperAsset("USDT", 100))
+		order, err := wallet.CreateOrderLimit(model.SideTypeSell, "BTCUSDT", 1, 100)
+		require.NoError(t, err)
+
+		// create order and lock values
+		require.Len(t, wallet.orders, 1)
+		require.Equal(t, 1.0, order.Quantity)
+		require.Equal(t, 100.0, order.Price)
+		require.Equal(t, 0.0, wallet.assets["USDT"].Free)
+		require.Equal(t, 100.0, wallet.assets["USDT"].Lock)
+
+		// cancel limit order and it should unlock funds
+		err = wallet.Cancel(order)
+		require.NoError(t, err)
+
+		require.Equal(t, model.OrderStatusTypeCanceled, wallet.orders[0].Status)
+		require.Equal(t, 100.0, wallet.assets["USDT"].Free)
+		require.Equal(t, 0.0, wallet.assets["USDT"].Lock)
+		require.Equal(t, 0.0, wallet.assets["BTC"].Free)
+		require.Equal(t, 0.0, wallet.assets["BTC"].Lock)
+	})
 }
 
 func TestPaperWallet_OrderMarket(t *testing.T) {