Added getTokensByOwner function to NFR. More tests

Black-Pepper-Team · Jun 30, 2024 · d7fa82f · d7fa82f
1 parent 4f4cf01
commit d7fa82f
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 2 deletions.
diff --git a/contracts/mocks/ERC721Mock.sol b/contracts/mocks/ERC721Mock.sol
@@ -20,6 +20,15 @@ contract ERC721Mock is ERC721Enumerable, ERC721URIStorage, Ownable {
         _burn(tokenId_);
     }
 
+    function getTokensByOwner(address owner_) public view returns (uint256[] memory tokens_) {
+        uint256 tokenCount_ = balanceOf(owner_);
+
+        tokens_ = new uint256[](tokenCount_);
+        for (uint256 i = 0; i < tokenCount_; ++i) {
+            tokens_[i] = tokenOfOwnerByIndex(owner_, i);
+        }
+    }
+
     function tokenURI(
         uint256 tokenId
     ) public view override(ERC721URIStorage, ERC721) returns (string memory) {

diff --git a/test/Chat.test.ts b/test/Chat.test.ts
@@ -12,7 +12,14 @@ import { Reverter } from "@/test/helpers/reverter";
 
 import { buildCredentialId, getMessageHash, getPoseidon, normalizeProof } from "@/test/helpers/zkp";
 
-import { AuthenticationStorage, Chat, ERC1967Proxy__factory, ERC721Mock, PoseidonSMT } from "@ethers-v6";
+import {
+  AuthenticationStorage,
+  Chat,
+  ERC1967Proxy__factory,
+  ERC721Mock,
+  PoseidonSMT,
+  PostMessageVerifier,
+} from "@ethers-v6";
 
 describe("Chat", () => {
   const reverter = new Reverter();
@@ -27,6 +34,8 @@ describe("Chat", () => {
   let chat: Chat;
   let authStorage: AuthenticationStorage;
 
+  let postMessageVerifier: PostMessageVerifier;
+
   before(async () => {
     [SECOND] = await ethers.getSigners();
 
@@ -67,7 +76,7 @@ describe("Chat", () => {
     chat = chat.attach(await proxy.getAddress()) as Chat;
 
     const PostMessageVerifier = await ethers.getContractFactory("PostMessageVerifier");
-    const postMessageVerifier = await PostMessageVerifier.deploy();
+    postMessageVerifier = await PostMessageVerifier.deploy();
 
     await chat.__Chat_init(await tree.getAddress(), await postMessageVerifier.getAddress());
 
@@ -154,6 +163,30 @@ describe("Chat", () => {
 
       const formattedProof = normalizeProof(data);
 
+      expect(
+        await circuit.verifyProof({
+          proof: {
+            pi_a: [String(formattedProof.a[0]), String(formattedProof.a[1])],
+            pi_b: [
+              [String(formattedProof.b[0][1]), String(formattedProof.b[0][0])],
+              [String(formattedProof.b[1][1]), String(formattedProof.b[1][0])],
+            ],
+            pi_c: [String(formattedProof.c[0]), String(formattedProof.c[1])],
+            protocol: "groth16",
+            curve: "bn128",
+          },
+          publicSignals: [data.publicSignals[0], data.publicSignals[1], data.publicSignals[2], data.publicSignals[3]],
+        }),
+      ).to.be.true;
+      expect(
+        await postMessageVerifier.verifyProof(formattedProof.a, formattedProof.b, formattedProof.c, [
+          data.publicSignals[0],
+          data.publicSignals[1],
+          data.publicSignals[2],
+          data.publicSignals[3],
+        ]),
+      ).to.be.true;
+
       await expect(chat.postMessage(erc721.getAddress(), message, proof.root, deadline, formattedProof))
         .to.emit(chat, "MessagePosted")
         .withArgs(await erc721.getAddress(), message);