From 81f4cff49f462ba0efd9d96bb74dac31c71811bd Mon Sep 17 00:00:00 2001
From: bweick <brian.weickmann@gmail.com>
Date: Mon, 26 Jul 2021 09:39:36 -0700
Subject: [PATCH] Brian/gim trade info update (#115)

Change require to check that component is part of rebalance instead of current component. Allows trade info fetching for new components.
---
 .../protocol/modules/GeneralIndexModule.sol   |  5 ++-
 package.json                                  |  2 +-
 .../modules/generalIndexModule.spec.ts        | 36 ++++++++++++++++---
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/contracts/protocol/modules/GeneralIndexModule.sol b/contracts/protocol/modules/GeneralIndexModule.sol
index 82bd4ade6..f73b7abdd 100644
--- a/contracts/protocol/modules/GeneralIndexModule.sol
+++ b/contracts/protocol/modules/GeneralIndexModule.sol
@@ -576,7 +576,10 @@ contract GeneralIndexModule is ModuleBase, ReentrancyGuard {
         onlyValidAndInitializedSet(_setToken)
         returns (bool, uint256)
     {
-        require(_setToken.isComponent(address(_component)), "Component not recognized");
+        require(
+            rebalanceInfo[_setToken].rebalanceComponents.contains(address(_component)),
+            "Component not recognized"
+        );
         uint256 totalSupply = _setToken.totalSupply();
         return _calculateTradeSizeAndDirection(_setToken, _component, totalSupply);
     }
diff --git a/package.json b/package.json
index a97e44025..b118b8632 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@setprotocol/set-protocol-v2",
-  "version": "0.0.47",
+  "version": "0.0.48",
   "description": "",
   "main": "dist",
   "types": "dist/types",
diff --git a/test/protocol/modules/generalIndexModule.spec.ts b/test/protocol/modules/generalIndexModule.spec.ts
index d4d48abfd..ab08c32d1 100644
--- a/test/protocol/modules/generalIndexModule.spec.ts
+++ b/test/protocol/modules/generalIndexModule.spec.ts
@@ -20,7 +20,8 @@ import {
   ether,
   preciseDiv,
   preciseMul,
-  preciseMulCeil
+  preciseMulCeil,
+  usdc
 } from "@utils/index";
 import {
   cacheBeforeEach,
@@ -2185,9 +2186,9 @@ describe("GeneralIndexModule", () => {
 
       before(async () => {
         // current units [ether(86.9565217), bitcoin(.01111111), ether(100), ZERO]
-        newComponents = [];
+        newComponents = [setup.usdc.address];
         oldTargetUnits = [ether("60.869565780223716593"), bitcoin(.02), ether(55)];
-        newTargetUnits = [];
+        newTargetUnits = [usdc(100)];
         issueAmount = ether("20.000000000000000001");
       });
 
@@ -2208,6 +2209,8 @@ describe("GeneralIndexModule", () => {
       };
 
       beforeEach(async () => {
+        await indexModule.setTradeMaximums(index.address, [setup.usdc.address], [usdc(3000)]);
+
         initializeSubjectVariables();
 
         await startRebalance();
@@ -2242,6 +2245,29 @@ describe("GeneralIndexModule", () => {
         expect(isSendTokenFixed).to.be.true;
       });
 
+      describe("when the component is being added to the Set", async () => {
+        beforeEach(async () => {
+          subjectComponent = setup.usdc.address;
+        });
+
+        it("the correct trade direction and size should be returned", async () => {
+          const totalSupply = await subjectSetToken.totalSupply();
+          const currentUsdcUnit = await subjectSetToken.getDefaultPositionRealUnit(setup.usdc.address);
+          const expectedUsdcSize = preciseDiv(
+            preciseMulCeil(usdc(100), totalSupply).sub(preciseMul(currentUsdcUnit, totalSupply)),
+            PRECISE_UNIT.sub(feePercentage)
+          );
+
+          const [
+            isSendTokenFixed,
+            componentQuantity,
+          ] = await subject();
+
+          expect(componentQuantity).to.eq(expectedUsdcSize);
+          expect(isSendTokenFixed).to.be.false;
+        });
+      });
+
       describe("and the buy component does not meet the max trade size", async () => {
         beforeEach(async () => {
           await indexModule.startRebalance(
@@ -2283,9 +2309,9 @@ describe("GeneralIndexModule", () => {
         });
       });
 
-      describe("when the component is not valid", async () => {
+      describe("when the component is not part of the rebalance", async () => {
         beforeEach(() => {
-          subjectComponent = ADDRESS_ZERO;
+          subjectComponent = setup.weth.address;
         });
 
         it("should revert", async () => {