Update docs

test/utils/structs/EnumerableSet.test.js

@@ -1,46 +1,79 @@
-const { BN } = require('@openzeppelin/test-helpers');
-
-const EnumerableBytes32SetMock = artifacts.require('EnumerableBytes32SetMock');
-const EnumerableAddressSetMock = artifacts.require('EnumerableAddressSetMock');
-const EnumerableUintSetMock = artifacts.require('EnumerableUintSetMock');
+const EnumerableSet = artifacts.require('$EnumerableSet');
+const { mapValues } = require('../../helpers/map-values');
 
 const { shouldBehaveLikeSet } = require('./EnumerableSet.behavior');
 
+const getMethods = ms => {
+  return mapValues(
+    ms,
+    m =>
+      (self, ...args) =>
+        self.methods[m](0, ...args),
+  );
+};
+
+// Get the name of the library. In the transpiled code it will be EnumerableSetUpgradeable.
+const library = EnumerableSet._json.contractName.replace(/^\$/, '');
+
 contract('EnumerableSet', function (accounts) {
+  beforeEach(async function () {
+    this.set = await EnumerableSet.new();
+  });
+
   // Bytes32Set
   describe('EnumerableBytes32Set', function () {
-    const bytesA = '0xdeadbeef'.padEnd(66, '0');
-    const bytesB = '0x0123456789'.padEnd(66, '0');
-    const bytesC = '0x42424242'.padEnd(66, '0');
-
-    beforeEach(async function () {
-      this.set = await EnumerableBytes32SetMock.new();
-    });
-
-    shouldBehaveLikeSet(bytesA, bytesB, bytesC);
+    shouldBehaveLikeSet(
+      ['0xdeadbeef', '0x0123456789', '0x42424242'].map(e => e.padEnd(66, '0')),
+      getMethods({
+        add: '$add(uint256,bytes32)',
+        remove: '$remove(uint256,bytes32)',
+        contains: '$contains(uint256,bytes32)',
+        length: `$length_${library}_Bytes32Set(uint256)`,
+        at: `$at_${library}_Bytes32Set(uint256,uint256)`,
+        values: `$values_${library}_Bytes32Set(uint256)`,
+      }),
+      {
+        addReturn: `return$add_${library}_Bytes32Set_bytes32`,
+        removeReturn: `return$remove_${library}_Bytes32Set_bytes32`,
+      },
+    );
   });
 
   // AddressSet
   describe('EnumerableAddressSet', function () {
-    const [accountA, accountB, accountC] = accounts;
-
-    beforeEach(async function () {
-      this.set = await EnumerableAddressSetMock.new();
-    });
-
-    shouldBehaveLikeSet(accountA, accountB, accountC);
+    shouldBehaveLikeSet(
+      accounts,
+      getMethods({
+        add: '$add(uint256,address)',
+        remove: '$remove(uint256,address)',
+        contains: '$contains(uint256,address)',
+        length: `$length_${library}_AddressSet(uint256)`,
+        at: `$at_${library}_AddressSet(uint256,uint256)`,
+        values: `$values_${library}_AddressSet(uint256)`,
+      }),
+      {
+        addReturn: `return$add_${library}_AddressSet_address`,
+        removeReturn: `return$remove_${library}_AddressSet_address`,
+      },
+    );
   });
 
   // UintSet
   describe('EnumerableUintSet', function () {
-    const uintA = new BN('1234');
-    const uintB = new BN('5678');
-    const uintC = new BN('9101112');
-
-    beforeEach(async function () {
-      this.set = await EnumerableUintSetMock.new();
-    });
-
-    shouldBehaveLikeSet(uintA, uintB, uintC);
+    shouldBehaveLikeSet(
+      [1234, 5678, 9101112].map(e => web3.utils.toBN(e)),
+      getMethods({
+        add: '$add(uint256,uint256)',
+        remove: '$remove(uint256,uint256)',
+        contains: '$contains(uint256,uint256)',
+        length: `$length_${library}_UintSet(uint256)`,
+        at: `$at_${library}_UintSet(uint256,uint256)`,
+        values: `$values_${library}_UintSet(uint256)`,
+      }),
+      {
+        addReturn: `return$add_${library}_UintSet_uint256`,
+        removeReturn: `return$remove_${library}_UintSet_uint256`,
+      },
+    );
   });
 });