First set of tests in suite

forta-network · Nov 10, 2023 · 07cd629 · 07cd629
1 parent 88258bd
commit 07cd629
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 12 deletions.
diff --git a/contracts/components/threat_oracle/ThreatOracleCore.sol b/contracts/components/threat_oracle/ThreatOracleCore.sol
@@ -31,8 +31,9 @@ abstract contract ThreatOracleCore is BaseComponentUpgradeable {
         emit AddressThreatLevelSet(_address, threatLevel);
     }
 
-    function getThreatLevel(address _address) public view returns (uint) {
-        return _addressThreatLevel.get(_address);
+    function getThreatLevel(address _address) public view returns (uint256) {
+        (,uint256 fetchedThreatLevel) = _addressThreatLevel.tryGet(_address);
+        return fetchedThreatLevel;
     }
 
     function isRegistered(address _address) public view returns (bool) {

diff --git a/test/components/threat.oracle.test.js b/test/components/threat.oracle.test.js
@@ -15,20 +15,121 @@ function createAddresses(addressAmount) {
     return generatedAddresses;
 }
 
-const mockAddresses = [...createAddresses(10)];
+function createThreatLevels(threatLevelAmount) {
+    const generatedThreatLevels = [];
+    for (let i = 0; i < threatLevelAmount; i++) {
+        const threatLevel = Math.floor(Math.random() * 5) + 1;
+        generatedThreatLevels.push(threatLevel);
+    }
+    return generatedThreatLevels;
+}
+
+const mockAddresses = createAddresses(10);
+const mockThreatLevels = createThreatLevels(10);
 
 describe('Threat Oracles', async function () {
     prepare(/*{ stake: { agents: { min: '100', max: '500', activated: true } } }*/);
 
-    it('registers multiple addresses to the threat oracle', async function () {
-        // TODO: Update `getThreatOracle` to use `tryGet` instead
-        console.log(await this.threatOracle.getThreatLevel(mockAddresses[0]));
-        // expect(await this.threatOracle.getThreatLevel()).to.be.equal();
+    it('registers a single address', async function () {
+        const initialAddressesRegistered = await this.threatOracle.totalAddressesRegistered();
+
+        expect(await this.threatOracle.getThreatLevel(mockAddresses[0])).to.be.equal(0);
+        expect(await this.threatOracle.isRegistered(mockAddresses[0])).to.be.equal(false);
+        expect(initialAddressesRegistered).to.be.equal(0);
+
+        await expect(this.threatOracle.connect(this.accounts.manager).setThreatLevels([mockAddresses[0]], [mockThreatLevels[0]]))
+            .to.emit(this.threatOracle, 'AddressThreatLevelSet')
+            .withArgs(mockAddresses[0], mockThreatLevels[0]);
+
+        expect(await this.threatOracle.getThreatLevel(mockAddresses[0])).to.be.equal(mockThreatLevels[0]);
+        expect(await this.threatOracle.isRegistered(mockAddresses[0])).to.be.equal(true);
+        expect(await this.threatOracle.totalAddressesRegistered()).to.be.equal(initialAddressesRegistered + 1);
+    });
+
+    it('registers multiple addresses', async function () {
+        const initialAddressesRegistered = await this.threatOracle.totalAddressesRegistered();
+
+        for(let i = 0; i < mockAddresses.length; i++) {
+            const mockAddress = mockAddresses[i];
+            expect(await this.threatOracle.getThreatLevel(mockAddress)).to.be.equal(0);
+            expect(await this.threatOracle.isRegistered(mockAddress)).to.be.equal(false);
+        }
+        expect(initialAddressesRegistered).to.be.equal(0);
+
+        await this.threatOracle.connect(this.accounts.manager).setThreatLevels(mockAddresses, mockThreatLevels);
+
+        for(let i = 0; i < mockAddresses.length; i++) {
+            const mockAddress = mockAddresses[i];
+            const mockThreatLevel = mockThreatLevels[i];
+
+            expect(await this.threatOracle.getThreatLevel(mockAddress)).to.be.equal(mockThreatLevel);
+            expect(await this.threatOracle.isRegistered(mockAddress)).to.be.equal(true);
+        }
+        expect(await this.threatOracle.totalAddressesRegistered()).to.be.equal(initialAddressesRegistered + mockAddresses.length);
+    });
+
+    it('does not allow to register addresses and threat levels if they are not equal in amount', async function () {
+        const initialAddressesRegistered = await this.threatOracle.totalAddressesRegistered();
+
+        for(let i = 0; i < mockAddresses.length; i++) {
+            const mockAddress = mockAddresses[i];
+
+            expect(await this.threatOracle.getThreatLevel(mockAddress)).to.be.equal(0);
+            expect(await this.threatOracle.isRegistered(mockAddress)).to.be.equal(false);
+        }
+        expect(initialAddressesRegistered).to.be.equal(0);
+
+        const highThreatLevelsAmount = [...mockThreatLevels, 4, 5];
+
+        await expect(this.threatOracle.connect(this.accounts.manager).setThreatLevels(mockAddresses, highThreatLevelsAmount))
+            .to.be.revertedWith(`IncorrectAmounts(${mockAddresses.length}, ${highThreatLevelsAmount.length})`);
+
+        for(let i = 0; i < mockAddresses.length; i++) {
+            const mockAddress = mockAddresses[i];
+
+            expect(await this.threatOracle.getThreatLevel(mockAddress)).to.be.equal(0);
+            expect(await this.threatOracle.isRegistered(mockAddress)).to.be.equal(false);
+        }
+        expect(initialAddressesRegistered).to.be.equal(0);
+    });
+
+    it('does not allow an account without access control to register addresses', async function () {
+        const initialAddressesRegistered = await this.threatOracle.totalAddressesRegistered();
+
+        expect(await this.threatOracle.getThreatLevel(mockAddresses[0])).to.be.equal(0);
+        expect(await this.threatOracle.isRegistered(mockAddresses[0])).to.be.equal(false);
+        expect(initialAddressesRegistered).to.be.equal(0);
+
+        await expect(this.threatOracle.connect(this.accounts.other).setThreatLevels([mockAddresses[0]], [mockThreatLevels[0]]))
+            .to.be.revertedWith(`MissingRole("${this.roles.THREAT_ORACLE_ADMIN}", "${this.accounts.other.address}")`);
+
+        expect(await this.threatOracle.getThreatLevel(mockAddresses[0])).to.be.equal(0);
+        expect(await this.threatOracle.isRegistered(mockAddresses[0])).to.be.equal(false);
+        expect(initialAddressesRegistered).to.be.equal(0);
+    });
+
+    it('allows FP to be corrected - lower existing threat level to zero', async function () {
+        const initialAddressesRegistered = await this.threatOracle.totalAddressesRegistered();
+
+        expect(await this.threatOracle.getThreatLevel(mockAddresses[0])).to.be.equal(0);
+        expect(await this.threatOracle.isRegistered(mockAddresses[0])).to.be.equal(false);
+        expect(initialAddressesRegistered).to.be.equal(0);
+
+        await expect(this.threatOracle.connect(this.accounts.manager).setThreatLevels([mockAddresses[0]], [mockThreatLevels[0]]))
+            .to.emit(this.threatOracle, 'AddressThreatLevelSet')
+            .withArgs(mockAddresses[0], mockThreatLevels[0]);
+
+        expect(await this.threatOracle.getThreatLevel(mockAddresses[0])).to.be.equal(mockThreatLevels[0]);
+        expect(await this.threatOracle.isRegistered(mockAddresses[0])).to.be.equal(true);
+        expect(await this.threatOracle.totalAddressesRegistered()).to.be.equal(initialAddressesRegistered + 1);
+
+        // FP correction
+        await expect(this.threatOracle.connect(this.accounts.manager).setThreatLevels([mockAddresses[0]], [0]))
+            .to.emit(this.threatOracle, 'AddressThreatLevelSet')
+            .withArgs(mockAddresses[0], 0);
 
-        /*
-        await expect(this.agents.connect(this.accounts.manager).setFrontRunningDelay('1800'))
-            .to.emit(this.agents, 'FrontRunningDelaySet')
-            .withArgs(ethers.BigNumber.from('1800'));
-        */
+        expect(await this.threatOracle.getThreatLevel(mockAddresses[0])).to.be.equal(0);
+        expect(await this.threatOracle.isRegistered(mockAddresses[0])).to.be.equal(true);
+        expect(await this.threatOracle.totalAddressesRegistered()).to.be.equal(initialAddressesRegistered + 1);
     });
 })