From b2d4a667d5f9b940f137af7854ff5a62424b90eb Mon Sep 17 00:00:00 2001
From: cgewecke <christophergewecke@gmail.com>
Date: Tue, 5 Jan 2021 19:10:38 -0800
Subject: [PATCH] Add modifierWhitelist option (#597)

---
 README.md               |  1 +
 lib/instrumenter.js     |  3 ++-
 lib/parse.js            |  3 ++-
 lib/registrar.js        | 10 +++++++++-
 lib/validator.js        |  5 +++++
 test/units/modifiers.js | 24 +++++++++++++++++++++++-
 test/units/validator.js |  1 +
 7 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index ed90615..16d7303 100644
--- a/README.md
+++ b/README.md
@@ -102,6 +102,7 @@ module.exports = {
 | measureFunctionCoverage | *boolean* | `true` | Computes function coverage. [More...][34] |
 | measureModifierCoverage | *boolean* | `true` | Computes each modifier invocation as a code branch. [More...][34] |
 | matrixOutputPath | *String* | `./testMatrix.json` | Relative path to write test matrix JSON object to. [More...][39] |
+| modifierWhitelist | *String[]* | `[]` | List of modifier names (ex: "onlyOwner") to exclude from branch measurement. (Useful for modifiers which prepare something instead of acting as a gate.)) |
 | istanbulFolder | *String* | `./coverage` |  Folder location for Istanbul coverage reports. |
 | istanbulReporter | *Array* | `['html', 'lcov', 'text', 'json']` | [Istanbul coverage reporters][2]  |
 | mocha | *Object* | `{ }` | [Mocha options][3] to merge into existing mocha config. `grep` and `invert` are useful for skipping certain tests under coverage using tags in the test descriptions.|
diff --git a/lib/instrumenter.js b/lib/instrumenter.js
index 90f316f..988f90e 100644
--- a/lib/instrumenter.js
+++ b/lib/instrumenter.js
@@ -15,6 +15,7 @@ class Instrumenter {
   constructor(config={}){
     this.instrumentationData = {};
     this.injector = new Injector();
+    this.modifierWhitelist = config.modifierWhitelist || [];
     this.enabled = {
       statements: (config.measureStatementCoverage === false) ? false : true,
       functions: (config.measureFunctionCoverage === false) ? false: true,
@@ -62,7 +63,7 @@ class Instrumenter {
     const contract = {};
 
     this.injector.resetModifierMapping();
-    parse.configure(this.enabled);
+    parse.configure(this.enabled, this.modifierWhitelist);
 
     contract.source = contractSource;
     contract.instrumented = contractSource;
diff --git a/lib/parse.js b/lib/parse.js
index b7e7c02..e4128ba 100644
--- a/lib/parse.js
+++ b/lib/parse.js
@@ -11,8 +11,9 @@ const FILE_SCOPED_ID = "fileScopedId";
 const parse = {};
 
 // Utilities
-parse.configure = function(_enabled){
+parse.configure = function(_enabled, _whitelist){
   register.enabled = Object.assign(register.enabled, _enabled);
+  register.modifierWhitelist = _whitelist;
 }
 
 // Nodes
diff --git a/lib/registrar.js b/lib/registrar.js
index 5c6da7c..f44af4f 100644
--- a/lib/registrar.js
+++ b/lib/registrar.js
@@ -19,6 +19,8 @@ class Registrar {
       branches: true,
       lines: true
     }
+
+    this.modifierWhitelist = [];
   }
 
   /**
@@ -129,7 +131,13 @@ class Registrar {
         }
 
         // Add modifier branch coverage
-        if (!this.enabled.modifiers || expression.isConstructor) continue;
+        if (
+          !this.enabled.modifiers   ||
+          expression.isConstructor  ||
+          this.modifierWhitelist.includes(modifier.name)
+        ) {
+          continue;
+        }
 
         this.addNewBranch(contract, modifier);
         this._createInjectionPoint(
diff --git a/lib/validator.js b/lib/validator.js
index e3f6a9a..1cf4f9e 100644
--- a/lib/validator.js
+++ b/lib/validator.js
@@ -42,6 +42,11 @@ const configSchema = {
       type: "array",
       items: {type: "string"}
     },
+
+    modifierWhitelist: {
+      type: "array",
+      items: {type: "string"}
+    }
   },
 };
 
diff --git a/test/units/modifiers.js b/test/units/modifiers.js
index 4ff0420..de8a53d 100644
--- a/test/units/modifiers.js
+++ b/test/units/modifiers.js
@@ -13,7 +13,10 @@ describe('modifiers', () => {
     api = new Api({silent: true});
     await api.ganache(client);
   })
-  beforeEach(() => coverage = new Coverage());
+  beforeEach(() => {
+    api.config = {};
+    coverage = new Coverage()
+  });
   after(async() => await api.finish());
 
   async function setupAndRun(solidityFile){
@@ -97,6 +100,25 @@ describe('modifiers', () => {
     });
   });
 
+  // Same test as above - should have 2 fewer branches
+  it('should exclude whitelisted modifiers', async function() {
+    api.config.modifierWhitelist = ['mmm', 'nnn'];
+    const mapping = await setupAndRun('modifiers/multiple-mods-same-fn');
+
+    assert.deepEqual(mapping[util.filePath].l, {
+      5: 1, 6: 1, 10: 1, 11: 1, 15: 1
+    });
+    assert.deepEqual(mapping[util.filePath].b, {
+      1: [1, 0], 2: [1, 0]
+    });
+    assert.deepEqual(mapping[util.filePath].s, {
+      1: 1, 2: 1, 3: 1
+    });
+    assert.deepEqual(mapping[util.filePath].f, {
+      1: 1, 2: 1, 3: 1
+    });
+  });
+
   it('should cover multiple functions which use the same modifier', async function() {
     const contract = await util.bootstrapCoverage('modifiers/multiple-fns-same-mod', api);
     coverage.addContract(contract.instrumented, util.filePath);
diff --git a/test/units/validator.js b/test/units/validator.js
index bb81f68..1cb74c0 100644
--- a/test/units/validator.js
+++ b/test/units/validator.js
@@ -119,6 +119,7 @@ describe('config validation', () => {
     const options =  [
       "skipFiles",
       "istanbulReporter",
+      "modifierWhitelist"
     ]
 
     options.forEach(name => {