From 8ca866ae980dc855535ac388b5c5176175ab6bbf Mon Sep 17 00:00:00 2001
From: Atul Agarwal <21087753+asquare08@users.noreply.github.com>
Date: Fri, 22 Mar 2024 02:35:53 +0530
Subject: [PATCH] add test for getRequiredMargin

---
 .../contracts/juror/matching_validation.go    |  2 +-
 .../traderviewer/limit_orders_test.go         | 41 +++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/precompile/contracts/juror/matching_validation.go b/precompile/contracts/juror/matching_validation.go
index b993b71483..75354879e3 100644
--- a/precompile/contracts/juror/matching_validation.go
+++ b/precompile/contracts/juror/matching_validation.go
@@ -443,7 +443,7 @@ func getRequiredMargin(bibliophile b.BibliophileClient, order ILimitOrderBookOrd
 
 func getRequiredMarginNew(bibliophile b.BibliophileClient, baseAsset *big.Int, price *big.Int, marketId int64, trader *common.Address) *big.Int {
 	marketAddress := bibliophile.GetMarketAddressFromMarketID(marketId)
-	marginFraction := bibliophile.GetTraderMarginFraction(marketAddress,trader)
+	marginFraction := bibliophile.GetTraderMarginFraction(marketAddress, trader)
 	quoteAsset := hu.Div1e18(hu.Mul(hu.Abs(baseAsset), price))
 	return hu.Div1e6(hu.Mul(quoteAsset, marginFraction))
 }
diff --git a/precompile/contracts/traderviewer/limit_orders_test.go b/precompile/contracts/traderviewer/limit_orders_test.go
index 28936fe421..e46f8e41b3 100644
--- a/precompile/contracts/traderviewer/limit_orders_test.go
+++ b/precompile/contracts/traderviewer/limit_orders_test.go
@@ -438,3 +438,44 @@ func getOrderHash(order ILimitOrderBookOrder) common.Hash {
 	}
 	return orderHash
 }
+
+func TestGetRequiredMargin(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	mockBibliophile := b.NewMockBibliophileClient(ctrl)
+	ammIndex := big.NewInt(0)
+	longBaseAssetQuantity := big.NewInt(5000000000000000000)
+	shortBaseAssetQuantity := big.NewInt(-5000000000000000000)
+	price := big.NewInt(100000000)
+	trader := common.HexToAddress("0x3C44CdDdB6a900fa2b585dd299e03d12FA4293BC")
+	ammAddress := common.HexToAddress("0x70997970C51812dc3A010C7d01b50e0d17dc79C8")
+	marginFraction := big.NewInt(2e5) // 0.2
+	t.Run("when order is long", func(t *testing.T) {
+		input := getRequiredMarginInput(longBaseAssetQuantity, price, ammIndex, trader)
+		mockBibliophile.EXPECT().GetMarketAddressFromMarketID(ammIndex.Int64()).Return(ammAddress).Times(1)
+		mockBibliophile.EXPECT().GetTraderMarginFraction(ammAddress, &trader).Return(marginFraction).Times(1)
+		output := GetRequiredMargin(mockBibliophile, &input)
+		expectedMargin := hu.Div1e18(hu.Mul(longBaseAssetQuantity, price))
+		expectedMargin = hu.Div1e6(hu.Mul(expectedMargin, marginFraction))
+		assert.Equal(t, expectedMargin, output)
+	})
+
+	t.Run("when order is short", func(t *testing.T) {
+		input := getRequiredMarginInput(shortBaseAssetQuantity, price, ammIndex, trader)
+		mockBibliophile.EXPECT().GetMarketAddressFromMarketID(ammIndex.Int64()).Return(ammAddress).Times(1)
+		mockBibliophile.EXPECT().GetTraderMarginFraction(ammAddress, &trader).Return(marginFraction).Times(1)
+		output := GetRequiredMargin(mockBibliophile, &input)
+		expectedMargin := hu.Div1e18(hu.Mul(hu.Abs(shortBaseAssetQuantity), price))
+		expectedMargin = hu.Div1e6(hu.Mul(expectedMargin, marginFraction))
+		assert.Equal(t, expectedMargin, output)
+	})
+}
+
+func getRequiredMarginInput(baseAssetQuantity *big.Int, price *big.Int, ammIndex *big.Int, trader common.Address) GetRequiredMarginInput {
+	return GetRequiredMarginInput{
+		BaseAssetQuantity: baseAssetQuantity,
+		Price:             price,
+		AmmIndex:          ammIndex,
+		Trader:            trader,
+	}
+}