grid2: add basic investment check test checkRequiredInvestmentByQuantity

pkg/strategy/grid2/strategy.go

@@ -176,38 +176,54 @@ func (s *Strategy) Run(ctx context.Context, orderExecutor bbgo.OrderExecutor, se
 	return nil
 }
 
-func (s *Strategy) checkRequiredInvestmentByQuantity(baseInvestment, quoteInvestment, totalBaseBalance, totalQuoteBalance, quantity, lastPrice fixedpoint.Value, pins []Pin) error {
+type InvestmentBudget struct {
+	baseInvestment  fixedpoint.Value
+	quoteInvestment fixedpoint.Value
+	baseBalance     fixedpoint.Value
+	quoteBalance    fixedpoint.Value
+}
+
+func (s *Strategy) checkRequiredInvestmentByQuantity(baseInvestment, quoteInvestment, baseBalance, quoteBalance, quantity, lastPrice fixedpoint.Value, pins []Pin) error {
+	if baseInvestment.Compare(baseBalance) > 0 {
+		return fmt.Errorf("baseInvestment setup %f is greater than the total base balance %f", baseInvestment.Float64(), baseBalance.Float64())
+	}
+
+	if quoteInvestment.Compare(quoteBalance) > 0 {
+		return fmt.Errorf("quoteInvestment setup %f is greater than the total quote balance %f", quoteInvestment.Float64(), quoteBalance.Float64())
+	}
+
+	// check more investment budget details
 	requiredBase := fixedpoint.Zero
 	requiredQuote := fixedpoint.Zero
-	for i := len(s.grid.Pins) - 1; i >= 0; i++ {
+	for i := len(pins) - 1; i >= 0; i++ {
-		pin := s.grid.Pins[i]
+		pin := pins[i]
 		price := fixedpoint.Value(pin)
 
-		// TODO: add fee if we don't have the platform token. BNB, OKEX or MAX...
+		// TODO: add fee if we don't have the platform token. BNB, OKB or MAX...
 		if price.Compare(lastPrice) >= 0 {
 			// for orders that sell
 			// if we still have the base balance
-			if requiredBase.Compare(totalBaseBalance) < 0 {
+			if requiredBase.Compare(baseBalance) < 0 {
 				requiredBase = requiredBase.Add(quantity)
-			} else if i > 0 {
+			} else if i > 0 { // we do not want to sell at i == 0
-				// convert buy quote to requiredQuote
+				// convert sell to buy quote and add to requiredQuote
 				nextLowerPin := s.grid.Pins[i-1]
 				nextLowerPrice := fixedpoint.Value(nextLowerPin)
 				requiredQuote = requiredQuote.Add(quantity.Mul(nextLowerPrice))
-
 			}
 		} else {
+			// for orders that buy
 			requiredQuote = requiredQuote.Add(quantity.Mul(price))
 		}
 	}
 
-	if requiredBase.Compare(totalBaseBalance) < 0 && requiredQuote.Compare(totalQuoteBalance) < 0 {
+	if requiredBase.Compare(baseBalance) > 0 && requiredQuote.Compare(quoteBalance) > 0 {
 		return fmt.Errorf("both base balance (%f %s) and quote balance (%f %s) are not enought",
-			totalBaseBalance.Float64(), s.Market.BaseCurrency,
+			baseBalance.Float64(), s.Market.BaseCurrency,
-			totalQuoteBalance.Float64(), s.Market.QuoteCurrency)
+			quoteBalance.Float64(), s.Market.QuoteCurrency)
 	}
 
-	if requiredBase.Compare(totalBaseBalance) < 0 {
+	if requiredBase.Compare(baseBalance) < 0 {
 		// see if we can convert some quotes to base
 	}
 

pkg/strategy/grid2/strategy_test.go

@@ -0,0 +1,27 @@
+package grid2
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestStrategy_checkRequiredInvestmentByQuantity(t *testing.T) {
+	s := &Strategy{}
+
+	t.Run("basic base balance check", func(t *testing.T) {
+		err := s.checkRequiredInvestmentByQuantity(number(2.0), number(10_000.0),
+			number(1.0), number(10_000.0),
+			number(0.1), number(19000.0), []Pin{})
+		assert.Error(t, err)
+		assert.EqualError(t, err, "baseInvestment setup 2.000000 is greater than the total base balance 1.000000")
+	})
+
+	t.Run("basic quote balance check", func(t *testing.T) {
+		err := s.checkRequiredInvestmentByQuantity(number(1.0), number(10_000.0),
+			number(1.0), number(100.0),
+			number(0.1), number(19_000.0), []Pin{})
+		assert.Error(t, err)
+		assert.EqualError(t, err, "quoteInvestment setup 10000.000000 is greater than the total quote balance 100.000000")
+	})
+}