From bada13ed95b377d3ce1ae377385710cf98791481 Mon Sep 17 00:00:00 2001
From: TomokiMiyauci <development.operation.6.6.6@gmail.com>
Date: Mon, 13 Dec 2021 20:31:05 +0900
Subject: [PATCH] feat(mock): add `onceImplementation` function to mock object

---
 mock/fn.ts      | 52 ++++++++++++++++++++++++---------
 mock/fn_test.ts | 76 ++++++++++++++++++++++++++++++++++++++++++++++++-
 mock/mock.ts    |  3 ++
 3 files changed, 116 insertions(+), 15 deletions(-)

diff --git a/mock/fn.ts b/mock/fn.ts
index f4bfceb..2084012 100644
--- a/mock/fn.ts
+++ b/mock/fn.ts
@@ -3,25 +3,36 @@
 
 import { Mock } from "./mock.ts";
 import { incrementalNumber } from "./utils.ts";
-import { head } from "../matcher/utils.ts";
 import type { MockObject } from "./mock.ts";
 
-/** make mock object */
+/** store fn internal implementation */
+class MockFnStore {
+  private onceImplementations: ((...args: unknown[]) => unknown)[] = [];
+  constructor(
+    private defaultImplementation?: ((...args: unknown[]) => unknown),
+  ) {}
+
+  pickImplementation(): ((...args: unknown[]) => unknown) | undefined {
+    return this.onceImplementations.shift() ?? this.defaultImplementation;
+  }
+}
+
+/** make mock object with implementation function */
 function fn<A extends readonly unknown[], R>(
   implementation: (...args: A) => R,
 ): MockObject<A, R>;
+/** make mock object */
 function fn(): MockObject;
 function fn(
   implementation?: (...args: unknown[]) => unknown,
 ): MockObject {
   const mock = new Mock();
-  let implementations:
-    ({ implementation: (...args: unknown[]) => unknown } & IsDefault)[] =
-      implementation ? [{ isDefault: true, implementation }] : [];
+  const mockFnStore = new MockFnStore(implementation);
 
   /** Calls a method of an object, substituting another object for the current object.  */
   const call = (...args: unknown[]): unknown => {
-    const value = head(implementations)?.implementation?.(...args);
+    const implementation = mockFnStore.pickImplementation();
+    const value = implementation?.(...args);
     mock.add({
       args,
       result: {
@@ -33,10 +44,22 @@ function fn(
     return value;
   };
 
+  /** Sets the default mock function. The set function will be called when the mock object is called.  */
   const setImplementation = (
     implementation: (...args: unknown[]) => unknown,
   ): MockObject => {
-    implementations = [{ implementation, isDefault: true }];
+    mockFnStore["defaultImplementation"] = implementation;
+    return call as MockObject;
+  };
+
+  /** Sets a mock function to be called only once.
+   * This takes precedence over the default mock function.
+   * If there is more than one once implementation, they will be called in the order of registration.
+   */
+  const onceImplementation = (
+    implementation: (...args: unknown[]) => unknown,
+  ): MockObject => {
+    mockFnStore["onceImplementations"].unshift(implementation);
     return call as MockObject;
   };
 
@@ -47,16 +70,17 @@ function fn(
     },
   });
 
-  Object.defineProperty(call, "setImplementation", {
-    value: setImplementation,
+  Object.defineProperties(call, {
+    setImplementation: {
+      value: setImplementation,
+    },
+    onceImplementation: {
+      value: onceImplementation,
+    },
   });
 
   return call as MockObject;
 }
 
-export { fn };
+export { fn, MockFnStore };
 export type { MockObject };
-
-type IsDefault = {
-  isDefault: boolean;
-};
diff --git a/mock/fn_test.ts b/mock/fn_test.ts
index 1ba129e..7dc9786 100644
--- a/mock/fn_test.ts
+++ b/mock/fn_test.ts
@@ -1,6 +1,6 @@
 // Copyright 2021-Present the Unitest authors. All rights reserved. MIT license.
 
-import { fn } from "./fn.ts";
+import { fn, MockFnStore } from "./fn.ts";
 import { isFunction } from "../deps.ts";
 import { assert, assertEquals, assertExists } from "../dev_deps.ts";
 
@@ -77,3 +77,77 @@ Deno.test("setImplementation", () => {
   assertEquals(f1.mock.calls.length, 1);
   assertEquals(f2.mock.calls.length, 1);
 });
+
+Deno.test("onceImplementation should call only one time", () => {
+  assertExists(fn().onceImplementation);
+
+  const mockObject = fn();
+
+  const onceImplementation = fn();
+  mockObject.onceImplementation(onceImplementation);
+
+  assertEquals(mockObject(), undefined);
+  assertEquals(onceImplementation.mock.calls.length, 1);
+  mockObject();
+  assertEquals(onceImplementation.mock.calls.length, 1);
+});
+
+Deno.test("onceImplementation should be called in preference to default implementation", () => {
+  const defaultImplementation = fn();
+  const onceImplementation = fn();
+  const mockObject = fn(defaultImplementation).onceImplementation(
+    onceImplementation,
+  );
+
+  assertEquals(defaultImplementation.mock.calls.length, 0);
+  assertEquals(onceImplementation.mock.calls.length, 0);
+
+  mockObject();
+
+  assertEquals(defaultImplementation.mock.calls.length, 0);
+  assertEquals(onceImplementation.mock.calls.length, 1);
+
+  mockObject();
+  assertEquals(defaultImplementation.mock.calls.length, 1);
+  assertEquals(onceImplementation.mock.calls.length, 1);
+
+  mockObject();
+  assertEquals(defaultImplementation.mock.calls.length, 2);
+  assertEquals(onceImplementation.mock.calls.length, 1);
+});
+
+Deno.test("MockFnStore", () => {
+  const store = new MockFnStore();
+  assertExists(store["pickImplementation"]);
+  assertEquals(store["onceImplementations"], []);
+  assertEquals(store["defaultImplementation"], undefined);
+  assertEquals(store.pickImplementation(), undefined);
+});
+
+Deno.test("MockFnStore should return picked implementation", () => {
+  const mockObject = fn();
+  const store = new MockFnStore(mockObject);
+
+  assertEquals(
+    store["defaultImplementation"],
+    mockObject,
+  );
+  assertEquals(
+    store.pickImplementation(),
+    mockObject,
+  );
+  assertEquals(
+    store.pickImplementation(),
+    mockObject,
+  );
+  const mockObject2 = fn();
+  store["onceImplementations"].unshift(mockObject2);
+  assertEquals(
+    store.pickImplementation(),
+    mockObject2,
+  );
+  assertEquals(
+    store.pickImplementation(),
+    mockObject,
+  );
+});
diff --git a/mock/mock.ts b/mock/mock.ts
index 85a8a10..576b2c2 100644
--- a/mock/mock.ts
+++ b/mock/mock.ts
@@ -38,6 +38,9 @@ interface MockObject<A extends readonly unknown[] = any[], R = unknown> {
   setImplementation(
     implementation: (...args: A) => R,
   ): MockObject<A, R>;
+  onceImplementation(
+    implementation: (...args: A) => R,
+  ): MockObject<A, R>;
 }
 
 /** mock result store */