Add MockERC20 token minting and transactions

contracts/SampleShop1.sol

@@ -3,20 +3,38 @@ pragma solidity >=0.6.0 <0.9.0;
 
 import "@openzeppelin/contracts/token/ERC721/extensions/ERC721Enumerable.sol";
 import "@openzeppelin/contracts/access/Ownable.sol";
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import "./interface/ShopInterface.sol";
 import "hardhat/console.sol";
 
-contract SampleShop1 is Shop {
-  mapping(address => uint256) public reputations;
+contract SampleShop1 is Shop, ERC20 {
+  mapping(address => uint256) public balances;
+  uint256 public supply;
 
-  constructor() {}
+  constructor(
+    string memory _name,
+    string memory _symbol,
+    uint256 _supply
+  ) ERC20(_name, _symbol) {
+    supply = _supply;
+  }
 
   event ExampleInteractionEvent(uint256 _value, address _sender);
 
+  // Main interact function the VendingMachine contract is calling
   function interact(address _sender) public payable override {
-    require(msg.value == 0, "Invalid amount");
-    reputations[_sender] += 1;
-    // sample interaction
+    require(msg.value > 0, "Invalid amount");
+    require(totalSupply() + msg.value * 2 <= supply, "Over supply limit");
+    approve(_sender, msg.value * 2);
+    balances[_sender] += msg.value * 2; // sample interaction. TODO: Add safe math
     emit ExampleInteractionEvent(msg.value, msg.sender);
   }
+
+  // Withdraw token
+  function withdraw() public {
+    uint256 amount = balances[msg.sender];
+    require(amount != 0, "Balance is empty");
+    _mint(msg.sender, amount);
+    balances[msg.sender] = 0;
+  }
 }

contracts/VendingMachine.sol

@@ -25,7 +25,6 @@ contract VendingMachine is Ownable {
 
   // Transact
   // user can transact with shop if within range defined by zkSNARK
-  // TODO: Add mock token
   function transact(
     uint256[2] memory _a,
     uint256[2][2] memory _b,
@@ -47,6 +46,7 @@ contract VendingMachine is Ownable {
   ) public payable onlyInPosition(_a, _b, _c, _input) {
     uint256 _locationHash = _input[0];
     require(msg.value > overtakeFees[_locationHash], "Need more stake amount");
+    balances[shopOwners[_locationHash]] += overtakeFees[_locationHash]; // return previous take amount to previous owner, similar to bidding
     overtakeFees[_locationHash] = msg.value;
     shopOwners[_locationHash] = msg.sender;
   }

test/helper.ts

@@ -1,8 +1,10 @@
-export enum revertMessages {
+export enum REVERT_MESSAGES {
   INVALID_PROOF = "Proof is not valid",
   INVALID_OWNER = "Not shop owner",
   INVALID_RANGE = "Not in range",
   INVALID_STAKE = "Need more stake amount",
+  EMPTY_BALANCE = "Balance is empty",
+  OVERFLOW_TOKEN_SUPPLY = "Over supply limit",
 }
 
 export const sampleRangeProof = [

test/main.test.ts

@@ -1,7 +1,7 @@
 import { expect } from "chai";
 import { ethers as Ethers } from "ethers";
 import { ethers } from "hardhat";
-import { revertMessages } from "./helper";
+import { REVERT_MESSAGES } from "./helper";
 import { sampleRangeProof } from "./helper";
 
 const machineLocationHash = sampleRangeProof[3][0];
@@ -17,7 +17,7 @@ describe("ZK Vending Machine", () => {
 
   before(async () => {
     zkvmContract = await (await ethers.getContractFactory("VendingMachine")).deploy([machineLocationHash], [INITIAL_OVERTAKE_FEE]);
-    sampleShop1Contract = await (await ethers.getContractFactory("SampleShop1")).deploy();
+    sampleShop1Contract = await (await ethers.getContractFactory("SampleShop1")).deploy("DIAMOND", "DIA", 100);
     [addr1, addr2, addr3] = await ethers.getSigners();
   });
 
@@ -33,18 +33,36 @@ describe("ZK Vending Machine", () => {
   });
 
   it("Transact with new shop", async () => {
-    await zkvmContract.transact(...sampleRangeProof);
-    expect(await sampleShop1Contract.reputations(addr1.address)).to.be.equal(1);
+    await zkvmContract.transact(...sampleRangeProof, { value: 10 });
+    expect(await sampleShop1Contract.balances(addr1.address)).to.be.equal(20);
   });
 
   it("Overtake shop", async () => {
-    await zkvmContract.connect(addr2).overtake(...sampleRangeProof, { value: 20 });
+    expect(zkvmContract.connect(addr2).overtake(...sampleRangeProof, { value: 10 })).to.be.revertedWith(REVERT_MESSAGES.INVALID_STAKE); // fail to overtake shop: insufficient stake amount
+
+    await zkvmContract.connect(addr2).overtake(...sampleRangeProof, { value: 20 }); // successful overtake with correct stake amount
+    expect(await zkvmContract.balances(addr1.address)).to.be.equal(10);
     expect(await zkvmContract.shopOwners(machineLocationHash)).to.be.equal(addr2.address);
     expect(await zkvmContract.overtakeFees(machineLocationHash)).to.be.equal(20);
   });
 
   it("Transact with overtaken shop", async () => {
-    await zkvmContract.connect(addr2).transact(...sampleRangeProof);
-    expect(await sampleShop1Contract.reputations(addr2.address)).to.be.equal(1);
+    await zkvmContract.connect(addr2).transact(...sampleRangeProof, { value: 20 });
+    expect(await sampleShop1Contract.balances(addr2.address)).to.be.equal(40);
+
+    expect(sampleShop1Contract.connect(addr2).interact(addr2.address, { value: 200 })).to.be.revertedWith(REVERT_MESSAGES.OVERFLOW_TOKEN_SUPPLY); // fail to transact: overflow balance
+  });
+
+  it("Withdraw staked amount", async () => {
+    expect(zkvmContract.connect(addr3).withdraw()).to.be.revertedWith(REVERT_MESSAGES.EMPTY_BALANCE); // empty balance
+    await zkvmContract.connect(addr1).withdraw();
+  });
+
+  it("Withdraw tokens from shop1", async () => {
+    expect(sampleShop1Contract.connect(addr3).withdraw()).to.be.revertedWith(REVERT_MESSAGES.EMPTY_BALANCE); // cannot withdraw empty balance
+
+    await sampleShop1Contract.connect(addr2).withdraw();
+    expect(await sampleShop1Contract.balanceOf(addr2.address)).to.be.equal(40);
+    expect(await sampleShop1Contract.balances(addr2.address)).to.be.equal(0);
   });
 });