From d926649d29d34bc7d6048ac91be9d9bd3e67327e Mon Sep 17 00:00:00 2001
From: Ryo Igarashi <n33t5hin@gmail.com>
Date: Thu, 22 Feb 2024 16:30:07 +0900
Subject: [PATCH] feat: Implement Disposable for mastodon.streaming.Client

---
 examples/timeline-with-streaming.ts       |  5 +++--
 jest.config.js                            |  1 +
 src/adapters/action/dispatcher-ws.spec.ts | 16 +++++++++++++++-
 src/adapters/action/dispatcher-ws.ts      |  4 ++++
 src/adapters/action/proxy.spec.ts         | 14 ++++++++++++++
 src/adapters/action/proxy.ts              |  3 +++
 src/interfaces/action.ts                  |  1 +
 src/mastodon/streaming/client.ts          | 11 +++--------
 test-utils/jest-setup-after-env-unit.ts   |  4 ++++
 9 files changed, 48 insertions(+), 11 deletions(-)
 create mode 100644 test-utils/jest-setup-after-env-unit.ts

diff --git a/examples/timeline-with-streaming.ts b/examples/timeline-with-streaming.ts
index 264663d3..09f769bf 100644
--- a/examples/timeline-with-streaming.ts
+++ b/examples/timeline-with-streaming.ts
@@ -1,14 +1,15 @@
 import { createStreamingAPIClient } from "masto";
 
 const subscribe = async (): Promise<void> => {
-  const masto = createStreamingAPIClient({
+  using masto = createStreamingAPIClient({
     streamingApiUrl: "<STREAMING API URL>",
     accessToken: "<TOKEN>",
   });
 
+  using events = masto.hashtag.subscribe({ tag: "mastojs" });
   console.log("subscribed to #mastojs");
 
-  for await (const event of masto.hashtag.subscribe({ tag: "mastojs" })) {
+  for await (const event of events) {
     switch (event.event) {
       case "update": {
         console.log("new post", event.payload.content);
diff --git a/jest.config.js b/jest.config.js
index d4bb93b3..1b159809 100644
--- a/jest.config.js
+++ b/jest.config.js
@@ -6,6 +6,7 @@ export default {
       testEnvironment: "node",
       testMatch: ["<rootDir>/src/**/*.spec.ts"],
       transform: { "^.+\\.tsx?$": "ts-jest" },
+      setupFilesAfterEnv: ["<rootDir>/test-utils/jest-setup-after-env-unit.ts"],
     },
     {
       displayName: "e2e",
diff --git a/src/adapters/action/dispatcher-ws.spec.ts b/src/adapters/action/dispatcher-ws.spec.ts
index 21ecb9c0..bb98574b 100644
--- a/src/adapters/action/dispatcher-ws.spec.ts
+++ b/src/adapters/action/dispatcher-ws.spec.ts
@@ -22,6 +22,20 @@ describe("DispatcherWs", () => {
         data: undefined,
         meta: {},
       });
-    }).toThrowError(MastoUnexpectedError);
+    }).toThrow(MastoUnexpectedError);
+  });
+
+  it("can be disposed", () => {
+    const connector = new WebSocketConnectorImpl({
+      constructorParameters: ["wss://example.com"],
+    });
+    const dispatcher = new WebSocketActionDispatcher(
+      connector,
+      new SerializerNativeImpl(),
+      createLogger("error"),
+    );
+
+    dispatcher[Symbol.dispose]();
+    expect(connector.canAcquire()).toBe(false);
   });
 });
diff --git a/src/adapters/action/dispatcher-ws.ts b/src/adapters/action/dispatcher-ws.ts
index 73eda2c6..07d69a70 100644
--- a/src/adapters/action/dispatcher-ws.ts
+++ b/src/adapters/action/dispatcher-ws.ts
@@ -45,4 +45,8 @@ export class WebSocketActionDispatcher
       { ...data },
     ) as T;
   }
+
+  [Symbol.dispose](): void {
+    this.connector.close();
+  }
 }
diff --git a/src/adapters/action/proxy.spec.ts b/src/adapters/action/proxy.spec.ts
index 8c332453..b4338ae5 100644
--- a/src/adapters/action/proxy.spec.ts
+++ b/src/adapters/action/proxy.spec.ts
@@ -102,4 +102,18 @@ describe("RequestBuilder", () => {
     expect(() => api()).toThrow(TypeError);
     expect(() => api.close()).not.toThrow(TypeError);
   });
+
+  it("can be disposed", () => {
+    const dispatcher = {
+      dispatch: async <T>(_: AnyAction) => {
+        return {} as T;
+      },
+      [Symbol.dispose]: jest.fn(),
+    };
+    const api: any = createActionProxy(dispatcher, { context: ["root"] });
+
+    api[Symbol.dispose]();
+
+    expect(dispatcher[Symbol.dispose]).toHaveBeenCalled();
+  });
 });
diff --git a/src/adapters/action/proxy.ts b/src/adapters/action/proxy.ts
index 8878bf74..4011d696 100644
--- a/src/adapters/action/proxy.ts
+++ b/src/adapters/action/proxy.ts
@@ -60,6 +60,9 @@ const get =
     if (typeof property === "string" && SPECIAL_PROPERTIES.has(property)) {
       return;
     }
+    if (property === Symbol.dispose) {
+      return actionDispatcher[Symbol.dispose];
+    }
     if (typeof property === "symbol") {
       return;
     }
diff --git a/src/interfaces/action.ts b/src/interfaces/action.ts
index 9612578b..f99d3a66 100644
--- a/src/interfaces/action.ts
+++ b/src/interfaces/action.ts
@@ -12,6 +12,7 @@ export type AnyAction = Action<string>;
 
 export interface ActionDispatcher<T extends AnyAction> {
   dispatch<U>(action: T): U | Promise<U>;
+  [Symbol.dispose]?(): void;
 }
 
 export interface ActionDispatcherHook<T extends AnyAction, U = unknown> {
diff --git a/src/mastodon/streaming/client.ts b/src/mastodon/streaming/client.ts
index 4a3cc62f..9f44985e 100644
--- a/src/mastodon/streaming/client.ts
+++ b/src/mastodon/streaming/client.ts
@@ -8,18 +8,13 @@ export interface SubscribeHashtagParams {
   readonly tag: string;
 }
 
-export interface Subscription {
-  unsubscribe(): void;
+export interface Subscription extends AsyncIterable<Event>, Disposable {
   values(): AsyncIterableIterator<Event>;
-  [Symbol.asyncIterator](): AsyncIterator<Event, undefined>;
 
-  /**
-   * @experimental This is an experimental API.
-   */
-  [Symbol.dispose](): void;
+  unsubscribe(): void;
 }
 
-export interface Client {
+export interface Client extends Disposable {
   public: {
     subscribe(): Subscription;
     media: {
diff --git a/test-utils/jest-setup-after-env-unit.ts b/test-utils/jest-setup-after-env-unit.ts
new file mode 100644
index 00000000..3b6d869d
--- /dev/null
+++ b/test-utils/jest-setup-after-env-unit.ts
@@ -0,0 +1,4 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+(Symbol as any).dispose ??= Symbol("Symbol.dispose");
+(Symbol as any).asyncDispose ??= Symbol("Symbol.asyncDispose");
+/* eslint-enable @typescript-eslint/no-explicit-any */