Clean up driver tests (#321)

drips-network · Oct 11, 2023 · 19e7b67 · 19e7b67
1 parent 394a634
commit 19e7b67
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 17 deletions.
diff --git a/test/AddressDriver.t.sol b/test/AddressDriver.t.sol
@@ -38,10 +38,9 @@ contract AddressDriverTest is Test {
         // Make AddressDriver's driver ID non-0 to test if it's respected by AddressDriver
         drips.registerDriver(address(1));
         drips.registerDriver(address(1));
-        uint32 driverId = drips.registerDriver(address(this));
-        AddressDriver driverLogic = new AddressDriver(drips, address(caller), driverId);
+        AddressDriver driverLogic = new AddressDriver(drips, address(caller), drips.nextDriverId());
         driver = AddressDriver(address(new ManagedProxy(driverLogic, admin)));
-        drips.updateDriverAddress(driverId, address(driver));
+        drips.registerDriver(address(driver));
 
         thisId = driver.calcAccountId(address(this));
         accountId = driver.calcAccountId(user);
@@ -108,11 +107,10 @@ contract AddressDriverTest is Test {
 
         assertEq(erc20.balanceOf(address(this)), balance - amt, "Invalid balance after top-up");
         assertEq(erc20.balanceOf(address(drips)), amt, "Invalid Drips balance after top-up");
-        (,,, uint128 streamsBalance,) = drips.streamsState(thisId, erc20);
+        (bytes32 streamsHash,,, uint128 streamsBalance,) = drips.streamsState(thisId, erc20);
+        assertEq(streamsHash, drips.hashStreams(receivers), "Invalid streams hash after top-up");
         assertEq(streamsBalance, amt, "Invalid streams balance after top-up");
         assertEq(realBalanceDelta, int128(amt), "Invalid streams balance delta after top-up");
-        (bytes32 streamsHash,,,,) = drips.streamsState(thisId, erc20);
-        assertEq(streamsHash, drips.hashStreams(receivers), "Invalid streams hash after top-up");
 
         // Withdraw
         balance = erc20.balanceOf(address(user));

diff --git a/test/NFTDriver.t.sol b/test/NFTDriver.t.sol
@@ -43,10 +43,9 @@ contract NFTDriverTest is Test {
         // Make NFTDriver's driver ID non-0 to test if it's respected by NFTDriver
         drips.registerDriver(address(1));
         drips.registerDriver(address(1));
-        uint32 driverId = drips.registerDriver(address(this));
-        NFTDriver driverLogic = new NFTDriver(drips, address(caller), driverId);
+        NFTDriver driverLogic = new NFTDriver(drips, address(caller), drips.nextDriverId());
         driver = NFTDriver(address(new ManagedProxy(driverLogic, admin)));
-        drips.updateDriverAddress(driverId, address(driver));
+        drips.registerDriver(address(driver));
 
         tokenId = driver.mint(address(this), noMetadata());
         tokenId1 = driver.mint(address(this), noMetadata());
@@ -219,7 +218,6 @@ contract NFTDriverTest is Test {
         uint128 amt = 5;
 
         // Top-up
-
         StreamReceiver[] memory receivers = new StreamReceiver[](1);
         receivers[0] =
             StreamReceiver(tokenId2, StreamConfigImpl.create(0, drips.minAmtPerSec(), 0, 0));
@@ -231,12 +229,10 @@ contract NFTDriverTest is Test {
 
         assertEq(erc20.balanceOf(address(this)), balance - amt, "Invalid balance after top-up");
         assertEq(erc20.balanceOf(address(drips)), amt, "Invalid Drips balance after top-up");
-        (,,, uint128 streamsBalance,) = drips.streamsState(tokenId1, erc20);
+        (bytes32 streamsHash,,, uint128 streamsBalance,) = drips.streamsState(tokenId1, erc20);
+        assertEq(streamsHash, drips.hashStreams(receivers), "Invalid streams hash after top-up");
         assertEq(streamsBalance, amt, "Invalid streams balance after top-up");
-
         assertEq(realBalanceDelta, int128(amt), "Invalid streams balance delta after top-up");
-        (bytes32 streamsHash,,,,) = drips.streamsState(tokenId1, erc20);
-        assertEq(streamsHash, drips.hashStreams(receivers), "Invalid streams hash after top-up");
 
         // Withdraw
         balance = erc20.balanceOf(address(user));
@@ -338,6 +334,14 @@ contract NFTDriverTest is Test {
         driver.safeMint(user, noMetadata());
     }
 
+    function testMintWithSaltCanBePaused() public canBePausedTest {
+        driver.mintWithSalt(0, user, noMetadata());
+    }
+
+    function testSafeMintWithSaltCanBePaused() public canBePausedTest {
+        driver.safeMintWithSalt(0, user, noMetadata());
+    }
+
     function testCollectCanBePaused() public canBePausedTest {
         driver.collect(0, erc20, user);
     }

diff --git a/test/RepoDriver.t.sol b/test/RepoDriver.t.sol
@@ -568,6 +568,7 @@ contract RepoDriverTest is Test {
 
     function testSetStreams() public {
         uint128 amt = 5;
+
         // Top-up
         StreamReceiver[] memory receivers = new StreamReceiver[](1);
         receivers[0] =
@@ -578,11 +579,11 @@ contract RepoDriverTest is Test {
         );
         assertEq(erc20.balanceOf(address(this)), balance - amt, "Invalid balance after top-up");
         assertEq(erc20.balanceOf(address(drips)), amt, "Invalid Drips balance after top-up");
-        (,,, uint128 streamsBalance,) = drips.streamsState(accountId1, erc20);
+        (bytes32 streamsHash,,, uint128 streamsBalance,) = drips.streamsState(accountId1, erc20);
+        assertEq(streamsHash, drips.hashStreams(receivers), "Invalid streams hash after top-up");
         assertEq(streamsBalance, amt, "Invalid streams balance after top-up");
         assertEq(realBalanceDelta, int128(amt), "Invalid streams balance delta after top-up");
-        (bytes32 streamsHash,,,,) = drips.streamsState(accountId1, erc20);
-        assertEq(streamsHash, drips.hashStreams(receivers), "Invalid streams hash after top-up");
+
         // Withdraw
         balance = erc20.balanceOf(address(user));
         realBalanceDelta = driver.setStreams(