dify/web/app/components/header/plan-badge/__tests__/index.spec.tsx

import type { Mock } from 'vitest'
import { fireEvent, render, screen } from '@testing-library/react'
import { vi } from 'vitest'
import { createMockProviderContextValue } from '@/__mocks__/provider-context'
import { useProviderContext } from '@/context/provider-context'
import { Plan } from '../../../billing/type'
import { PlanBadge } from '../index'

const mockConfig = vi.hoisted(() => ({
  isCloudEdition: true,
}))

vi.mock('@/config', () => ({
  get IS_CLOUD_EDITION() {
    return mockConfig.isCloudEdition
  },
}))

vi.mock('@/context/provider-context', () => ({
  useProviderContext: vi.fn(),
  baseProviderContextValue: {},
}))

describe('PlanBadge', () => {
  const mockUseProviderContext = useProviderContext as Mock

  beforeEach(() => {
    vi.clearAllMocks()
    mockConfig.isCloudEdition = true
  })

  it('should return null if isFetchedPlan is false', () => {
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: false }),
    )
    const { container } = render(<PlanBadge plan={Plan.sandbox} />)
    expect(container.firstChild).toBeNull()
  })

  it('should render upgrade badge when plan is sandbox and sandboxAsUpgrade is true', () => {
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )
    render(<PlanBadge plan={Plan.sandbox} sandboxAsUpgrade={true} />)
    expect(
      screen.getByText('billing.upgradeBtn.encourageShort'),
    ).toBeInTheDocument()
    expect(screen.queryByRole('button')).not.toBeInTheDocument()
  })

  it('should render upgrade action as a button when onClick is provided', () => {
    const handleClick = vi.fn()
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )

    render(<PlanBadge plan={Plan.sandbox} sandboxAsUpgrade={true} onClick={handleClick} />)

    const button = screen.getByRole('button', { name: 'billing.upgradeBtn.encourageShort' })
    fireEvent.click(button)
    expect(handleClick).toHaveBeenCalledTimes(1)
  })

  it('should render sandbox badge instead of upgrade badge in self-hosted edition', () => {
    mockConfig.isCloudEdition = false
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )

    render(<PlanBadge plan={Plan.sandbox} sandboxAsUpgrade={true} />)

    expect(screen.getByText(Plan.sandbox)).toBeInTheDocument()
    expect(screen.queryByText('billing.upgradeBtn.encourageShort')).not.toBeInTheDocument()
    expect(screen.queryByRole('button')).not.toBeInTheDocument()
  })

  it('should render sandbox badge when plan is sandbox and sandboxAsUpgrade is false', () => {
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )
    render(<PlanBadge plan={Plan.sandbox} sandboxAsUpgrade={false} />)
    expect(screen.getByText(Plan.sandbox)).toBeInTheDocument()
    expect(screen.queryByRole('button')).not.toBeInTheDocument()
  })

  it('should render professional badge when plan is professional', () => {
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )
    render(<PlanBadge plan={Plan.professional} />)
    expect(screen.getByText('pro')).toBeInTheDocument()
  })

  it('should render graduation icon when isEducationWorkspace is true and plan is professional', () => {
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({
        isFetchedPlan: true,
        isEducationWorkspace: true,
      }),
    )
    const { container } = render(<PlanBadge plan={Plan.professional} />)

    expect(container.querySelector('svg')).toBeInTheDocument()
    expect(screen.getByText('pro')).toBeInTheDocument()
  })

  it('should render team badge when plan is team', () => {
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )
    render(<PlanBadge plan={Plan.team} />)
    expect(screen.getByText(Plan.team)).toBeInTheDocument()
  })

  it('should return null when plan is enterprise', () => {
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )
    const { container } = render(<PlanBadge plan={Plan.enterprise} />)
    expect(container.firstChild).toBeNull()
  })

  it('should trigger onClick when clicked', () => {
    const handleClick = vi.fn()
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )
    render(<PlanBadge plan={Plan.team} onClick={handleClick} />)
    fireEvent.click(screen.getByRole('button', { name: Plan.team }))
    expect(handleClick).toHaveBeenCalledTimes(1)
  })

  it('should handle allowHover prop', () => {
    mockUseProviderContext.mockReturnValue(
      createMockProviderContextValue({ isFetchedPlan: true }),
    )
    const { container } = render(
      <PlanBadge plan={Plan.team} allowHover={true} />,
    )

    expect(container.firstChild).not.toBeNull()
  })
})