Skip to content

Commit

Permalink
Update tests to new contract logic
Browse files Browse the repository at this point in the history
  • Loading branch information
@RCantu92
RCantu92 committed Nov 16, 2023
1 parent de93aa2 commit d86dd67
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 64 deletions.
28 changes: 17 additions & 11 deletions contracts/components/threat_oracle/ThreatOracleCore.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,22 +3,26 @@

pragma solidity ^0.8.9;

import "@openzeppelin/contracts/utils/structs/EnumerableMap.sol";
import "../BaseComponentUpgradeable.sol";

abstract contract ThreatOracleCore is BaseComponentUpgradeable {
struct ThreatProperties {
string category;
uint256 confidenceScore;
uint8 confidenceScore;
}

/// Confidence scores will be between 0-1.00,
// converted to unsigned integers
uint8 constant MAX_CONFIDENCE_SCORE = 100;

/// Map of addresses to their corresponding threat properties.
mapping (address => ThreatAddressProperties) private _addressThreatProperties;
mapping (address => ThreatProperties) private _addressThreatProperties;

event AddressRegistered(address indexed _address, string indexed category, uint256 indexed confidenceScore);
event AddressRegistered(address indexed _address, string indexed category, uint8 indexed confidenceScore);
event AddressDeregistered(address indexed _address);

error UnevenAmounts(uint256 addressesAmount, uint256 categoriesAmount, uint256 confidenceScoresAmount);
error ConfidenceScoreExceedsMax(uint8 maxLimit, uint8 exceedingScore);

/**
* @notice Method to register addresses with their threat properties.
Expand All @@ -33,7 +37,7 @@ abstract contract ThreatOracleCore is BaseComponentUpgradeable {
function registerAddresses(
address[] calldata addresses,
string[] calldata categories,
uint256[] calldata confidenceScores
uint8[] calldata confidenceScores
) external onlyRole(THREAT_ORACLE_ADMIN_ROLE) {
uint256 addressesAmount = addresses.length;
uint256 categoriesAmount = categories.length;
Expand All @@ -44,7 +48,7 @@ abstract contract ThreatOracleCore is BaseComponentUpgradeable {
addressesAmount != confidenceScoresAmount ||
categoriesAmount != confidenceScoresAmount
) {
revert UnevenAmounts(addressesAmount, threatLevelsAmount, confidenceScoresAmount);
revert UnevenAmounts(addressesAmount, categoriesAmount, confidenceScoresAmount);
}

for (uint256 i = 0; i < addressesAmount; i++) {
Expand All @@ -70,16 +74,18 @@ abstract contract ThreatOracleCore is BaseComponentUpgradeable {
* a category (e.g. 'exploit') and confidence score (0-1.0).
* @param _address Address of interest.
* @return category of the given address.
* @return confidence score of the given address.
* @return confidenceScore of the given address.
*/
function getThreatCategoryAndConfidence(address _address) public view returns (string, uint256) {
ThreatProperties threatProperties = _addressThreatProperties[_address];
function getThreatCategoryAndConfidence(address _address) public view returns (string memory category, uint8 confidenceScore) {
ThreatProperties memory threatProperties = _addressThreatProperties[_address];
return (threatProperties.category, threatProperties.confidenceScore);
}

function _registerAddress(address _address, string category, uint256 confidenceScore) private {
function _registerAddress(address _address, string calldata category, uint8 confidenceScore) private {
if(confidenceScore > MAX_CONFIDENCE_SCORE) revert ConfidenceScoreExceedsMax(MAX_CONFIDENCE_SCORE, confidenceScore);

_addressThreatProperties[_address] = ThreatProperties({ category: category, confidenceScore: confidenceScore });
emit AddressRegistered(_address, threatProperties.category, threatProperties.confidenceScore);
emit AddressRegistered(_address, category, confidenceScore);
}

function _deregisterAddress(address _address) private {
Expand Down
122 changes: 69 additions & 53 deletions test/components/threat.oracle.test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ const { ethers } = require('hardhat');
const { expect } = require('chai');
const { prepare } = require('../fixture');

const THREAT_CATEGORIES = ["exploit", "MEV"];

function createAddress(address) {
const paddedAddress = ethers.utils.hexZeroPad(address, 20);
return paddedAddress.toLowerCase();
Expand All @@ -15,6 +17,7 @@ function createAddresses(addressAmount) {
return generatedAddresses;
}

// TODO: Delete as this will be unnecessary.
function createThreatLevels(threatLevelAmount) {
const generatedThreatLevels = [];
for (let i = 0; i < threatLevelAmount; i++) {
Expand All @@ -24,91 +27,100 @@ function createThreatLevels(threatLevelAmount) {
return generatedThreatLevels;
}

function createThreatCategories(categoriesAmount) {
const generatedThreatCategories = [];
for (let i = 0; i < categoriesAmount; i++) {
const randomInt = Math.floor(Math.random() * THREAT_CATEGORIES.length);
const category = THREAT_CATEGORIES[randomInt];
generatedThreatCategories.push(category);
}
return generatedThreatCategories;
}

function createConfidenceScores(scoresAmount) {
const generatedConfidenceScores = [];
for (let i = 0; i < scoresAmount; i++) {
const randomFloat = Math.floor(Math.random() * 100);
const confidenceScore = randomFloat + 1;
generatedConfidenceScores.push(confidenceScore);
}
return generatedConfidenceScores;
}

const mockAddresses = createAddresses(10);
const mockThreatLevels = createThreatLevels(10);
const mockCategories = createThreatCategories(10);
const mockConfidenceScores = createConfidenceScores(10);

describe('Threat Oracle', async function () {
prepare(/*{ stake: { agents: { min: '100', max: '500', activated: true } } }*/);

it('registers a single address', async function () {
const initialAddressesRegistered = await this.threatOracle.totalAddressesRegistered();
let { category, confidenceScore } = await this.threatOracle.getThreatCategoryAndConfidence(mockAddresses[0]);
expect(category).to.be.equal("");
expect(confidenceScore).to.be.equal(0);

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).registerAddresses([mockAddresses[0]], [mockCategories[0]], [mockConfidenceScores[0]]))
.to.emit(this.threatOracle, 'AddressRegistered')
.withArgs(mockAddresses[0], mockCategories[0], mockConfidenceScores[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);
({ category, confidenceScore } = await this.threatOracle.getThreatCategoryAndConfidence(mockAddresses[0]));
expect(category).to.be.equal(mockCategories[0]);
expect(confidenceScore).to.be.equal(mockConfidenceScores[0]);
});

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);
for (let i = 0; i < mockAddresses.length; i++) {
let { category, confidenceScore } = await this.threatOracle.getThreatCategoryAndConfidence(mockAddresses[i]);
expect(category).to.be.equal("");
expect(confidenceScore).to.be.equal(0);
}
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];
await expect(this.threatOracle.connect(this.accounts.manager).registerAddresses(mockAddresses, mockCategories, mockConfidenceScores));

expect(await this.threatOracle.getThreatLevel(mockAddress)).to.be.equal(mockThreatLevel);
expect(await this.threatOracle.isRegistered(mockAddress)).to.be.equal(true);
for (let i = 0; i < mockAddresses.length; i++) {
({ category, confidenceScore } = await this.threatOracle.getThreatCategoryAndConfidence(mockAddresses[i]));
expect(category).to.be.equal(mockCategories[i]);
expect(confidenceScore).to.be.equal(mockConfidenceScores[i]);
}
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);
for (let i = 0; i < mockAddresses.length; i++) {
let { category, confidenceScore } = await this.threatOracle.getThreatCategoryAndConfidence(mockAddresses[i]);
expect(category).to.be.equal("");
expect(confidenceScore).to.be.equal(0);
}
expect(initialAddressesRegistered).to.be.equal(0);

const highThreatLevelsAmount = [...mockThreatLevels, 4, 5];
const highConfidenceScoresAmount = [...mockConfidenceScores, 40, 50];

await expect(this.threatOracle.connect(this.accounts.manager).setThreatLevels(mockAddresses, highThreatLevelsAmount))
.to.be.revertedWith(`UnevenAmounts(${mockAddresses.length}, ${highThreatLevelsAmount.length})`);
await expect(this.threatOracle.connect(this.accounts.manager).registerAddresses(mockAddresses, mockCategories, highConfidenceScoresAmount))
.to.be.revertedWith(`UnevenAmounts(${mockAddresses.length}, ${mockCategories.length}, ${highConfidenceScoresAmount.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);
let { category, confidenceScore } = await this.threatOracle.getThreatCategoryAndConfidence(mockAddresses[i]);
expect(category).to.be.equal("");
expect(confidenceScore).to.be.equal(0);
}
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);
for (let i = 0; i < mockAddresses.length; i++) {
let { category, confidenceScore } = await this.threatOracle.getThreatCategoryAndConfidence(mockAddresses[i]);
expect(category).to.be.equal("");
expect(confidenceScore).to.be.equal(0);
}

await expect(this.threatOracle.connect(this.accounts.other).setThreatLevels([mockAddresses[0]], [mockThreatLevels[0]]))
await expect(this.threatOracle.connect(this.accounts.other).registerAddresses(mockAddresses, mockCategories, mockConfidenceScores))
.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);
for (let i = 0; i < mockAddresses.length; i++) {
let { category, confidenceScore } = await this.threatOracle.getThreatCategoryAndConfidence(mockAddresses[i]);
expect(category).to.be.equal("");
expect(confidenceScore).to.be.equal(0);
}
});

it('allows FP to be corrected - lower existing threat level to zero', async function () {
it.only('allows FP to be corrected, i.e. remove address from block list', async function () {
const initialAddressesRegistered = await this.threatOracle.totalAddressesRegistered();

expect(await this.threatOracle.getThreatLevel(mockAddresses[0])).to.be.equal(0);
Expand All @@ -133,6 +145,9 @@ describe('Threat Oracle', async function () {
expect(await this.threatOracle.totalAddressesRegistered()).to.be.equal(initialAddressesRegistered + 1);
});

// Come back to this one once others have been updated
it('does not allow a confidence score that is too high to be registered', async function () {});

describe('Multicall', async function () {
it('allows to register a high number of addresses via multicall', async function () {
// Max amount before breaking was between `407` & `413`,
Expand Down Expand Up @@ -234,6 +249,7 @@ describe('Threat Oracle', async function () {
);
});

// Was testing this before course correcting
it.only('does not allow addresses to be registered if they and threat levels are uneven in amount - with multicall', async function () {
// Max amount before breaking was between `407` & `413`,
// and we want it to revert initially for this test,
Expand Down

0 comments on commit d86dd67

Please sign in to comment.