Refine startPolling time option handling

README.md

@@ -22,6 +22,21 @@ Task.startPolling();
 await Task.schedule('sayHello', new Date(Date.now() + 1000));
 ```
 
+## Polling Options
+
+You can pass `getCurrentTime()` to `startPolling()` to supply a custom clock for polling,
+which is useful for demos or tests that advance time automatically.
+
+```javascript
+let currentTime = new Date('2023-06-01T00:00:00.000Z');
+
+Task.startPolling({
+  getCurrentTime: () => new Date(currentTime)
+});
+
+currentTime = new Date(currentTime.valueOf() + 60_000);
+```
+
 ## Params
 
 The 2nd param to `Task.schedule()` is an object that this framework will call the handler function with.
@@ -57,4 +72,4 @@ await Task.schedule(
   { name: 'Friend' },
   5000
 );
-```
+```

src/taskSchema.js

@@ -116,14 +116,20 @@ taskSchema.methods.sideEffect = async function sideEffect(fn, params) {
 taskSchema.statics.startPolling = function startPolling(options) {
   const interval = options?.interval ?? 1000;
   const workerName = options?.workerName;
-  const pollOptions = workerName ? { workerName } : null;
+  const getCurrentTime = options?.getCurrentTime;
+  const pollOptions = {
+    ...(workerName ? { workerName } : {}),
+    ...(getCurrentTime ? { getCurrentTime } : {})
+  };
   let cancelled = false;
   let timeout = null;
+  const Task = this;
   if (!this._cancel) {
     doPoll.call(this);
     this._cancel = () => {
       cancelled = true;
       clearTimeout(timeout);
+      Task._cancel = null;
     };
   }
   return this._cancel;
@@ -136,7 +142,7 @@ taskSchema.statics.startPolling = function startPolling(options) {
     const Task = this;
 
     // Expire tasks that have timed out (refactored to separate function)
-    await Task.expireTimedOutTasks();
+    await Task.expireTimedOutTasks({ getCurrentTime });
 
     this._currentPoll = this.poll(pollOptions);
     await this._currentPoll.then(
@@ -152,8 +158,9 @@ taskSchema.statics.startPolling = function startPolling(options) {
 };
 
 // Refactor logic for expiring timed out tasks here
-taskSchema.statics.expireTimedOutTasks = async function expireTimedOutTasks() {
-  const now = time.now();
+taskSchema.statics.expireTimedOutTasks = async function expireTimedOutTasks(options = {}) {
+  const getCurrentTime = options.getCurrentTime;
+  const now = typeof getCurrentTime === 'function' ? getCurrentTime() : time.now();
   const Task = this;
   while (true) {
     const task = await Task.findOneAndUpdate(
@@ -251,13 +258,14 @@ taskSchema.statics.removeAllHandlers = function removeAllHandlers() {
 taskSchema.statics.poll = async function poll(opts) {
   const parallel = (opts && opts.parallel) || 1;
   const workerName = opts?.workerName;
+  const getCurrentTime = opts?.getCurrentTime;
 
   const additionalParams = workerName ? { workerName } : {};
 
   while (true) {
     const tasksInProgress = [];
     for (let i = 0; i < parallel; ++i) {
-      const now = time.now();
+      const now = typeof getCurrentTime === 'function' ? getCurrentTime() : time.now();
       const task = await this.findOneAndUpdate(
         { status: 'pending', scheduledAt: { $lte: now } },
         {

test/task.test.js

@@ -199,6 +199,37 @@ describe('Task', function() {
     cancel();
   });
 
+  it('allows startPolling() to use getCurrentTime()', async function() {
+    let resolve;
+    const p = new Promise((_resolve) => {
+      resolve = _resolve;
+    });
+    Task.registerHandler('getAnswer', (params, task) => {
+      resolve({ task, params });
+      return 42;
+    });
+
+    await Task.schedule('getAnswer', new Date(time.now().valueOf() + 1000), {
+      question: 'calculating...'
+    });
+
+    const getCurrentTime = sinon.stub().callsFake(() => new Date(now.valueOf() + 2000));
+    cancel = Task.startPolling({ interval: 100, getCurrentTime });
+
+    const res = await p;
+    assert.ok(getCurrentTime.called);
+    assert.deepEqual(res.params, { question: 'calculating...' });
+
+    await Task._currentPoll;
+
+    const task = await Task.findById(res.task._id);
+    assert.ok(task);
+    assert.equal(task.status, 'succeeded');
+    assert.strictEqual(task.result, 42);
+
+    cancel();
+  });
+
   it('catches errors in task', async function() {
     let resolve;
     let reject;