fix test

dydxprotocol · Aug 20, 2024 · 0c550b3 · 0c550b3
1 parent 374379a
commit 0c550b3
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 98 deletions.
diff --git a/protocol/daemons/liquidation/client/sub_task_runner_test.go b/protocol/daemons/liquidation/client/sub_task_runner_test.go
@@ -88,8 +88,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Can get liquidatable subaccount with long position": {
@@ -148,8 +147,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Skip well collateralized subaccounts": {
@@ -209,8 +207,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Skip subaccounts with no open positions": {
@@ -259,8 +256,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 					NegativeTncSubaccountIds:   []satypes.SubaccountId{},
 					SubaccountOpenPositionInfo: []clobtypes.SubaccountOpenPositionInfo{},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Can get subaccount that become undercollateralized with funding payments (short)": {
@@ -340,8 +336,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Can get subaccount that become liquidatable with funding payments (long)": {
@@ -421,8 +416,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Skips subaccount that become well-collateralized with funding payments (short)": {
@@ -500,8 +494,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Skips subaccount that become well-collateralized with funding payments (long)": {
@@ -579,8 +572,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Can get negative tnc subaccount with short position": {
@@ -642,8 +634,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 		"Can get negative tnc subaccount with long position": {
@@ -705,8 +696,7 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 						},
 					},
 				}
-				response3 := &api.LiquidateSubaccountsResponse{}
-				mck.On("LiquidateSubaccounts", ctx, req).Return(response3, nil)
+				setupMockLiquidateSubaccountRequests(mck, ctx, req)
 			},
 		},
 	}
@@ -739,3 +729,58 @@ func TestRunLiquidationDaemonTaskLoop(t *testing.T) {
 		})
 	}
 }
+
+func setupMockLiquidateSubaccountRequests(
+	mck *mocks.QueryClient,
+	ctx context.Context,
+	request *api.LiquidateSubaccountsRequest,
+) {
+	req := &api.LiquidateSubaccountsRequest{
+		BlockHeight:               request.BlockHeight,
+		LiquidatableSubaccountIds: request.LiquidatableSubaccountIds,
+	}
+	response := &api.LiquidateSubaccountsResponse{}
+	mck.On("LiquidateSubaccounts", ctx, req).Return(response, nil)
+
+	req = &api.LiquidateSubaccountsRequest{
+		BlockHeight:              request.BlockHeight,
+		NegativeTncSubaccountIds: request.NegativeTncSubaccountIds,
+	}
+	mck.On("LiquidateSubaccounts", ctx, req).Return(response, nil)
+
+	if len(request.SubaccountOpenPositionInfo) == 0 {
+		req = &api.LiquidateSubaccountsRequest{
+			BlockHeight:                request.BlockHeight,
+			SubaccountOpenPositionInfo: []clobtypes.SubaccountOpenPositionInfo{},
+		}
+		mck.On("LiquidateSubaccounts", ctx, req).Return(response, nil)
+	} else {
+		for _, info := range request.SubaccountOpenPositionInfo {
+			if len(info.SubaccountsWithLongPosition) > 0 {
+				req = &api.LiquidateSubaccountsRequest{
+					BlockHeight: request.BlockHeight,
+					SubaccountOpenPositionInfo: []clobtypes.SubaccountOpenPositionInfo{
+						{
+							PerpetualId:                 info.PerpetualId,
+							SubaccountsWithLongPosition: info.SubaccountsWithLongPosition,
+						},
+					},
+				}
+				mck.On("LiquidateSubaccounts", ctx, req).Return(response, nil)
+			}
+
+			if len(info.SubaccountsWithShortPosition) > 0 {
+				req = &api.LiquidateSubaccountsRequest{
+					BlockHeight: request.BlockHeight,
+					SubaccountOpenPositionInfo: []clobtypes.SubaccountOpenPositionInfo{
+						{
+							PerpetualId:                  info.PerpetualId,
+							SubaccountsWithShortPosition: info.SubaccountsWithShortPosition,
+						},
+					},
+				}
+				mck.On("LiquidateSubaccounts", ctx, req).Return(response, nil)
+			}
+		}
+	}
+}
diff --git a/protocol/daemons/server/types/liquidations/daemon_liquidation_info.go b/protocol/daemons/server/types/liquidations/daemon_liquidation_info.go
@@ -34,10 +34,19 @@ func (ls *DaemonLiquidationInfo) Update(
 ) {
 	ls.Lock()
 	defer ls.Unlock()
+
+	if blockHeight > ls.blockHeight {
+		ls.liquidatableSubaccountIds = make([]satypes.SubaccountId, 0)
+		ls.negativeTncSubaccountIds = make([]satypes.SubaccountId, 0)
+		ls.subaccountsWithPositions = make(map[uint32]*clobtypes.SubaccountOpenPositionInfo)
+	} else if blockHeight < ls.blockHeight {
+		panic("UpdateLiquidatableSubaccountIds: block height cannot be less than the current block height")
+	}
+	ls.UpdateBlockHeight(blockHeight)
+
 	ls.UpdateLiquidatableSubaccountIds(liquidatableSubaccountIds, blockHeight)
 	ls.UpdateNegativeTncSubaccountIds(negativeTncSubaccountIds, blockHeight)
 	ls.UpdateSubaccountsWithPositions(subaccountsWithPositions, blockHeight)
-	ls.UpdateBlockHeight(blockHeight)
 }
 
 // UpdateBlockHeight updates the struct with the given block height.
@@ -51,14 +60,7 @@ func (ls *DaemonLiquidationInfo) UpdateLiquidatableSubaccountIds(
 	updates []satypes.SubaccountId,
 	blockHeight uint32,
 ) {
-	if blockHeight > ls.blockHeight {
-		ls.liquidatableSubaccountIds = make([]satypes.SubaccountId, len(updates))
-		copy(ls.liquidatableSubaccountIds, updates)
-	} else if blockHeight == ls.blockHeight {
-		ls.liquidatableSubaccountIds = append(ls.liquidatableSubaccountIds, updates...)
-	} else {
-		panic("UpdateLiquidatableSubaccountIds: block height cannot be less than the current block height")
-	}
+	ls.liquidatableSubaccountIds = append(ls.liquidatableSubaccountIds, updates...)
 }
 
 // UpdateNegativeTncSubaccountIds updates the struct with the given a list of subaccount ids
@@ -67,55 +69,31 @@ func (ls *DaemonLiquidationInfo) UpdateNegativeTncSubaccountIds(
 	updates []satypes.SubaccountId,
 	blockHeight uint32,
 ) {
-	if blockHeight > ls.blockHeight {
-		ls.negativeTncSubaccountIds = make([]satypes.SubaccountId, len(updates))
-		copy(ls.negativeTncSubaccountIds, updates)
-	} else if blockHeight == ls.blockHeight {
-		ls.negativeTncSubaccountIds = append(ls.negativeTncSubaccountIds, updates...)
-	} else {
-		panic("UpdateLiquidatableSubaccountIds: block height cannot be less than the current block height")
-	}
+	ls.negativeTncSubaccountIds = append(ls.negativeTncSubaccountIds, updates...)
 }
 
 // UpdateSubaccountsWithPositions updates the struct with the given a list of subaccount ids with open positions.
 func (ls *DaemonLiquidationInfo) UpdateSubaccountsWithPositions(
 	subaccountsWithPositions []clobtypes.SubaccountOpenPositionInfo,
 	blockHeight uint32,
 ) {
-	if blockHeight > ls.blockHeight {
-		// Reset the map if the block height has changed.
-		ls.subaccountsWithPositions = make(map[uint32]*clobtypes.SubaccountOpenPositionInfo)
-		for _, info := range subaccountsWithPositions {
-			clone := &clobtypes.SubaccountOpenPositionInfo{
+	// Append to the current map if the block height not changed.
+	for _, info := range subaccountsWithPositions {
+		if _, ok := ls.subaccountsWithPositions[info.PerpetualId]; !ok {
+			ls.subaccountsWithPositions[info.PerpetualId] = &clobtypes.SubaccountOpenPositionInfo{
 				PerpetualId:                  info.PerpetualId,
-				SubaccountsWithLongPosition:  make([]satypes.SubaccountId, len(info.SubaccountsWithLongPosition)),
-				SubaccountsWithShortPosition: make([]satypes.SubaccountId, len(info.SubaccountsWithShortPosition)),
-			}
-			copy(clone.SubaccountsWithLongPosition, info.SubaccountsWithLongPosition)
-			copy(clone.SubaccountsWithShortPosition, info.SubaccountsWithShortPosition)
-			ls.subaccountsWithPositions[info.PerpetualId] = clone
-		}
-	} else if blockHeight == ls.blockHeight {
-		// Append to the current map if the block height not changed.
-		for _, info := range subaccountsWithPositions {
-			if _, ok := ls.subaccountsWithPositions[info.PerpetualId]; !ok {
-				ls.subaccountsWithPositions[info.PerpetualId] = &clobtypes.SubaccountOpenPositionInfo{
-					PerpetualId:                  info.PerpetualId,
-					SubaccountsWithLongPosition:  make([]satypes.SubaccountId, 0),
-					SubaccountsWithShortPosition: make([]satypes.SubaccountId, 0),
-				}
+				SubaccountsWithLongPosition:  make([]satypes.SubaccountId, 0),
+				SubaccountsWithShortPosition: make([]satypes.SubaccountId, 0),
 			}
-			ls.subaccountsWithPositions[info.PerpetualId].SubaccountsWithLongPosition = append(
-				ls.subaccountsWithPositions[info.PerpetualId].SubaccountsWithLongPosition,
-				info.SubaccountsWithLongPosition...,
-			)
-			ls.subaccountsWithPositions[info.PerpetualId].SubaccountsWithShortPosition = append(
-				ls.subaccountsWithPositions[info.PerpetualId].SubaccountsWithShortPosition,
-				info.SubaccountsWithShortPosition...,
-			)
 		}
-	} else {
-		panic("UpdateLiquidatableSubaccountIds: block height cannot be less than the current block height")
+		ls.subaccountsWithPositions[info.PerpetualId].SubaccountsWithLongPosition = append(
+			ls.subaccountsWithPositions[info.PerpetualId].SubaccountsWithLongPosition,
+			info.SubaccountsWithLongPosition...,
+		)
+		ls.subaccountsWithPositions[info.PerpetualId].SubaccountsWithShortPosition = append(
+			ls.subaccountsWithPositions[info.PerpetualId].SubaccountsWithShortPosition,
+			info.SubaccountsWithShortPosition...,
+		)
 	}
 }