From 3710338f718d1202665117e02aaf0a3cb3fd334e Mon Sep 17 00:00:00 2001
From: qedk <1994constant@gmail.com>
Date: Sun, 24 Dec 2023 00:02:43 +0400
Subject: [PATCH] test: add WrappedAvail tests

---
 src/WrappedAvail.sol        | 10 ++++----
 test/WrappedAvailTest.t.sol | 47 +++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100644 test/WrappedAvailTest.t.sol

diff --git a/src/WrappedAvail.sol b/src/WrappedAvail.sol
index 36e90a3..a55eb0b 100644
--- a/src/WrappedAvail.sol
+++ b/src/WrappedAvail.sol
@@ -13,25 +13,25 @@ import {IWrappedAvail} from "src/interfaces/IWrappedAvail.sol";
 contract WrappedAvail is ERC20Permit, IWrappedAvail {
     address public immutable bridge;
 
-    error OnlyBridge();
+    error OnlyAvailBridge();
 
     constructor(address _bridge) ERC20Permit("Wrapped Avail") ERC20("WAVL", "Wrapped Avail") {
         // slither-disable-next-line missing-zero-check
         bridge = _bridge;
     }
 
-    modifier onlyBridge() {
+    modifier onlyAvailBridge() {
         if (msg.sender != bridge) {
-            revert OnlyBridge();
+            revert OnlyAvailBridge();
         }
         _;
     }
 
-    function mint(address destination, uint256 amount) external onlyBridge {
+    function mint(address destination, uint256 amount) external onlyAvailBridge {
         _mint(destination, amount);
     }
 
-    function burn(address from, uint256 amount) external onlyBridge {
+    function burn(address from, uint256 amount) external onlyAvailBridge {
         _burn(from, amount);
     }
 }
diff --git a/test/WrappedAvailTest.t.sol b/test/WrappedAvailTest.t.sol
new file mode 100644
index 0000000..95773a9
--- /dev/null
+++ b/test/WrappedAvailTest.t.sol
@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.23;
+
+import {WrappedAvail} from "src/WrappedAvail.sol";
+import {Vm, Test} from "forge-std/Test.sol";
+
+contract WrappedAvailTest is Test {
+    WrappedAvail public avail;
+    address public bridge;
+
+    function setUp() external {
+        bridge = makeAddr("bridge");
+        avail = new WrappedAvail(bridge);
+    }
+
+    function testRevertOnlyAvailBridge_mint(address sender, address dest, uint256 amount) external {
+        vm.assume(dest != address(0) && sender != bridge);
+        vm.prank(sender);
+        vm.expectRevert(WrappedAvail.OnlyAvailBridge.selector);
+        avail.mint(dest, amount);
+        assertEq(avail.balanceOf(dest), 0);
+    }
+
+    function test_mint(address dest, uint256 amount) external {
+        vm.assume(dest != address(0));
+        vm.prank(bridge);
+        avail.mint(dest, amount);
+        assertEq(avail.balanceOf(dest), amount);
+    }
+
+    function testRevertOnlyAvailBridge_burn(address sender, address dest, uint256 amount) external {
+        vm.assume(dest != address(0) && sender != bridge);
+        vm.prank(sender);
+        vm.expectRevert(WrappedAvail.OnlyAvailBridge.selector);
+        avail.burn(dest, amount);
+        assertEq(avail.balanceOf(dest), 0);
+    }
+
+    function test_burn(address dest, uint256 mintAmount, uint256 burnAmount) external {
+        vm.assume(dest != address(0) && burnAmount < mintAmount);
+        vm.startPrank(bridge);
+        avail.mint(dest, mintAmount);
+        assertEq(avail.balanceOf(dest), mintAmount);
+        avail.burn(dest, burnAmount);
+        assertEq(avail.balanceOf(dest), mintAmount - burnAmount);
+    }
+}