Fix/borrow logic fix

compiled_contracts/Comptroller/step_000_cont_0_sizes.csv

@@ -1,2 +1,2 @@
-storage,21967
+storage,22007
 contract,7159

contracts/Comptroller.py

@@ -224,7 +224,8 @@ def borrowAllowed(self, params):
             # only cTokens may call borrowAllowed if borrower not in market
             sp.verify(sp.sender == params.cToken,
                       EC.CMPT_INVALID_BORROW_SENDER)
-            self.addToLoans(sp.sender, params.borrower)
+        sp.if sp.sender == params.cToken:
+            self.addToLoans(params.cToken, params.borrower)
         self.checkInsuffLiquidityInternal(
             params.cToken, params.borrower, params.borrowAmount)
         self.invalidateLiquidity(params.borrower)

contracts/tests/CTokenTest.py

@@ -99,8 +99,36 @@ def test():
                         "symbol": sp.utils.bytes_of_string("cToken"),
                         "decimals": sp.utils.bytes_of_string("x"),
                     })
+    c2 = TestCToken(comptroller_=cmpt.address,
+                    interestRateModel_=irm.address,
+                    initialExchangeRateMantissa_=sp.nat(exchange_rate),
+                    administrator_=admin.address,
+                    metadata_=sp.big_map({
+                        "": sp.utils.bytes_of_string("tezos-storage:data"),
+                        "data": sp.utils.bytes_of_string(json.dumps({
+                            "name": "...",
+                            "description": "...",
+                            "version": "1.0.0",
+                            "authors": ["ewqenqjw"],
+                            "homepage": "https://some-website.com",
+                            "interfaces": ["TZIP-007"],
+                            "license": {"name": "..."}
+                        }))
+                    }),
+                    token_metadata_={
+                        "name": sp.utils.bytes_of_string("Compound token"),
+                        "symbol": sp.utils.bytes_of_string("cToken"),
+                        "decimals": sp.utils.bytes_of_string("x"),
+                    })
 
     scenario += c1
+    scenario += c2
+
+    scenario.h2("Set initial state")
+    scenario += cmpt.enterMarkets(c1.address).run(sender=alice, level=bLevel.current())
+    scenario += cmpt.enterMarkets(c2.address).run(sender=alice, level=bLevel.current())
+    scenario.verify(cmpt.data.markets.contains(c1.address))
+    scenario.verify(cmpt.data.markets.contains(c2.address))
 
     scenario.h2("Try borrow when there is no cash")
     scenario += c1.borrow(100).run(sender=alice, level=bLevel.current(), valid=False)
@@ -147,6 +175,15 @@ def test():
     scenario += c1.borrow(1100 * ctoken_decimals + 1).run(sender=carl, level=bLevel.current(), valid=False)
     scenario.h3("Try borrow in callback")
     scenario += c1.getCashLegacy(sp.pair(sp.unit, c1.typed.borrow)).run(sender=alice, level=bLevel.current(), valid=False)
+    scenario.h3("Add liquidity to c2 for borrowing tests")
+    DataRelevance.updateAccrueInterest(scenario, bLevel, bob, c2)
+    scenario += c2.mint(1000).run(sender=bob, level=bLevel.current()) 
+    scenario.h3("Verify that loans are added correctly")
+    DataRelevance.updateAllRelevance(scenario, bLevel, carl, c2, cmpt, c2.address, carl.address)
+    scenario += cmpt.setBorrowAllowed(sp.bool(True))
+    scenario += c2.borrow(10).run(sender=alice, level=bLevel.current())
+    scenario.verify(cmpt.data.loans[alice.address].contains(c1.address))
+    scenario.verify(cmpt.data.loans[alice.address].contains(c2.address))
 
     scenario.h2("Test Redeem")
     scenario.h3("Redeem allowed")

contracts/tests/mock/ComptrollerMock.py

@@ -11,12 +11,16 @@ def __init__(self, **extra_storage):
             borrow_allowed = sp.bool(True),
             redeem_allowed = sp.bool(True),
             repay_borrow_allowed = sp.bool(True),
+            markets = sp.set(t=sp.TAddress),
+            loans=sp.big_map(l={}, tkey=sp.TAddress,
+                             tvalue=sp.TSet(sp.TAddress)),
             **extra_storage
         )
 
     @sp.entry_point
     def enterMarkets(self, params):
-        sp.set_type(params, sp.TUnit)
+        sp.set_type(params, sp.TAddress)
+        self.data.markets.add(params)
 
     @sp.entry_point
     def exitMarket(self, params):
@@ -44,7 +48,16 @@ def setRedeemAllowed(self, params):
     def borrowAllowed(self, params):
         sp.set_type(params, CMPTInterface.TBorrowAllowedParams)
         sp.verify(self.data.borrow_allowed)
-
+        sp.if ~ self.data.loans.contains(params.borrower):
+            # only cTokens may call borrowAllowed if borrower not in market
+            sp.verify(sp.sender == params.cToken,
+                     "Invalid sender for borrowAllowed")
+        sp.if sp.sender == params.cToken:
+            sp.if self.data.loans.contains(params.borrower):
+                self.data.loans[params.borrower].add(params.cToken)
+            sp.else:
+                self.data.loans[params.borrower] = sp.set([params.cToken])
+
     @sp.entry_point
     def setBorrowAllowed(self, params):
         self.data.borrow_allowed = params