cozsweet-frontend-nextjs/src/stores/chat/__tests__/chat-machine.test-utils.ts

import { fromCallback, fromPromise } from "xstate";

import { DEFAULT_CHARACTER_ID } from "@/data/constants/character";
import {
  ChatSendResponseSchema,
  UnlockPrivateResponseSchema,
  type ChatSendResponse,
  type UnlockPrivateResponse,
  type UnlockPrivateResponseInput,
} from "@/data/schemas/chat";
import type { ChatEvent } from "@/stores/chat/chat-events";
import { chatMachine } from "@/stores/chat/chat-machine";
import type {
  UnlockMessageOutput as MachineUnlockMessageOutput,
  UnlockMessageRequest,
} from "@/stores/chat/helper/unlock";
import type { ChatHistoryActorInput } from "@/stores/chat/machine/actors/history";
import type { UiMessage } from "@/stores/chat/ui-message";

export interface SendMessageHttpOutput {
  response: ChatSendResponse;
  reply: UiMessage | null;
}

export interface UnlockHistoryOutput {
  unlocked: boolean;
  reason: string;
  shortfallCredits: number;
  messages: UiMessage[];
}

export interface TestUnlockMessageOutput {
  messageId: string;
  response: UnlockPrivateResponse;
}

export const TEST_CHAT_MACHINE_INPUT = {
  characterId: DEFAULT_CHARACTER_ID,
};

export function makeChatSendResponse(): ChatSendResponse {
  return ChatSendResponseSchema.parse({
    reply: "",
    audioUrl: "",
    messageId: "msg-1",
    isGuest: false,
    timestamp: Date.now(),
    image: { type: null, url: null },
    lockDetail: {
      locked: false,
      reason: null,
    },
  });
}

export function makeUnlockPrivateResponse(
  overrides: Partial<UnlockPrivateResponseInput> = {},
): UnlockPrivateResponse {
  return UnlockPrivateResponseSchema.parse({
    unlocked: true,
    content: "unlocked content",
    audioUrl: "",
    reason: "ok",
    creditBalance: 90,
    creditsCharged: 10,
    requiredCredits: 10,
    shortfallCredits: 0,
    ...overrides,
  });
}

export function createTestChatMachine(
  options: {
    historyMessages?: UiMessage[];
    sendMessageHttpError?: Error;
    unlockHistoryOutput?: UnlockHistoryOutput;
    unlockHistoryError?: Error;
    unlockMessageOutput?: TestUnlockMessageOutput;
  } = {},
) {
  return chatMachine.provide({
    actors: {
      loadHistory: createLoadHistoryCallback(options.historyMessages ?? []),
      sendMessageHttp: fromPromise<
        SendMessageHttpOutput,
        { characterId: string; content: string }
      >(async () => {
        if (options.sendMessageHttpError) {
          throw options.sendMessageHttpError;
        }
        return {
          response: makeChatSendResponse(),
          reply: null,
        };
      }),
      httpMessageQueue: fromCallback<ChatEvent, { characterId: string }>(
        ({ receive, sendBack }) => {
          receive((event) => {
            if (event.type !== "ChatSendMessage") return;
            sendBack({ type: "ChatQueuedSendStarted" });
            sendBack({
              type: "ChatQueuedHttpDone",
              output: {
                response: makeChatSendResponse(),
                reply: null,
              },
            });
          });
          return () => undefined;
        },
      ),
      unlockHistory: fromPromise<UnlockHistoryOutput, { characterId: string }>(
        async () => {
          if (options.unlockHistoryError) {
            throw options.unlockHistoryError;
          }
          return {
            unlocked: true,
            reason: "ok",
            shortfallCredits: 0,
            messages: [],
            ...options.unlockHistoryOutput,
          };
        },
      ),
      unlockMessage: fromPromise<
        MachineUnlockMessageOutput,
        UnlockMessageRequest & { characterId: string }
      >(async ({ input }) => {
        const output = options.unlockMessageOutput ?? {
          messageId: input.messageId ?? input.displayMessageId,
          response: makeUnlockPrivateResponse(),
        };
        return {
          displayMessageId: input.displayMessageId,
          request: input,
          response: output.response,
        };
      }),
    },
  });
}

export function createLoadHistoryCallback(
  messages: UiMessage[],
  networkMessages: UiMessage[] = messages,
  pagination: { total: number; limit: number } = {
    total: networkMessages.length,
    limit: 50,
  },
) {
  return fromCallback<ChatEvent, ChatHistoryActorInput>(({ sendBack }) => {
    sendBack({
      type: "ChatLocalHistoryLoaded",
      output: {
        messages,
        localCount: messages.length,
      },
    });
    sendBack({
      type: "ChatNetworkHistoryLoaded",
      output: {
        messages: networkMessages,
        localOverwritten: true,
        localCount: messages.length,
        networkCount: networkMessages.length,
        total: pagination.total,
        limit: pagination.limit,
      },
    });
    return () => undefined;
  });
}