Concurrency-safe stubbing for TypeScript / Node.js

stub.ts

type QueueItem = {
  filename: string
  resolve: () => void
  originalValue: any
  stubValue: any
}
type Stub = {
  queue: QueueItem[]
  filename: string
  originalValue: any
}

const currentStubs = new Map<any, Record<string, Stub | undefined>>()

export function makeStub(filename: string) {
  return function stub<T, K extends keyof T>(obj: T, key: K, stubValue: T[K]) {
    let stubs = currentStubs.get(obj)
    if (!stubs) {
      stubs = {}
      currentStubs.set(obj, stubs)
    }

    const stub = stubs[key as string]
    const originalValue = key in obj ? obj[key] : '#delete#'
    if (stub) {
      //
      // Another parallel test is currently stubbing this object+key.
      // Add self to the queue.
      //

      return new Promise(resolve => {
        stub.queue.push({ filename, originalValue, stubValue, resolve })
      })
    }
    else {
      stubs[key as string] = {
        queue: [],
        filename,
        originalValue,
      }
    }
    // Perform the stub.
    obj[key] = stubValue;
    return
  }
}

export function resetStubs(filename: string) {
  for (let [obj, objStubs] of currentStubs.entries()) {
    for (let [key, stub] of Object.entries(objStubs)) {
      if (!stub || stub.filename !== filename) continue
      if (stub.originalValue === '#delete#') {
        delete obj[key];
      }
      else {
        obj[key] = stub.originalValue;
      }

      const next = stub.queue.shift()
      if (next) {
        //
        // Another test might be waiting for this stub to complete;
        // release now that we're done.
        // No need to batch since this function is synchronous,
        // and the queue is an array of Promise resolve functions.
        //
        stub.filename = next.filename
        stub.originalValue = next.originalValue
        obj[key] = next.stubValue
        next.resolve()
      }
      else {
        //
        // Queue was empty; no other tests are waiting for this key.
        // Clean up key so that stub() will find it empty on next call
        //
        delete objStubs[key]
      }
    }
    //
    // Cleanup
    //
    if (Object.keys(objStubs).length === 0) {
      currentStubs.delete(obj)
    }
  }
}