Lock.sol Contract Refactor

contracts/Lock.sol

@@ -1,34 +1,95 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.28;
 
-// Uncomment this line to use console.log
-// import "hardhat/console.sol";
-
+/**
+ * @title Lock Contract
+ * @dev A simple time-locked wallet that allows withdrawal after a specified time
+ * @notice Funds are locked until the specified unlock time
+ */
 contract Lock {
-    uint public unlockTime;
-    address payable public owner;
+    // === State Variables ===
+    uint256 public immutable i_unlockTime;
+    address payable public immutable i_owner;
 
-    event Withdrawal(uint amount, uint when);
+    // === Events ===
+    event Withdrawal(uint256 amount, uint256 when);
+    event Deposit(address indexed from, uint256 amount);
+
+    // === Custom Errors ===
+    error Lock__UnlockTimeNotInFuture();
+    error Lock__NoFundsProvided();
+    error Lock__WithdrawalTooEarly(uint256 currentTime, uint256 unlockTime);
+    error Lock__NotOwner(address caller, address owner);
+    error Lock__TransferFailed();
 
-    constructor(uint _unlockTime) payable {
-        require(
-            block.timestamp < _unlockTime,
-            "Unlock time should be in the future"
-        );
+    /**
+     * @dev Contract constructor
+     * @param _unlockTime Unix timestamp when funds will be unlockable
+     */
+    constructor(uint256 _unlockTime) payable {
+        if (block.timestamp >= _unlockTime) {
+            revert Lock__UnlockTimeNotInFuture();
+        }
+        if (msg.value <= 0) {
+            revert Lock__NoFundsProvided();
+        }
 
-        unlockTime = _unlockTime;
-        owner = payable(msg.sender);
+        i_unlockTime = _unlockTime;
+        i_owner = payable(msg.sender);
+
+        if (msg.value > 0) {
+            emit Deposit(msg.sender, msg.value);
+        }
     }
 
-    function withdraw() public {
-        // Uncomment this line, and the import of "hardhat/console.sol", to print a log in your terminal
-        // console.log("Unlock time is %o and block timestamp is %o", unlockTime, block.timestamp);
+    /**
+     * @dev Allow deposits to the time lock
+     */
+    receive() external payable {
+        emit Deposit(msg.sender, msg.value);
+    }
 
-        require(block.timestamp >= unlockTime, "You can't withdraw yet");
-        require(msg.sender == owner, "You aren't the owner");
+    /**
+     * @dev Withdraws all funds from the contract
+     * @notice Can only be called by the owner after unlock time
+     */
+    function withdraw() external {
+        // Check conditions
+        if (block.timestamp < i_unlockTime) {
+            revert Lock__WithdrawalTooEarly(block.timestamp, i_unlockTime);
+        }
+        if (msg.sender != i_owner) {
+            revert Lock__NotOwner(msg.sender, i_owner);
+        }
+
+        uint256 amount = address(this).balance;
+
+        // Emit event before external call (follows checks-effects-interactions)
+        emit Withdrawal(amount, block.timestamp);
 
-        emit Withdrawal(address(this).balance, block.timestamp);
+        // Transfer funds
+        (bool success, ) = i_owner.call{value: amount}("");
+        if (!success) {
+            revert Lock__TransferFailed();
+        }
+    }
+
+    /**
+     * @dev Returns the current balance of the contract
+     * @return uint256 The contract balance in wei
+     */
+    function getBalance() external view returns (uint256) {
+        return address(this).balance;
+    }
 
-        owner.transfer(address(this).balance);
+    /**
+     * @dev Returns the time remaining until unlock
+     * @return uint256 Seconds until unlock (0 if already unlocked)
+     */
+    function getTimeRemaining() external view returns (uint256) {
+        if (block.timestamp >= i_unlockTime) {
+            return 0;
+        }
+        return i_unlockTime - block.timestamp;
     }
 }

test/Lock.ts

@@ -30,13 +30,13 @@ describe("Lock", function () {
     it("Should set the right unlockTime", async function () {
       const { lock, unlockTime } = await loadFixture(deployOneYearLockFixture);
 
-      expect(await lock.unlockTime()).to.equal(unlockTime);
+      expect(await lock.i_unlockTime()).to.equal(unlockTime);
     });
 
     it("Should set the right owner", async function () {
       const { lock, owner } = await loadFixture(deployOneYearLockFixture);
 
-      expect(await lock.owner()).to.equal(owner.address);
+      expect(await lock.i_owner()).to.equal(owner.address);
     });
 
     it("Should receive and store the funds to lock", async function () {
@@ -53,8 +53,19 @@ describe("Lock", function () {
       // We don't use the fixture here because we want a different deployment
       const latestTime = await time.latest();
       const Lock = await hre.ethers.getContractFactory("Lock");
-      await expect(Lock.deploy(latestTime, { value: 1 })).to.be.revertedWith(
-        "Unlock time should be in the future"
+      await expect(Lock.deploy(latestTime, { value: 1 })).to.be.revertedWithCustomError(
+        Lock,
+        "Lock__UnlockTimeNotInFuture"
+      );
+    });
+
+    it("Should fail if no funds are provided", async function () {
+      const ONE_YEAR_IN_SECS = 365 * 24 * 60 * 60;
+      const unlockTime = (await time.latest()) + ONE_YEAR_IN_SECS;
+      const Lock = await hre.ethers.getContractFactory("Lock");
+      await expect(Lock.deploy(unlockTime, { value: 0 })).to.be.revertedWithCustomError(
+        Lock,
+        "Lock__NoFundsProvided"
       );
     });
   });
@@ -64,8 +75,9 @@ describe("Lock", function () {
       it("Should revert with the right error if called too soon", async function () {
         const { lock } = await loadFixture(deployOneYearLockFixture);
 
-        await expect(lock.withdraw()).to.be.revertedWith(
-          "You can't withdraw yet"
+        await expect(lock.withdraw()).to.be.revertedWithCustomError(
+          lock,
+          "Lock__WithdrawalTooEarly"
         );
       });
 
@@ -78,8 +90,9 @@ describe("Lock", function () {
         await time.increaseTo(unlockTime);
 
         // We use lock.connect() to send a transaction from another account
-        await expect(lock.connect(otherAccount).withdraw()).to.be.revertedWith(
-          "You aren't the owner"
+        await expect(lock.connect(otherAccount).withdraw()).to.be.revertedWithCustomError(
+          lock,
+          "Lock__NotOwner"
         );
       });
 
@@ -124,4 +137,68 @@ describe("Lock", function () {
       });
     });
   });
+
+  describe("Additional Functions", function () {
+    it("Should return the correct balance", async function () {
+      const { lock, lockedAmount } = await loadFixture(deployOneYearLockFixture);
+
+      expect(await lock.getBalance()).to.equal(lockedAmount);
+    });
+
+    it("Should return the correct time remaining", async function () {
+      const { lock, unlockTime } = await loadFixture(deployOneYearLockFixture);
+
+      const currentTime = await time.latest();
+      const expectedTimeRemaining = unlockTime - currentTime;
+
+      expect(await lock.getTimeRemaining()).to.be.closeTo(expectedTimeRemaining, 2);
+    });
+
+    it("Should return 0 time remaining after unlock time", async function () {
+      const { lock, unlockTime } = await loadFixture(deployOneYearLockFixture);
+
+      await time.increaseTo(unlockTime + 1);
+
+      expect(await lock.getTimeRemaining()).to.equal(0);
+    });
+  });
+
+  describe("Deposits", function () {
+    it("Should emit Deposit event on construction", async function () {
+      const ONE_YEAR_IN_SECS = 365 * 24 * 60 * 60;
+      const ONE_GWEI = 1_000_000_000;
+      const unlockTime = (await time.latest()) + ONE_YEAR_IN_SECS;
+      const [owner] = await hre.ethers.getSigners();
+
+      const Lock = await hre.ethers.getContractFactory("Lock");
+      const lock = await Lock.deploy(unlockTime, { value: ONE_GWEI });
+      const receipt = await lock.deploymentTransaction()?.wait();
+
+      // Check that a Deposit event was emitted during deployment
+      const depositEvents = receipt?.logs.filter(log => {
+        try {
+          const parsed = lock.interface.parseLog(log);
+          return parsed?.name === 'Deposit';
+        } catch {
+          return false;
+        }
+      });
+
+      expect(depositEvents).to.have.length(1);
+    });
+
+    it("Should accept additional deposits via receive function", async function () {
+      const { lock, owner, lockedAmount } = await loadFixture(deployOneYearLockFixture);
+      const additionalAmount = 500_000_000;
+
+      await expect(owner.sendTransaction({
+        to: lock.target,
+        value: additionalAmount
+      }))
+        .to.emit(lock, "Deposit")
+        .withArgs(owner.address, additionalAmount);
+
+      expect(await lock.getBalance()).to.equal(lockedAmount + additionalAmount);
+    });
+  });
 });