feat: ensure data utils work well with legacy relationship proxies (#9317)

* feat: ensure data utils work well with legacy relationship proxies

* fixup types

* fixup

* fixup

* fix lint

Author: runspired

packages/diagnostic/src/-define.ts

@@ -64,9 +64,11 @@ export function test<TC extends TestContext = TestContext>(name: string, cb: Tes
   assert(`Cannot add the same test name twice: ${name}`, !currentModule.tests.byName.has(name));
   Config.totals.tests++;
 
+  const testName = currentModule.moduleName + ' > ' + name;
   const testInfo = {
-    id: generateHash(currentModule.moduleName + ' > ' + name),
+    id: generateHash(testName),
     name,
+    testName,
     cb,
     skip: false,
     todo: false,
@@ -83,9 +85,11 @@ export function todo<TC extends TestContext = TestContext>(name: string, cb: Tes
   assert(`Cannot add the same test name twice: ${name}`, !currentModule.tests.byName.has(name));
   Config.totals.todo++;
 
+  const testName = currentModule.moduleName + ' > ' + name;
   const testInfo = {
-    id: generateHash(currentModule.moduleName + ' > ' + name),
+    id: generateHash(testName),
     name,
+    testName,
     cb,
     skip: false,
     todo: true,
@@ -102,9 +106,11 @@ export function skip<TC extends TestContext = TestContext>(name: string, cb: Tes
   assert(`Cannot add the same test name twice: ${name}`, !currentModule.tests.byName.has(name));
   Config.totals.skipped++;
 
+  const testName = currentModule.moduleName + ' > ' + name;
   const testInfo = {
-    id: generateHash(currentModule.moduleName + ' > ' + name),
+    id: generateHash(testName),
     name,
+    testName,
     cb,
     skip: true,
     todo: false,

packages/diagnostic/src/-types.ts

@@ -100,8 +100,12 @@ export type ModuleCallback<TC extends TestContext> = ((hooks: Hooks<TC>) => void
 export type TestCallback<TC extends TestContext> = (this: TC, assert: Diagnostic) => void | Promise<void>;
 
 export interface TestInfo<TC extends TestContext> {
+  /* A unique id for the test based on the hash of the full testName */
   id: string;
+  /* The name of the test, not including moduleName */
   name: string;
+  /* The full name of the test including moduleName */
+  testName: string;
   cb: TestCallback<TC>;
   skip: boolean;
   todo: boolean;

packages/diagnostic/src/internals/run.ts

@@ -16,7 +16,7 @@ export async function runTest<TC extends TestContext>(
   const testContext = {
     [PublicTestInfo]: {
       id: test.id,
-      name: test.name,
+      name: test.testName,
     },
   } as unknown as TC;
   const testReport: TestReport = {

packages/ember/src/-private/promise-state.ts

@@ -44,12 +44,24 @@ export class PromiseState<T = unknown, E = unknown> {
   }
 }
 
+const LegacyPromiseProxy = Symbol.for('LegacyPromiseProxy');
+type LegacyAwaitable<T, E> = { promise: Promise<T> | Awaitable<T, E>; [LegacyPromiseProxy]: true };
+
+function isLegacyAwaitable<T, E>(promise: object): promise is LegacyAwaitable<T, E> {
+  return LegacyPromiseProxy in promise && 'promise' in promise && promise[LegacyPromiseProxy] === true;
+}
+
+function getPromise<T, E>(promise: Promise<T> | Awaitable<T, E> | LegacyAwaitable<T, E>): Promise<T> | Awaitable<T, E> {
+  return isLegacyAwaitable(promise) ? promise.promise : promise;
+}
+
 export function getPromiseState<T = unknown, E = unknown>(promise: Promise<T> | Awaitable<T, E>): PromiseState<T, E> {
-  let state = PromiseCache.get(promise) as PromiseState<T, E> | undefined;
+  const _promise = getPromise(promise);
+  let state = PromiseCache.get(_promise) as PromiseState<T, E> | undefined;
 
   if (!state) {
-    state = new PromiseState(promise);
-    PromiseCache.set(promise, state);
+    state = new PromiseState(_promise);
+    PromiseCache.set(_promise, state);
   }
 
   return state;

packages/model/src/-private/promise-belongs-to.ts

@@ -23,6 +23,8 @@ export interface BelongsToProxyCreateArgs<T = unknown> {
   _belongsToState: BelongsToProxyMeta<T>;
 }
 
+export const LegacyPromiseProxy = Symbol.for('LegacyPromiseProxy');
+
 interface PromiseObjectType<T> extends PromiseProxyMixin<T | null>, ObjectProxy<T> {
   // eslint-disable-next-line @typescript-eslint/no-misused-new
   new <PT>(...args: unknown[]): PromiseObjectType<PT>;
@@ -80,6 +82,8 @@ class PromiseBelongsTo<T = unknown> extends Extended<T> {
     await legacySupport.reloadBelongsTo(key, options);
     return this;
   }
+
+  [LegacyPromiseProxy] = true as const;
 }
 
 export { PromiseBelongsTo };

packages/model/src/-private/promise-many-array.ts

@@ -6,6 +6,7 @@ import { compat } from '@ember-data/tracking';
 import { defineSignal } from '@ember-data/tracking/-private';
 
 import type ManyArray from './many-array';
+import { LegacyPromiseProxy } from './promise-belongs-to';
 
 export interface HasManyProxyCreateArgs<T = unknown> {
   promise: Promise<ManyArray<T>>;
@@ -195,6 +196,8 @@ export default class PromiseManyArray<T = unknown> {
   static create<T>({ promise, content }: HasManyProxyCreateArgs<T>): PromiseManyArray<T> {
     return new this(promise, content);
   }
+
+  [LegacyPromiseProxy] = true as const;
 }
 defineSignal(PromiseManyArray.prototype, 'content', null);
 defineSignal(PromiseManyArray.prototype, 'isPending', false);

tests/warp-drive__ember/tests/integration/get-promise-state-test.gts

@@ -1,11 +1,38 @@
 import { rerender, settled } from '@ember/test-helpers';
 
-import { createDeferred, setPromiseResult } from '@ember-data/request';
+import { createDeferred, setPromiseResult, type Awaitable } from '@ember-data/request';
 import type { RenderingTestContext } from '@warp-drive/diagnostic/ember';
 import { module, setupRenderingTest, test } from '@warp-drive/diagnostic/ember';
 import { getPromiseState } from '@warp-drive/ember';
 
 type PromiseState<T, E> = ReturnType<typeof getPromiseState<T, E>>;
+const SecretSymbol = Symbol.for('LegacyPromiseProxy');
+
+interface PromiseProxy<T, E> extends Promise<T> {}
+class PromiseProxy<T, E> {
+  [SecretSymbol]: true;
+  promise: Awaitable<T, E>;
+
+  constructor(promise: Awaitable<T, E>) {
+    this[SecretSymbol] = true;
+    this.promise = promise;
+  }
+
+  then<T1, T2>(
+    onFulfilled?: ((value: T) => unknown) | undefined | null,
+    onRejected?: ((error: E) => T2 | Promise<T2>) | undefined | null
+  ): Promise<T1 | T2> {
+    return this.promise.then(onFulfilled!, onRejected!) as Promise<T1 | T2>;
+  }
+
+  catch<T2>(onRejected: ((error: E) => T2 | Promise<T2>) | undefined | null): Promise<T2> {
+    return this.promise.catch(onRejected!) as Promise<T2>;
+  }
+
+  finally(onFinally: () => void): Promise<T> {
+    return this.promise.finally(onFinally) as Promise<T>;
+  }
+}
 
 module('Integration | get-promise-state', function (hooks) {
   setupRenderingTest(hooks);
@@ -196,4 +223,88 @@ module('Integration | get-promise-state', function (hooks) {
     assert.equal(counter, 1);
     assert.equal(this.element.textContent?.trim(), 'Our Error\n          Count:\n          1');
   });
+
+  test('it unwraps promise-proxies that utilize the secret symbol for error states', async function (this: RenderingTestContext, assert) {
+    const _promise = Promise.resolve().then(() => {
+      throw new Error('Our Error');
+    });
+    const promise = new PromiseProxy<never, Error>(_promise);
+
+    try {
+      getPromiseState(promise);
+      await promise;
+    } catch {
+      // do nothing
+    }
+
+    let state: PromiseState<string, Error>;
+    function _getPromiseState<T>(p: Promise<T>): PromiseState<T, Error> {
+      state = getPromiseState(p) as PromiseState<string, Error>;
+      return state as PromiseState<T, Error>;
+    }
+    let counter = 0;
+    function countFor(_result: unknown, _error: unknown) {
+      return ++counter;
+    }
+
+    await this.render(
+      <template>
+        {{#let (_getPromiseState promise) as |state|}}
+          {{#if state.isPending}}
+            Pending
+          {{else if state.isError}}
+            {{state.error.message}}
+          {{else if state.isSuccess}}
+            Invalid Success Reached
+          {{/if}}
+          <br />Count:
+          {{countFor state.result state.error}}{{/let}}
+      </template>
+    );
+
+    assert.equal(state!.result, null);
+    assert.true(state!.error instanceof Error);
+    assert.equal((state!.error as Error | undefined)?.message, 'Our Error');
+    assert.equal(counter, 1);
+    assert.equal(this.element.textContent?.trim(), 'Our Error\n          Count:\n          1');
+    await rerender();
+    assert.equal(state!.result, null);
+    assert.true(state!.error instanceof Error);
+    assert.equal((state!.error as Error | undefined)?.message, 'Our Error');
+    assert.equal(counter, 1);
+    assert.equal(this.element.textContent?.trim(), 'Our Error\n          Count:\n          1');
+    assert.equal(state!, getPromiseState(_promise));
+  });
+
+  test('it unwraps promise-proxies that utilize the secret symbol for success states', async function (this: RenderingTestContext, assert) {
+    const _promise = Promise.resolve().then(() => 'Our Data');
+    const promise = new PromiseProxy<string, Error>(_promise);
+    getPromiseState(promise);
+    await promise;
+
+    let state: PromiseState<string, Error>;
+    function _getPromiseState<T>(p: Promise<T>): PromiseState<T, Error> {
+      state = getPromiseState(p) as PromiseState<string, Error>;
+      return state as PromiseState<T, Error>;
+    }
+    let counter = 0;
+    function countFor(_result: unknown) {
+      return ++counter;
+    }
+
+    await this.render(
+      <template>
+        {{#let (_getPromiseState promise) as |state|}}
+          {{state.result}}<br />Count:
+          {{countFor state.result}}
+        {{/let}}
+      </template>
+    );
+
+    assert.equal(this.element.textContent?.trim(), 'Our DataCount:\n          1');
+    await settled();
+
+    assert.equal(this.element.textContent?.trim(), 'Our DataCount:\n          1');
+    assert.equal(state!, getPromiseState(_promise));
+  });
 });