doc_ai_frontend/src/contexts/__tests__/AuthContext.oauth.test.tsx

import { useEffect } from 'react';
import { act, render, waitFor } from '@testing-library/react';
import { describe, expect, it, vi, beforeEach } from 'vitest';

import { AuthProvider, useAuth } from '../AuthContext';

const {
  loginMock,
  registerMock,
  logoutMock,
  restoreSessionMock,
  getGoogleOAuthUrlMock,
  exchangeGoogleCodeMock,
} = vi.hoisted(() => ({
  loginMock: vi.fn(),
  registerMock: vi.fn(),
  logoutMock: vi.fn(),
  restoreSessionMock: vi.fn(() => null),
  getGoogleOAuthUrlMock: vi.fn(),
  exchangeGoogleCodeMock: vi.fn(),
}));

vi.mock('../../lib/authService', () => ({
  authService: {
    login: loginMock,
    register: registerMock,
    logout: logoutMock,
    restoreSession: restoreSessionMock,
    getGoogleOAuthUrl: getGoogleOAuthUrlMock,
    exchangeGoogleCode: exchangeGoogleCodeMock,
  },
}));

function Harness({ onReady }: { onReady: (ctx: ReturnType<typeof useAuth>) => void }) {
  const auth = useAuth();

  useEffect(() => {
    onReady(auth);
  }, [auth, onReady]);

  return null;
}

function renderWithProvider(onReady: (ctx: ReturnType<typeof useAuth>) => void) {
  return render(
    <AuthProvider>
      <Harness onReady={onReady} />
    </AuthProvider>
  );
}

describe('AuthContext OAuth flow', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    sessionStorage.clear();
    localStorage.clear();
    restoreSessionMock.mockReturnValue(null);
  });

  it('beginGoogleOAuth writes state and redirect then redirects browser', async () => {
    getGoogleOAuthUrlMock.mockResolvedValue({ authUrl: 'https://accounts.google.com/o/oauth2/v2/auth' });
    let ctxRef: ReturnType<typeof useAuth> | null = null;

    renderWithProvider((ctx) => {
      ctxRef = ctx;
    });

    await waitFor(() => {
      expect(ctxRef).toBeTruthy();
    });

    await act(async () => {
      await (ctxRef as ReturnType<typeof useAuth>).beginGoogleOAuth();
    });

    expect(sessionStorage.getItem('texpixel_oauth_state')).toBeTruthy();
    expect(sessionStorage.getItem('texpixel_post_login_redirect')).toBe(window.location.href);
    expect(getGoogleOAuthUrlMock).toHaveBeenCalledTimes(1);
  });

  it('completeGoogleOAuth rejects when state mismatches', async () => {
    sessionStorage.setItem('texpixel_oauth_state', 'expected_state');
    let ctxRef: ReturnType<typeof useAuth> | null = null;

    renderWithProvider((ctx) => {
      ctxRef = ctx;
    });

    await waitFor(() => {
      expect(ctxRef).toBeTruthy();
    });

    let result: { error: Error | null } = { error: null };
    await act(async () => {
      result = await (ctxRef as ReturnType<typeof useAuth>).completeGoogleOAuth({
        code: 'abc',
        state: 'wrong_state',
        redirect_uri: 'http://localhost:5173/auth/google/callback',
      });
    });

    expect(result.error).toBeTruthy();
    expect(exchangeGoogleCodeMock).not.toHaveBeenCalled();
    expect(localStorage.getItem('texpixel_token')).toBeNull();
  });

  it('completeGoogleOAuth stores session on success', async () => {
    sessionStorage.setItem('texpixel_oauth_state', 'state_ok');
    exchangeGoogleCodeMock.mockImplementation(async () => {
      localStorage.setItem('texpixel_token', 'Bearer header.payload.sig');
      return {
        token: 'Bearer header.payload.sig',
        expiresAt: 1999999999,
        user: {
          user_id: 7,
          id: '7',
          email: 'oauth@example.com',
          exp: 1999999999,
          iat: 1111111,
        },
      };
    });

    let ctxRef: ReturnType<typeof useAuth> | null = null;

    renderWithProvider((ctx) => {
      ctxRef = ctx;
    });

    await waitFor(() => {
      expect(ctxRef).toBeTruthy();
    });

    let result: { error: Error | null } = { error: null };
    await act(async () => {
      result = await (ctxRef as ReturnType<typeof useAuth>).completeGoogleOAuth({
        code: 'code_ok',
        state: 'state_ok',
        redirect_uri: 'http://localhost:5173/auth/google/callback',
      });
    });

    expect(result.error).toBeNull();

    await waitFor(() => {
      expect((ctxRef as ReturnType<typeof useAuth>).isAuthenticated).toBe(true);
    });
    expect(localStorage.getItem('texpixel_token')).toBe('Bearer header.payload.sig');
  });

  it('completeGoogleOAuth deduplicates same code requests', async () => {
    sessionStorage.setItem('texpixel_oauth_state', 'state_ok');

    exchangeGoogleCodeMock.mockImplementation(
      () =>
        new Promise((resolve) => {
          setTimeout(() => {
            localStorage.setItem('texpixel_token', 'Bearer header.payload.sig');
            resolve({
              token: 'Bearer header.payload.sig',
              expiresAt: 1999999999,
              user: {
                user_id: 7,
                id: '7',
                email: 'oauth@example.com',
                exp: 1999999999,
                iat: 1111111,
              },
            });
          }, 20);
        })
    );

    let ctxRef: ReturnType<typeof useAuth> | null = null;

    renderWithProvider((ctx) => {
      ctxRef = ctx;
    });

    await waitFor(() => {
      expect(ctxRef).toBeTruthy();
    });

    let result1: { error: Error | null } = { error: null };
    let result2: { error: Error | null } = { error: null };

    await act(async () => {
      const p1 = (ctxRef as ReturnType<typeof useAuth>).completeGoogleOAuth({
        code: 'same_code',
        state: 'state_ok',
        redirect_uri: 'http://localhost:5173/auth/google/callback',
      });

      const p2 = (ctxRef as ReturnType<typeof useAuth>).completeGoogleOAuth({
        code: 'same_code',
        state: 'state_ok',
        redirect_uri: 'http://localhost:5173/auth/google/callback',
      });

      [result1, result2] = await Promise.all([p1, p2]);
    });

    expect(result1.error).toBeNull();
    expect(result2.error).toBeNull();
    expect(exchangeGoogleCodeMock).toHaveBeenCalledTimes(1);
  });
});