From 33e0197ad50103e9b9a58b5567f9f8436d2a587f Mon Sep 17 00:00:00 2001
From: Christopher Willis-Ford <cwillisf@media.mit.edu>
Date: Fri, 31 Aug 2018 12:38:06 -0700
Subject: [PATCH] Add tests for TokenBucket

---
 src/util/token-bucket.js       |  4 ++--
 test/fixtures/test-compare.js  | 15 +++++++++++++++
 test/unit/util_token-bucket.js | 31 +++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 test/fixtures/test-compare.js
 create mode 100644 test/unit/util_token-bucket.js

diff --git a/src/util/token-bucket.js b/src/util/token-bucket.js
index b5449bff9..0a1c20686 100644
--- a/src/util/token-bucket.js
+++ b/src/util/token-bucket.js
@@ -117,8 +117,8 @@ class TokenBucket {
         if (cost <= this._tokenCount) {
             return Promise.resolve();
         }
-        if (!(cost <= this._limit)) {
-            return Promise.reject(new Error('Task cost is greater than bucket limit'));
+        if (!(cost <= this._maxTokens)) {
+            return Promise.reject(new Error(`Task cost ${cost} is greater than bucket limit ${this._maxTokens}`));
         }
         return new Promise(resolve => {
             const tokensNeeded = this._tokenCount - cost;
diff --git a/test/fixtures/test-compare.js b/test/fixtures/test-compare.js
new file mode 100644
index 000000000..4aa625ea7
--- /dev/null
+++ b/test/fixtures/test-compare.js
@@ -0,0 +1,15 @@
+const testCompare = (t, lhs, op, rhs, message) => {
+    const details = `Expected: ${lhs} ${op} ${rhs}`;
+    const extra = {details};
+    switch (op) {
+    case '<': return t.ok(lhs < rhs, message, extra);
+    case '<=': return t.ok(lhs <= rhs, message, extra);
+    case '===': return t.ok(lhs === rhs, message, extra);
+    case '!==': return t.ok(lhs !== rhs, message, extra);
+    case '>=': return t.ok(lhs >= rhs, message, extra);
+    case '>': return t.ok(lhs > rhs, message, extra);
+    default: return t.fail(`Unrecognized op: ${op}`);
+    }
+};
+
+module.exports = testCompare;
diff --git a/test/unit/util_token-bucket.js b/test/unit/util_token-bucket.js
new file mode 100644
index 000000000..a70764de2
--- /dev/null
+++ b/test/unit/util_token-bucket.js
@@ -0,0 +1,31 @@
+const test = require('tap').test;
+
+const Timer = require('../../src/util/timer');
+const TokenBucket = require('../../src/util/token-bucket');
+
+const testCompare = require('../fixtures/test-compare');
+
+test('constructor', t => {
+    // Max tokens = 1000, refill 1000 tokens per second (1 per millisecond), and start with 0 tokens
+    const bukkit = new TokenBucket(1000, 1000, 0);
+
+    const timer = new Timer();
+    timer.start();
+
+    const taskResults = [];
+    const promises = [];
+    promises.push(
+        bukkit.do(() => taskResults.push('a'), 100).then(() =>
+            testCompare(t, timer.timeElapsed(), '>=', 100, 'Costly task must wait')
+        ),
+        bukkit.do(() => taskResults.push('b'), 0).then(() =>
+            testCompare(t, timer.timeElapsed(), '<', 150, 'Cheap task should run soon')
+        ),
+        bukkit.do(() => taskResults.push('c'), 101).then(() =>
+            testCompare(t, timer.timeElapsed(), '>=', 200, 'Tasks must run in serial')
+        )
+    );
+    return Promise.all(promises).then(() => {
+        t.deepEqual(taskResults, ['a', 'b', 'c'], 'All tasks must run in correct order');
+    });
+});