fix(check): fix check draft

unadlib · Mar 26, 2023 · 9f32123 · 9f32123
1 parent 5cb1885
commit 9f32123
Show file tree

Hide file tree

Showing 4 changed files with 125 additions and 20 deletions.
diff --git a/README.md b/README.md
@@ -304,7 +304,7 @@ expect(isDraftable(baseState.list)).toBeTruthy();
 
 ### `rawReturn()`
 
-It is used as a raw return value to ensure that this value replaces the finalized draft value or use it to return `undefined` explicitly.
+For return values that do not contain any drafts, you can use `rawReturn()` to wrap this return value to improve performance. It ensure that the return value is only returned explicitly.
 
 ```ts
 const baseState = { id: 'test' };
@@ -328,6 +328,29 @@ const state = create(baseState, (draft) => {
 });
 expect(state).toEqual({ a: 2, b: { c: 1 } });
 expect(isDraft(state.b)).toBeFalsy();
+```
+
+ If you use `rawReturn()`, we recommend that you enable `strict` mode in development.
+
+```ts
+const baseState = { a: 1, b: { c: 1 } };
+const state = create(
+  baseState,
+  (draft) => {
+    if (draft.b.c === 1) {
+      return rawReturn({
+        ...draft,
+        a: 2,
+      });
+    }
+  },
+  {
+    strict: true,
+  }
+);
+// it will warn `The return value contains drafts, please don't use 'rawReturn()' to wrap the return value.` in strict mode.
+expect(state).toEqual({ a: 2, b: { c: 1 } });
+expect(isDraft(state.b)).toBeFalsy();
 ```
 
 [View more API docs](./docs/README.md).

diff --git a/src/create.ts b/src/create.ts
@@ -149,13 +149,13 @@ function create(arg0: any, arg1: any, arg2?: any): any {
       if (rawReturnValue) {
         const _value = rawReturnValue[0];
         if (_options.strict && typeof value === 'object' && value !== null) {
-          handleReturnValue(proxyDraft, value, true);
+          handleReturnValue({ rootDraft: proxyDraft, value, useRawReturn: true });
         }
         return finalize([_value]);
       }
       if (value !== undefined) {
         if (typeof value === 'object' && value !== null) {
-          handleReturnValue(proxyDraft, value);
+          handleReturnValue({ rootDraft: proxyDraft, value });
         }
         return finalize([value]);
       }

diff --git a/src/current.ts b/src/current.ts
@@ -11,12 +11,14 @@ import {
   shallowCopy,
 } from './utils';
 
-export function handleReturnValue<T extends object>(
-  rootDraft: ProxyDraft<any> | undefined,
-  value: T,
-  warning = false
-) {
-  let isContainDraft = false;
+export function handleReturnValue<T extends object>(options: {
+  rootDraft: ProxyDraft<any> | undefined;
+  value: T;
+  useRawReturn?: boolean;
+  isContainDraft?: boolean;
+  isRoot?: boolean;
+}) {
+  const { rootDraft, value, useRawReturn = false, isRoot = true } = options;
   forEach(value, (key, item, source) => {
     const proxyDraft = getProxyDraft(item);
     // just handle the draft which is created by the same rootDraft
@@ -25,12 +27,7 @@ export function handleReturnValue<T extends object>(
       rootDraft &&
       proxyDraft.finalities === rootDraft.finalities
     ) {
-      isContainDraft = true;
-      if (__DEV__ && warning) {
-        console.warn(
-          `The return value contains drafts, please don't use 'rawReturn()' to wrap the return value.`
-        );
-      }
+      options.isContainDraft = true;
       const currentValue = proxyDraft.original;
       if (source instanceof Set) {
         const arr = Array.from(source);
@@ -42,13 +39,22 @@ export function handleReturnValue<T extends object>(
         set(source, key, currentValue);
       }
     } else if (typeof item === 'object' && item !== null) {
-      handleReturnValue(rootDraft, item, warning);
+      options.value = item;
+      options.isRoot = false;
+      handleReturnValue(options);
     }
   });
-  if (__DEV__ && warning && !isContainDraft) {
-    console.warn(
-      `The return value does not contain any draft, please use 'rawReturn()' to wrap the return value to improve performance.`
-    );
+  if (__DEV__ && isRoot) {
+    if (!options.isContainDraft)
+      console.warn(
+        `The return value does not contain any draft, please use 'rawReturn()' to wrap the return value to improve performance.`
+      );
+
+    if (useRawReturn) {
+      console.warn(
+        `The return value contains drafts, please don't use 'rawReturn()' to wrap the return value.`
+      );
+    }
   }
 }
 

diff --git a/test/rawReturn.test.ts b/test/rawReturn.test.ts
@@ -8,6 +8,10 @@
 /* eslint-disable @typescript-eslint/ban-ts-comment */
 import { create, isDraft, rawReturn } from '../src';
 
+afterEach(() => {
+  jest.clearAllMocks();
+});
+
 test('base', () => {
   const base = 3;
   expect(create(base, () => 4)).toBe(4);
@@ -418,6 +422,7 @@ test('does not finalize upvalue drafts', () => {
 });
 
 test('mixed draft', () => {
+  jest.spyOn(console, 'warn').mockImplementation(() => {});
   const baseState = { a: 1, b: { c: 1 } };
   const state = create(baseState, (draft) => {
     if (draft.b.c === 1) {
@@ -429,4 +434,75 @@ test('mixed draft', () => {
   });
   expect(state).toEqual({ a: 2, b: { c: 1 } });
   expect(isDraft(state.b)).toBeFalsy();
+  expect(console.warn).toBeCalledTimes(0);
+});
+
+test('mixed draft with rawReturn()', () => {
+  jest.spyOn(console, 'warn').mockImplementation(() => {});
+  const baseState = { a: 1, b: { c: 1 } };
+  const state = create(baseState, (draft) => {
+    if (draft.b.c === 1) {
+      return rawReturn({
+        ...draft,
+        a: 2,
+      });
+    }
+  });
+  expect(() => {
+    JSON.stringify(state);
+  }).toThrowError();
+  expect(() => {
+    isDraft(state.b);
+  }).toThrowError();
+  expect(console.warn).toBeCalledTimes(0);
+});
+
+test('mixed draft with rawReturn() and strict mode', () => {
+  jest.spyOn(console, 'warn').mockImplementation(() => {});
+  const baseState = { a: 1, b: { c: 1 } };
+  const state = create(
+    baseState,
+    (draft) => {
+      if (draft.b.c === 1) {
+        return rawReturn({
+          ...draft,
+          a: 2,
+        });
+      }
+    },
+    {
+      strict: true,
+    }
+  );
+  expect(state).toEqual({ a: 2, b: { c: 1 } });
+  expect(isDraft(state.b)).toBeFalsy();
+  expect(console.warn).toBeCalledTimes(1);
+  expect(console.warn).toBeCalledWith(
+    `The return value contains drafts, please don't use 'rawReturn()' to wrap the return value.`
+  );
+});
+
+test('no mixed draft with strict mode', () => {
+  jest.spyOn(console, 'warn').mockImplementation(() => {});
+  const baseState = { a: 1, b: { c: 1 } };
+  const state = create(
+    baseState,
+    (draft) => {
+      if (draft.b.c === 1) {
+        return {
+          ...baseState,
+          a: 2,
+        };
+      }
+    },
+    {
+      strict: true,
+    }
+  );
+  expect(state).toEqual({ a: 2, b: { c: 1 } });
+  expect(isDraft(state.b)).toBeFalsy();
+  expect(console.warn).toBeCalledTimes(1);
+  expect(console.warn).toBeCalledWith(
+    `The return value does not contain any draft, please use 'rawReturn()' to wrap the return value to improve performance.`
+  );
 });