From 4aa7cfe9330645ad1cccc8d2a21a3f8e3f36842a Mon Sep 17 00:00:00 2001
From: Elias Van Ootegem <elias@vega.xyz>
Date: Mon, 5 Aug 2024 14:05:36 +0100
Subject: [PATCH] test: fix checkbook logic

Signed-off-by: Elias Van Ootegem <elias@vega.xyz>
---
 core/matching/orderbook.go          |  4 ++-
 core/matching/orderbook_amm_test.go | 38 +++++++++++++++++++++++++----
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/core/matching/orderbook.go b/core/matching/orderbook.go
index 2a67167911..3bed8ad70d 100644
--- a/core/matching/orderbook.go
+++ b/core/matching/orderbook.go
@@ -725,8 +725,9 @@ func (b *OrderBook) CancelAllOrders(party string) ([]*types.OrderCancellationCon
 }
 
 func (b *OrderBook) CheckBook() bool {
-	checkOBB, checkOBS := true, true
+	checkOBB, checkOBS := false, false
 	if len(b.buy.levels) > 0 {
+		checkOBB = true
 		for _, o := range b.buy.levels[len(b.buy.levels)-1].orders {
 			if o.PeggedOrder == nil || o.PeggedOrder.Reference != types.PeggedReferenceBestBid {
 				checkOBB = false
@@ -735,6 +736,7 @@ func (b *OrderBook) CheckBook() bool {
 		}
 	}
 	if len(b.sell.levels) > 0 {
+		checkOBS = true
 		for _, o := range b.sell.levels[len(b.sell.levels)-1].orders {
 			if o.PeggedOrder == nil || o.PeggedOrder.Reference != types.PeggedReferenceBestAsk {
 				checkOBS = false
diff --git a/core/matching/orderbook_amm_test.go b/core/matching/orderbook_amm_test.go
index bf67896e96..567380be51 100644
--- a/core/matching/orderbook_amm_test.go
+++ b/core/matching/orderbook_amm_test.go
@@ -42,17 +42,45 @@ func TestOrderbookAMM(t *testing.T) {
 func testOnlyAMMOrders(t *testing.T) {
 	tst := getTestOrderBookWithAMM(t)
 	defer tst.ctrl.Finish()
+	sellPrice := num.NewUint(111)
+	buyPrice := num.NewUint(95)
+	// create some regular orders
+	sell := createOrder(t, tst, 10, sellPrice)
+	buy := createOrder(t, tst, 10, buyPrice)
+	buy.Side = types.SideBuy
+	tst.obs.EXPECT().SubmitOrder(gomock.Any(), gomock.Any(), gomock.Any()).Times(4).Return(nil)
+	tst.obs.EXPECT().NotifyFinished().Times(3)
+	_, err := tst.book.SubmitOrder(sell)
+	require.NoError(t, err)
+	_, err = tst.book.SubmitOrder(buy)
+	require.NoError(t, err)
+	// create some pegged orders
+	pob := createOrder(t, tst, 100, nil)
+	pob.Party = "B"
+	pob.Side = types.SideBuy
+	pob.PeggedOrder = &types.PeggedOrder{
+		Reference: types.PeggedReferenceBestBid,
+		Offset:    num.NewUint(10),
+	}
+	pos := createOrder(t, tst, 100, nil)
+	pos.Party = "S"
+	pos.PeggedOrder = &types.PeggedOrder{
+		Reference: types.PeggedReferenceBestAsk,
+		Offset:    num.NewUint(10),
+	}
+	require.NoError(t, err)
+	_, err = tst.book.SubmitOrder(pos)
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), tst.book.GetPeggedOrdersCount())
+	// now cancel the non-pegged orders
+	require.NoError(t, err)
 	one, zero := uint64(1), uint64(0)
+
 	// only buy orders
 	tst.obs.EXPECT().BestPricesAndVolumes().Times(1).Return(num.UintOne(), one, nil, zero)
 	check := tst.book.CheckBook()
 	require.False(t, check)
 
-	// only sell orders
-	tst.obs.EXPECT().BestPricesAndVolumes().Times(1).Return(nil, zero, num.UintOne(), one)
-	check = tst.book.CheckBook()
-	require.False(t, check)
-
 	// buy and sell orders
 	tst.obs.EXPECT().BestPricesAndVolumes().Times(1).Return(num.UintOne(), one, num.UintOne(), one)
 	check = tst.book.CheckBook()