From e6d2dfe2823e8b51a24c9a8a824b38ce996930de Mon Sep 17 00:00:00 2001
From: alsco77 <alsco77@pm.me>
Date: Sun, 10 Jan 2021 14:10:54 +0100
Subject: [PATCH] Add some basic cases

---
 contracts/savings/SavingsContract.sol    |  7 +--
 test/savings/TestSavingsContract.spec.ts | 63 +++++++++++++++++++++---
 2 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/contracts/savings/SavingsContract.sol b/contracts/savings/SavingsContract.sol
index 6014735d..bc9d427d 100644
--- a/contracts/savings/SavingsContract.sol
+++ b/contracts/savings/SavingsContract.sol
@@ -550,8 +550,8 @@ contract SavingsContract is
             }
 
             // 3. Level the assets to Fraction (connector) & 100-fraction (raw)
-            uint256 realSum = _data.rawBalance.add(connectorBalance);
-            uint256 ideal = realSum.mulTruncate(_data.fraction);
+            uint256 sum = _data.rawBalance.add(connectorBalance);
+            uint256 ideal = sum.mulTruncate(_data.fraction);
             //     If there is not enough mAsset in the connector, then deposit
             if(ideal > connectorBalance){
                 uint256 deposit = ideal.sub(connectorBalance);
@@ -563,6 +563,7 @@ contract SavingsContract is
                 // If fraction == 0, then withdraw everything
                 if(ideal == 0){
                     connector_.withdrawAll();
+                    sum = IERC20(underlying).balanceOf(address(this));
                 } else {
                     connector_.withdraw(connectorBalance.sub(ideal));
                 }
@@ -571,7 +572,7 @@ contract SavingsContract is
 
             // 4i. Refresh exchange rate and emit event
             lastBalance = ideal;
-            _refreshExchangeRate(realSum, _data.totalCredits, false);
+            _refreshExchangeRate(sum, _data.totalCredits, false);
             emit Poked(lastBalance_, ideal, connectorBalance.sub(lastBalance_));
 
         } else {
diff --git a/test/savings/TestSavingsContract.spec.ts b/test/savings/TestSavingsContract.spec.ts
index 16925144..98f01946 100644
--- a/test/savings/TestSavingsContract.spec.ts
+++ b/test/savings/TestSavingsContract.spec.ts
@@ -1059,25 +1059,74 @@ contract("SavingsContract", async (accounts) => {
     });
 
     describe("poking", () => {
-        it("allows only poker to poke");
-        it("fails if there are no credits");
-        it("should fail if the raw balance goes down somehow", async () => {
-            // _refreshExchangeRate
-            // ExchangeRate must increase
+        const deposit = simpleToExactAmount(100, 18);
+        before(async () => {
+            await createNewSavingsContract();
+        });
+        it("allows only poker to poke", async () => {
+            await expectRevert(
+                savingsContract.poke({ from: sa.governor }),
+                "Only poker can execute",
+            );
+        });
+        it("fails if there are no credits", async () => {
+            const credits = await savingsContract.totalSupply();
+            expect(credits).bignumber.eq(new BN(0));
+            await expectRevert(
+                savingsContract.poke({ from: sa.default }),
+                "Must have something to poke",
+            );
+        });
+        it("only allows pokes once every 4h", async () => {
+            await masset.approve(savingsContract.address, simpleToExactAmount(1, 21));
+            await savingsContract.preDeposit(deposit, alice);
+            await savingsContract.poke();
+            await time.increase(ONE_HOUR.muln(3));
+            await expectRevert(
+                savingsContract.poke({ from: sa.default }),
+                "Not enough time elapsed",
+            );
         });
-        it("only allows pokes once every 4h");
         context("with a connector", () => {
+            beforeEach(async () => {
+                await createNewSavingsContract();
+                const connector = await MockConnector.new(savingsContract.address, masset.address);
+
+                await masset.approve(savingsContract.address, simpleToExactAmount(1, 21));
+                await savingsContract.preDeposit(deposit, alice);
+
+                await savingsContract.setConnector(connector.address, { from: sa.governor });
+            });
+            afterEach(async () => {
+                const data = await getData(savingsContract, alice);
+                expect(exchangeRateHolds(data), "Exchange rate must hold");
+            });
             it("should do nothing if the fraction is 0");
+            it("should fail if the raw balance goes down somehow", async () => {
+                // _refreshExchangeRate
+                // ExchangeRate must increase
+            });
             it("should accrue interest and update exchange rate", async () => {
                 // check everything - lastPoke, lastBalance, etc
             });
-            it("should fail if the APY is too high");
+            it("should fail if the APY is too high", async () => {
+                // in 30 mins: "Interest protected from inflating past 10 Bps"
+                //  > 30 mins: "Interest protected from inflating past maxAPY"
+            });
+            it("should fail if the balance has gone down", async () => {
+                // "Invalid yield"
+            });
             it("should deposit to the connector if total supply lowers", async () => {
                 // deposit 1
                 // increase total balance
                 // deposit 2
             });
             it("should withdraw from the connector if total supply lowers");
+
+            // For ex. in yVault and due to slippage on Curve (+ price movement)
+            context("handling unavailable assets in the connector", () => {
+                it("does x");
+            });
         });
         context("with no connector", () => {
             it("simply updates the exchangeRate using the raw balance");