From 7b133723efefa99436dce696fed39099afb1cdb5 Mon Sep 17 00:00:00 2001
From: Chau Tran <nartc7789@gmail.com>
Date: Mon, 18 Sep 2023 11:20:30 -0500
Subject: [PATCH] fix(inject-destroy): add onDestroy to return value of
 injectDestroy (#69)

This PR augments the return type of `injectDestroy` to also include `onDestroy` so that folks don't have to `inject(DestroyRef)` when they need it to run arbitrary destroy logic using `DestroyRef#onDestroy`
---
 .../content/docs/utilities/inject-destroy.md  | 23 +++++++++++
 .../inject-destroy/src/inject-destroy.spec.ts | 40 ++++++++++++++++++-
 .../inject-destroy/src/inject-destroy.ts      | 14 +++++--
 3 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/docs/src/content/docs/utilities/inject-destroy.md b/docs/src/content/docs/utilities/inject-destroy.md
index 11ad4df2..10fc2234 100644
--- a/docs/src/content/docs/utilities/inject-destroy.md
+++ b/docs/src/content/docs/utilities/inject-destroy.md
@@ -52,6 +52,29 @@ export class MyComponent {
 
 As you can see, we don't need to implement `OnDestroy` anymore and we don't need to manually emit from the `Subject` when the component is destroyed.
 
+### `onDestroy`
+
+The value returned by `injectDestroy()` also includes `onDestroy()` function to register arbitrary destroy logic callbacks.
+
+```ts
+
+@Component({})
+export class MyComponent {
+  private dataService = inject(DataService);
+  private destroy$ = injectDestroy();
+
+  ngOnInit() {
+    this.dataService.getData()
+      .pipe(takeUntil(this.destroy$))
+      .subscribe(...);
+
+    this.destroy$.onDestroy(() => {
+      /* other destroy logics, similar to DestroyRef#onDestroy */
+    });
+  }
+}
+```
+
 ## How it works
 
 The helper functions injects the `DestroyRef` class from Angular, and on the `onDestroy` lifecycle hook, it emits from the `Subject` and completes it.
diff --git a/libs/ngxtension/inject-destroy/src/inject-destroy.spec.ts b/libs/ngxtension/inject-destroy/src/inject-destroy.spec.ts
index c95f3ba8..a43724df 100644
--- a/libs/ngxtension/inject-destroy/src/inject-destroy.spec.ts
+++ b/libs/ngxtension/inject-destroy/src/inject-destroy.spec.ts
@@ -9,7 +9,7 @@ import { interval, takeUntil } from 'rxjs';
 import { injectDestroy } from './inject-destroy';
 
 describe(injectDestroy.name, () => {
-	describe('emits when the component is destroyed', () => {
+	describe('emits when the component is destroyed using takeUntil', () => {
 		@Component({ standalone: true, template: '' })
 		class TestComponent implements OnInit {
 			destroy$ = injectDestroy();
@@ -45,4 +45,42 @@ describe(injectDestroy.name, () => {
 			expect(component.count).toBe(2);
 		}));
 	});
+
+	describe('emits when the component is destroyed using onDestroy', () => {
+		@Component({ standalone: true, template: '' })
+		class TestComponent implements OnInit {
+			destroy$ = injectDestroy();
+			count = 0;
+
+			ngOnInit() {
+				const sub = interval(1000).subscribe(() => this.count++);
+				this.destroy$.onDestroy(() => {
+					sub.unsubscribe();
+				});
+			}
+		}
+
+		let component: TestComponent;
+		let fixture: ComponentFixture<TestComponent>;
+
+		beforeEach(async () => {
+			fixture = TestBed.createComponent(TestComponent);
+			component = fixture.componentInstance;
+		});
+
+		it('should handle async stuff', fakeAsync(() => {
+			component.ngOnInit();
+
+			expect(component.count).toBe(0);
+			tick(1000);
+			expect(component.count).toBe(1);
+			tick(1000);
+			expect(component.count).toBe(2);
+
+			fixture.destroy(); // destroy the component here
+
+			tick(1000);
+			expect(component.count).toBe(2);
+		}));
+	});
 });
diff --git a/libs/ngxtension/inject-destroy/src/inject-destroy.ts b/libs/ngxtension/inject-destroy/src/inject-destroy.ts
index c87cfaff..8112d37a 100644
--- a/libs/ngxtension/inject-destroy/src/inject-destroy.ts
+++ b/libs/ngxtension/inject-destroy/src/inject-destroy.ts
@@ -1,8 +1,8 @@
 import {
 	DestroyRef,
 	inject,
-	Injector,
 	runInInjectionContext,
+	type Injector,
 } from '@angular/core';
 import { assertInjector } from 'ngxtension/assert-injector';
 import { ReplaySubject } from 'rxjs';
@@ -26,7 +26,9 @@ import { ReplaySubject } from 'rxjs';
  *   }
  * }
  */
-export const injectDestroy = (injector?: Injector) => {
+export const injectDestroy = (
+	injector?: Injector
+): ReplaySubject<void> & { onDestroy: DestroyRef['onDestroy'] } => {
 	injector = assertInjector(injectDestroy, injector);
 
 	return runInInjectionContext(injector, () => {
@@ -39,6 +41,12 @@ export const injectDestroy = (injector?: Injector) => {
 			subject$.complete();
 		});
 
-		return subject$;
+		Object.assign(subject$, {
+			onDestroy: destroyRef.onDestroy.bind(destroyRef),
+		});
+
+		return subject$ as ReplaySubject<void> & {
+			onDestroy: DestroyRef['onDestroy'];
+		};
 	});
 };