diff --git a/contracts/mocks/EnumerableSetMock.sol b/contracts/mocks/EnumerableSetMock.sol index 3ff690a0010..f9a2c9727da 100644 --- a/contracts/mocks/EnumerableSetMock.sol +++ b/contracts/mocks/EnumerableSetMock.sol @@ -4,7 +4,8 @@ pragma solidity ^0.6.0; import "../utils/EnumerableSet.sol"; -contract EnumerableSetMock { +// AddressSet +contract EnumerableAddressSetMock { using EnumerableSet for EnumerableSet.AddressSet; event OperationResult(bool result); @@ -33,3 +34,34 @@ contract EnumerableSetMock { return _set.at(index); } } + +// UintSet +contract EnumerableUintSetMock { + using EnumerableSet for EnumerableSet.UintSet; + + event OperationResult(bool result); + + EnumerableSet.UintSet private _set; + + function contains(uint256 value) public view returns (bool) { + return _set.contains(value); + } + + function add(uint256 value) public { + bool result = _set.add(value); + emit OperationResult(result); + } + + function remove(uint256 value) public { + bool result = _set.remove(value); + emit OperationResult(result); + } + + function length() public view returns (uint256) { + return _set.length(); + } + + function at(uint256 index) public view returns (uint256) { + return _set.at(index); + } +} \ No newline at end of file diff --git a/test/utils/EnumerableSet.behavior.js b/test/utils/EnumerableSet.behavior.js new file mode 100644 index 00000000000..7c4711c4e16 --- /dev/null +++ b/test/utils/EnumerableSet.behavior.js @@ -0,0 +1,116 @@ +const { expectEvent, expectRevert } = require('@openzeppelin/test-helpers'); +const { expect } = require('chai'); + +function shouldBehaveLikeSet (valueA, valueB, valueC) { + async function expectMembersMatch (set, values) { + await Promise.all(values.map(async value => + expect(await set.contains(value)).to.equal(true) + )); + + expect(await set.length()).to.bignumber.equal(values.length.toString()); + + // To compare values we convert to strings to workaround Chai + // limitations when dealing with nested arrays (required for BNs) + expect(await Promise.all([...Array(values.length).keys()].map(async (index) => { + const entry = await set.at(index); + return entry.toString(); + }))).to.have.same.members(values.map(v => v.toString())); + } + + it('starts empty', async function () { + expect(await this.set.contains(valueA)).to.equal(false); + + await expectMembersMatch(this.set, []); + }); + + it('adds a value', async function () { + const receipt = await this.set.add(valueA); + expectEvent(receipt, 'OperationResult', { result: true }); + + await expectMembersMatch(this.set, [valueA]); + }); + + it('adds several values', async function () { + await this.set.add(valueA); + await this.set.add(valueB); + + await expectMembersMatch(this.set, [valueA, valueB]); + expect(await this.set.contains(valueC)).to.equal(false); + }); + + it('returns false when adding values already in the set', async function () { + await this.set.add(valueA); + + const receipt = (await this.set.add(valueA)); + expectEvent(receipt, 'OperationResult', { result: false }); + + await expectMembersMatch(this.set, [valueA]); + }); + + it('reverts when retrieving non-existent elements', async function () { + await expectRevert(this.set.at(0), 'EnumerableSet: index out of bounds'); + }); + + it('removes added values', async function () { + await this.set.add(valueA); + + const receipt = await this.set.remove(valueA); + expectEvent(receipt, 'OperationResult', { result: true }); + + expect(await this.set.contains(valueA)).to.equal(false); + await expectMembersMatch(this.set, []); + }); + + it('returns false when removing values not in the set', async function () { + const receipt = await this.set.remove(valueA); + expectEvent(receipt, 'OperationResult', { result: false }); + + expect(await this.set.contains(valueA)).to.equal(false); + }); + + it('adds and removes multiple values', async function () { + // [] + + await this.set.add(valueA); + await this.set.add(valueC); + + // [A, C] + + await this.set.remove(valueA); + await this.set.remove(valueB); + + // [C] + + await this.set.add(valueB); + + // [C, B] + + await this.set.add(valueA); + await this.set.remove(valueC); + + // [A, B] + + await this.set.add(valueA); + await this.set.add(valueB); + + // [A, B] + + await this.set.add(valueC); + await this.set.remove(valueA); + + // [B, C] + + await this.set.add(valueA); + await this.set.remove(valueB); + + // [A, C] + + await expectMembersMatch(this.set, [valueA, valueC]); + + expect(await this.set.contains(valueB)).to.equal(false); + }); +} + +module.exports = { + shouldBehaveLikeSet, +}; diff --git a/test/utils/EnumerableSet.test.js b/test/utils/EnumerableSet.test.js index 8585b755dc2..ed292142b92 100644 --- a/test/utils/EnumerableSet.test.js +++ b/test/utils/EnumerableSet.test.js @@ -1,118 +1,33 @@ const { accounts, contract } = require('@openzeppelin/test-environment'); -const { expectEvent, expectRevert } = require('@openzeppelin/test-helpers'); -const { expect } = require('chai'); +const { BN } = require('@openzeppelin/test-helpers'); -const EnumerableSetMock = contract.fromArtifact('EnumerableSetMock'); +const EnumerableAddressSetMock = contract.fromArtifact('EnumerableAddressSetMock'); +const EnumerableUintSetMock = contract.fromArtifact('EnumerableUintSetMock'); -describe('EnumerableSet', function () { - const [ accountA, accountB, accountC ] = accounts; - - beforeEach(async function () { - this.set = await EnumerableSetMock.new(); - }); - - async function expectMembersMatch (set, values) { - await Promise.all(values.map(async account => - expect(await set.contains(account)).to.equal(true) - )); - - expect(await set.length()).to.bignumber.equal(values.length.toString()); - - expect(await Promise.all([...Array(values.length).keys()].map(index => - set.at(index) - ))).to.have.same.members(values); - } - - it('starts empty', async function () { - expect(await this.set.contains(accountA)).to.equal(false); - - await expectMembersMatch(this.set, []); - }); - - it('adds a value', async function () { - const receipt = await this.set.add(accountA); - expectEvent(receipt, 'OperationResult', { result: true }); +const { shouldBehaveLikeSet } = require('./EnumerableSet.behavior'); - await expectMembersMatch(this.set, [accountA]); - }); - - it('adds several values', async function () { - await this.set.add(accountA); - await this.set.add(accountB); - - await expectMembersMatch(this.set, [accountA, accountB]); - expect(await this.set.contains(accountC)).to.equal(false); - }); - - it('returns false when adding values already in the set', async function () { - await this.set.add(accountA); - - const receipt = (await this.set.add(accountA)); - expectEvent(receipt, 'OperationResult', { result: false }); - - await expectMembersMatch(this.set, [accountA]); - }); - - it('reverts when retrieving non-existent elements', async function () { - await expectRevert(this.set.at(0), 'EnumerableSet: index out of bounds'); - }); - - it('removes added values', async function () { - await this.set.add(accountA); - - const receipt = await this.set.remove(accountA); - expectEvent(receipt, 'OperationResult', { result: true }); - - expect(await this.set.contains(accountA)).to.equal(false); - await expectMembersMatch(this.set, []); - }); +describe('EnumerableSet', function () { + // AddressSet + describe('EnumerableAddressSet', function () { + const [ accountA, accountB, accountC ] = accounts; - it('returns false when removing values not in the set', async function () { - const receipt = await this.set.remove(accountA); - expectEvent(receipt, 'OperationResult', { result: false }); + beforeEach(async function () { + this.set = await EnumerableAddressSetMock.new(); + }); - expect(await this.set.contains(accountA)).to.equal(false); + shouldBehaveLikeSet(accountA, accountB, accountC); }); - it('adds and removes multiple values', async function () { - // [] - - await this.set.add(accountA); - await this.set.add(accountC); - - // [A, C] - - await this.set.remove(accountA); - await this.set.remove(accountB); - - // [C] - - await this.set.add(accountB); - - // [C, B] - - await this.set.add(accountA); - await this.set.remove(accountC); - - // [A, B] - - await this.set.add(accountA); - await this.set.add(accountB); - - // [A, B] - - await this.set.add(accountC); - await this.set.remove(accountA); - - // [B, C] - - await this.set.add(accountA); - await this.set.remove(accountB); - - // [A, C] + // UintSet + describe('EnumerableUintSet', function () { + const uintA = new BN('1234'); + const uintB = new BN('5678'); + const uintC = new BN('9101112'); - await expectMembersMatch(this.set, [accountA, accountC]); + beforeEach(async function () { + this.set = await EnumerableUintSetMock.new(); + }); - expect(await this.set.contains(accountB)).to.equal(false); + shouldBehaveLikeSet(uintA, uintB, uintC); }); });