From 4f7181fa56d45ff72755e0f8988ccd3a3628b0a3 Mon Sep 17 00:00:00 2001
From: Ilya Ashikhmin <ivashikhmin@gmail.com>
Date: Sun, 29 Dec 2024 23:05:30 +0100
Subject: [PATCH] fix(signals): patch state methods do not correctly resolve
 generic collection names withEntities

Ensures type-safety of named entity collections with patchState helpers

fixes #4638
---
 .../entities/spec/types/models.types.spec.ts  | 36 +++++++++++++++++++
 modules/signals/entities/src/models.ts        | 11 ++++--
 2 files changed, 44 insertions(+), 3 deletions(-)
 create mode 100644 modules/signals/entities/spec/types/models.types.spec.ts

diff --git a/modules/signals/entities/spec/types/models.types.spec.ts b/modules/signals/entities/spec/types/models.types.spec.ts
new file mode 100644
index 0000000000..a67029e317
--- /dev/null
+++ b/modules/signals/entities/spec/types/models.types.spec.ts
@@ -0,0 +1,36 @@
+import { expecter } from 'ts-snippet';
+import { compilerOptions } from './helpers';
+
+describe('models', () => {
+  const expectSnippet = expecter(
+    (code) => `
+        import { patchState, signalStore, type, withMethods } from '@ngrx/signals';
+        import { addEntity, entityConfig, withEntities } from '@ngrx/signals/entities';
+        type User = { id: number; name: string };
+
+        ${code}
+      `,
+    compilerOptions()
+  );
+
+  it('succeeds when entity collection name is passed as the type parameter', () => {
+    const snippet = `
+    const storeFactory = <Collection extends string>(
+      collection: Collection
+    ) => {
+      const Store = signalStore(
+        withEntities({ entity: type<User>(), collection: collection }),
+        withMethods((store) => ({
+          addUsers(user: User): void {
+            patchState(store, addEntity(user, { collection: collection }));
+          },
+        }))
+      );
+
+      return new Store();
+    };
+`;
+
+    expectSnippet(snippet).toSucceed();
+  });
+});
diff --git a/modules/signals/entities/src/models.ts b/modules/signals/entities/src/models.ts
index 956311b50b..bc3b937b75 100644
--- a/modules/signals/entities/src/models.ts
+++ b/modules/signals/entities/src/models.ts
@@ -1,4 +1,5 @@
 import { Signal } from '@angular/core';
+import { Prettify } from '@ngrx/signals';
 
 export type EntityId = string | number;
 
@@ -9,9 +10,13 @@ export type EntityState<Entity> = {
   ids: EntityId[];
 };
 
-export type NamedEntityState<Entity, Collection extends string> = {
-  [K in keyof EntityState<Entity> as `${Collection}${Capitalize<K>}`]: EntityState<Entity>[K];
-};
+export type NamedEntityState<Entity, Collection extends string> = Prettify<
+  {
+    [K in `${Collection}EntityMap`]: EntityMap<Entity>;
+  } & {
+    [K in `${Collection}Ids`]: EntityId[];
+  }
+>;
 
 export type EntityProps<Entity> = {
   entities: Signal<Entity[]>;