SWR Subscription

src_index.ts

export { useSubscription } from "./useSubscription";
export { subscription } from "./subscriptionMiddleware";
export { SubscriptionConfig } from "./SubscriptionConfig";
export * from "./types";

src_SubscriptionConfig.tsx

import React, { createContext, useRef, useCallback, FC } from "react";
import type { SubscribeFn, Event, Key, Listener } from "./types";

export interface SubscriptionContextType {
  subscribe: (key: Key, event: Event, subscription: SubscribeFn) => void;
  unsubscribe: (key: Key) => void;
  emit: (key: Event, data: any) => void;
}

const SubscriptionContext = createContext<SubscriptionContextType | null>(null);

const SubscriptionConfig: FC = (props) => {
  const listeners = useRef(new Map<Key, Listener>());

  const subscribe = useCallback(
    (key: Key, event: Event, subscription: SubscribeFn) => {
      const listener = listeners.current.get(key);

      if (!listener) {
        listeners.current.set(key, {
          event,
          smartCounter: 1,
          callback: subscription,
        });

        return;
      }

      listener.smartCounter++;
      listener.callback = subscription;
    },
    []
  );

  const unsubscribe = useCallback((key: Key) => {
    const listener = listeners.current.get(key);

    if (!listener) {
      return console.warn(
        "an dispatch tryied to unsubscribe a nonexistent subscription"
      );
    }

    if (listener.smartCounter > 1) {
      listener.smartCounter--;

      return;
    }

    listeners.current.delete(key);
  }, []);

  const emit = useCallback((event: Event, data: any) => {
    listeners.current.forEach((listener) => {
      if (listener.event === event) {
        listener.callback(data);
      }
    });
  }, []);

  return (
    <SubscriptionContext.Provider
      value={{
        subscribe,
        unsubscribe,
        emit,
      }}
      {...props}
    />
  );
};

export { SubscriptionConfig, SubscriptionContext };

src_subscriptionMiddleware.ts

import { useEffect, useRef } from "react";
import { Middleware, SWRHook, unstable_serialize } from "swr";

import { useSubscription } from "./useSubscription";
import type { Event, SubscriptionFn } from "./types";

const subscriptionMiddleware =
  (event: Event, subscription: SubscriptionFn): Middleware =>
  (useSWRNext: SWRHook) => {
    return (key, fetcher, config) => {
      const listener = useSubscription();
      const swr = useSWRNext(key, fetcher, config);
      const latestSwr = useRef(swr);
      latestSwr.current = swr;

      const keyString = unstable_serialize(key);
      useEffect(() => {
        listener.subscribe(keyString, event, (data) =>
          subscription(data, latestSwr.current)
        );

        return () => listener.unsubscribe(keyString);
      }, [keyString]);

      return swr;
    };
  };

export { subscriptionMiddleware as subscription };

src_types.ts

import { SWRResponse } from "swr";

export type SubscribeFn = (...args: any[]) => void;
export type SubscriptionFn = (data: any, swr: SWRResponse) => void;
export type Event = string;
export type Key = string;

export interface Listener {
  event: Event;
  callback: SubscribeFn;
  smartCounter: number;
}

src_useSubscription.ts

import { useContext } from "react";
import { SubscriptionContext } from "./SubscriptionConfig";

const useSubscription = () => {
  const listener = useContext(SubscriptionContext);

  if (!listener) {
    throw new Error("You cannot call a hook outside the subscription Provider");
  }

  return listener;
};

export { useSubscription };

test_middleware.test.tsx

import useSWR, { useSWRConfig } from "swr";
import axios from "axios";
import {
  subscription,
  SubscriptionConfig,
  useSubscription,
} from "swr-subscription";

import { act, renderHook, waitFor } from "@testing-library/react";

import { server } from "./mocks/server";

const usePosts = (caller?: any) => {
  const { data, isValidating } = useSWR(
    "http://localhost/posts",
    async (url: string) => {
      const response = await axios.get<any[]>(url);

      return response.data;
    },
    {
      use: [
        subscription("post::update", (data, swr) => {
          caller?.();
          swr.mutate(
            swr.data.map((item) => (item.id === data.id ? data : item)),
            { revalidate: false }
          );
        }),
      ],
    }
  );
  const { emit } = useSubscription();

  return { data, isValidating, emit };
};

describe("subscription middleware", () => {
  beforeAll(() => server.listen());
  afterEach(() => {
    server.resetHandlers();
  });
  afterAll(() => server.close());

  test("updating data from subscription", async () => {
    const { result } = renderHook(usePosts, {
      wrapper: SubscriptionConfig,
    });

    await waitFor(() =>
      expect(result.current.data).toEqual([
        { id: "1", title: "Post 1" },
        { id: "2", title: "Post 2" },
      ])
    );

    act(() => {
      result.current.emit("post::update", { id: "2", title: "New Post 2" });
    });

    await waitFor(() =>
      expect(result.current.data).toEqual([
        { id: "1", title: "Post 1" },
        { id: "2", title: "New Post 2" },
      ])
    );
  });

  test("avoid calling subscription twice", async () => {
    const mockFnOne = jest.fn();
    const mockFnTwo = jest.fn();
    const { result: resultOne } = renderHook(() => usePosts(mockFnOne), {
      wrapper: SubscriptionConfig,
    });
    const { result: resultTwo } = renderHook(() => usePosts(mockFnTwo), {
      wrapper: SubscriptionConfig,
    });

    await waitFor(() =>
      expect(resultOne.current.data).toEqual([
        { id: "1", title: "Post 1" },
        { id: "2", title: "New Post 2" },
      ])
    );

    expect(resultTwo.current.data).toEqual([
      { id: "1", title: "Post 1" },
      { id: "2", title: "New Post 2" },
    ]);

    act(() => {
      resultOne.current.emit("post::update", { id: "1", title: "New Post 1" });
    });

    await waitFor(() =>
      expect(resultOne.current.data).toEqual([
        { id: "1", title: "New Post 1" },
        { id: "2", title: "New Post 2" },
      ])
    );

    expect(mockFnOne).toHaveBeenCalled();
    expect(mockFnTwo).not.toHaveBeenCalled();
  });
});

test_mocks_handlers.ts

import { rest } from "msw";

const handlers = [
  rest.get("http://localhost/posts", (req, res, ctx) => {
    return res(
      ctx.status(200),
      ctx.json([
        { id: "1", title: "Post 1" },
        { id: "2", title: "Post 2" },
      ])
    );
  }),
];

export { handlers };

test_use-subscription.test.tsx

import { renderHook } from "@testing-library/react";

import { SubscriptionConfig, useSubscription } from "swr-subscription";

describe("Subscription Config provider", () => {
  test("context definition", () => {
    const { result } = renderHook(useSubscription, {
      wrapper: SubscriptionConfig,
    });

    expect(result.current).toBeDefined();
  });

  test("attach a subscription", () => {
    const { result } = renderHook(useSubscription, {
      wrapper: SubscriptionConfig,
    });
    const handler = jest.fn();

    result.current.subscribe("key", "message", handler);
    result.current.emit("message", "hello world");

    expect(handler).toHaveBeenCalledWith("hello world");
  });

  test("attach two subscriptions", () => {
    const { result } = renderHook(useSubscription, {
      wrapper: SubscriptionConfig,
    });
    const handlerOne = jest.fn();
    const handlerTwo = jest.fn();

    result.current.subscribe("key-one", "message", handlerOne);
    result.current.subscribe("key-two", "message", handlerTwo);
    result.current.emit("message", "hello world");

    expect(handlerOne).toHaveBeenCalledWith("hello world");
    expect(handlerTwo).toHaveBeenCalledWith("hello world");
  });

  test("avoid duplicated subscription", () => {
    const { result } = renderHook(useSubscription, {
      wrapper: SubscriptionConfig,
    });
    const handlerOne = jest.fn();
    const handlerTwo = jest.fn();

    result.current.subscribe("key", "message", handlerOne);
    result.current.subscribe("key", "message", handlerTwo);
    result.current.emit("message", "hello world");

    expect(handlerOne).not.toHaveBeenCalled();
    expect(handlerTwo).toHaveBeenCalledWith("hello world");
  });

  test("attach two subscriptions with same key and handle one unsubscribe", () => {
    const { result } = renderHook(useSubscription, {
      wrapper: SubscriptionConfig,
    });
    const handlerOne = jest.fn();
    const handlerTwo = jest.fn();

    result.current.subscribe("key", "message", handlerOne);
    result.current.subscribe("key", "message", handlerTwo);
    result.current.unsubscribe("key");
    result.current.emit("message", "hello world");

    expect(handlerOne).not.toHaveBeenCalled();
    expect(handlerTwo).toHaveBeenCalledWith("hello world");
  });

  test("attach two subscriptions with same key and unsubscribe all", () => {
    const { result } = renderHook(useSubscription, {
      wrapper: SubscriptionConfig,
    });
    const handlerOne = jest.fn();
    const handlerTwo = jest.fn();

    result.current.subscribe("key", "message", handlerOne);
    result.current.subscribe("key", "message", handlerTwo);
    result.current.unsubscribe("key");
    result.current.unsubscribe("key");
    result.current.emit("message", "hello world");

    expect(handlerOne).not.toHaveBeenCalled();
    expect(handlerTwo).not.toHaveBeenCalled();
  });
});