refactor(web): migrate vector space billing query to contract client (#39558)

This commit is contained in:
Evan
2026-07-25 04:26:14 +00:00
committed by GitHub
parent fa000ae8a7
commit c56ac3983e
9 changed files with 76 additions and 68 deletions
@@ -16,18 +16,30 @@ import TriggerEventsLimitModal from '@/app/components/billing/trigger-events-lim
import { Plan } from '@/app/components/billing/type'
import UpgradeBtn from '@/app/components/billing/upgrade-btn'
import VectorSpaceFull from '@/app/components/billing/vector-space-full'
import { createConsoleQueryWrapper } from '@/test/console/query-data'
import { consoleQuery } from '@/service/client'
import { createConsoleQueryClient, createConsoleQueryWrapper } from '@/test/console/query-data'
import { render as renderWithConsoleState } from '@/test/console/render'
let mockProviderCtx: Record<string, unknown> = {}
let mockConsoleState: Record<string, unknown> = {}
const render = (ui: ReactElement, options: RenderOptions = {}) => {
const queryClient = createConsoleQueryClient()
const plan = mockProviderCtx.plan as {
usage: { vectorSpace: number }
total: { vectorSpace: number }
}
queryClient.setQueryData(consoleQuery.features.vectorSpace.get.queryOptions().queryKey, {
size: plan.usage.vectorSpace,
limit: plan.total.vectorSpace,
})
const { wrapper } = createConsoleQueryWrapper({
systemFeatures: { deployment_edition: 'CLOUD' },
queryClient,
})
return renderWithConsoleState(ui, { ...options, wrapper })
}
let mockProviderCtx: Record<string, unknown> = {}
let mockConsoleState: Record<string, unknown> = {}
const mockSetShowPricingModal = vi.fn()
const mockSetShowAccountSettingModal = vi.fn()
@@ -72,9 +84,6 @@ vi.mock('@/service/use-billing', () => ({
refetch: mockRefetch,
}),
useBindPartnerStackInfo: () => ({ mutateAsync: vi.fn() }),
useCurrentPlanVectorSpace: () => ({
data: undefined,
}),
}))
vi.mock('@/service/use-education', () => ({
@@ -16,19 +16,31 @@ import * as React from 'react'
import { defaultPlan } from '@/app/components/billing/config'
import PlanComp from '@/app/components/billing/plan'
import { Plan } from '@/app/components/billing/type'
import { createConsoleQueryWrapper } from '@/test/console/query-data'
import { consoleQuery } from '@/service/client'
import { createConsoleQueryClient, createConsoleQueryWrapper } from '@/test/console/query-data'
import { render as renderWithConsoleState } from '@/test/console/render'
let mockProviderCtx: Record<string, unknown> = {}
let mockConsoleState: Record<string, unknown> = {}
const render = (ui: ReactElement, options: RenderOptions = {}) => {
const queryClient = createConsoleQueryClient()
const plan = mockProviderCtx.plan as {
usage: { vectorSpace: number }
total: { vectorSpace: number }
}
queryClient.setQueryData(consoleQuery.features.vectorSpace.get.queryOptions().queryKey, {
size: plan.usage.vectorSpace,
limit: plan.total.vectorSpace,
})
const { wrapper } = createConsoleQueryWrapper({
systemFeatures: { deployment_edition: 'CLOUD' },
queryClient,
})
return renderWithConsoleState(ui, { ...options, wrapper })
}
// ─── Mock state ──────────────────────────────────────────────────────────────
let mockProviderCtx: Record<string, unknown> = {}
let mockConsoleState: Record<string, unknown> = {}
const mockSetShowPricingModal = vi.fn()
const mockSetShowAccountSettingModal = vi.fn()
const mockRouterPush = vi.fn()
@@ -77,9 +89,6 @@ vi.mock('@/service/use-billing', () => ({
isFetching: false,
refetch: vi.fn(),
}),
useCurrentPlanVectorSpace: () => ({
data: undefined,
}),
}))
// ─── Navigation mocks ───────────────────────────────────────────────────────
@@ -2,10 +2,11 @@
import type { FC } from 'react'
import type { BasicPlan } from '../type'
import { RiHardDrive3Line } from '@remixicon/react'
import { useQuery } from '@tanstack/react-query'
import * as React from 'react'
import { useTranslation } from 'react-i18next'
import { useProviderContext } from '@/context/provider-context'
import { useCurrentPlanVectorSpace } from '@/service/use-billing'
import { consoleQuery } from '@/service/client'
import { Plan } from '../type'
import UsageInfo from '../usage-info'
import { getPlanVectorSpaceLimitMB } from '../utils'
@@ -20,7 +21,7 @@ const STORAGE_THRESHOLD_MB = getPlanVectorSpaceLimitMB(Plan.sandbox)
const VectorSpaceInfo: FC<Props> = ({ className }) => {
const { t } = useTranslation()
const { plan } = useProviderContext()
const { data: vectorSpace } = useCurrentPlanVectorSpace()
const { data: vectorSpace } = useQuery(consoleQuery.features.vectorSpace.get.queryOptions())
const displayPlan = vectorSpace
? {
...plan,
@@ -4,11 +4,27 @@ import type { CrawlOptions, CrawlResultItem, DataSet, FileItem } from '@/models/
import { fireEvent, screen } from '@testing-library/react'
import { Plan } from '@/app/components/billing/type'
import { DataSourceType } from '@/models/datasets'
import { renderWithConsoleQuery } from '@/test/console/query-data'
import { consoleQuery } from '@/service/client'
import { createConsoleQueryClient, renderWithConsoleQuery } from '@/test/console/query-data'
import StepOne from '../index'
const render = (ui: React.ReactElement) =>
renderWithConsoleQuery(ui, { systemFeatures: { deployment_edition: 'CLOUD' } })
let mockPlan = {
type: Plan.professional,
usage: { vectorSpace: 50, buildApps: 0, documentsUploadQuota: 0, vectorStorageQuota: 0 },
total: { vectorSpace: 100, buildApps: 0, documentsUploadQuota: 0, vectorStorageQuota: 0 },
}
const render = (ui: React.ReactElement) => {
const queryClient = createConsoleQueryClient()
queryClient.setQueryData(consoleQuery.features.vectorSpace.get.queryOptions().queryKey, {
size: mockPlan.usage.vectorSpace,
limit: mockPlan.total.vectorSpace,
})
return renderWithConsoleQuery(ui, {
systemFeatures: { deployment_edition: 'CLOUD' },
queryClient,
})
}
// Mock config for website crawl features
vi.mock('@/config', async (importOriginal) => ({
@@ -29,11 +45,6 @@ vi.mock('@/context/dataset-detail', () => ({
}))
// Mock provider context
let mockPlan = {
type: Plan.professional,
usage: { vectorSpace: 50, buildApps: 0, documentsUploadQuota: 0, vectorStorageQuota: 0 },
total: { vectorSpace: 100, buildApps: 0, documentsUploadQuota: 0, vectorStorageQuota: 0 },
}
let mockEnableBilling = false
vi.mock('@/context/provider-context', () => ({
@@ -43,16 +54,6 @@ vi.mock('@/context/provider-context', () => ({
}),
}))
vi.mock('@/service/use-billing', () => ({
useCurrentPlanVectorSpace: () => ({
data: {
size: mockPlan.usage.vectorSpace,
limit: mockPlan.total.vectorSpace,
},
isFetching: false,
}),
}))
vi.mock('../../file-uploader', () => ({
default: ({ onPreview, fileList }: { onPreview: (file: File) => void; fileList: FileItem[] }) => (
<div data-testid="file-uploader">
@@ -5,6 +5,7 @@ import type { DataSourceProvider, NotionPage } from '@/models/common'
import type { CrawlOptions, CrawlResultItem, FileItem } from '@/models/datasets'
import { cn } from '@langgenius/dify-ui/cn'
import { RiFolder6Line } from '@remixicon/react'
import { useQuery } from '@tanstack/react-query'
import { useBoolean } from 'ahooks'
import { useCallback, useMemo } from 'react'
import { useTranslation } from 'react-i18next'
@@ -15,7 +16,7 @@ import VectorSpaceFull from '@/app/components/billing/vector-space-full'
import { useDatasetDetailContextWithSelector } from '@/context/dataset-detail'
import { useProviderContext } from '@/context/provider-context'
import { DataSourceType } from '@/models/datasets'
import { useCurrentPlanVectorSpace } from '@/service/use-billing'
import { consoleQuery } from '@/service/client'
import EmptyDatasetCreationModal from '../empty-dataset-creation-modal'
import FileUploader from '../file-uploader'
import Website from '../website'
@@ -134,8 +135,9 @@ const StepOne = ({
const allFileLoaded = files.length > 0 && files.every((file) => file.file.id)
const hasNotion = notionPages.length > 0
const shouldCheckVectorSpace = enableBilling && (allFileLoaded || hasNotion)
const { data: vectorSpace, isFetching: isFetchingVectorSpacePlan } =
useCurrentPlanVectorSpace(shouldCheckVectorSpace)
const { data: vectorSpace, isFetching: isFetchingVectorSpacePlan } = useQuery(
consoleQuery.features.vectorSpace.get.queryOptions({ enabled: shouldCheckVectorSpace }),
)
const isCheckingVectorSpace = shouldCheckVectorSpace && !vectorSpace && isFetchingVectorSpacePlan
const isVectorSpaceFull =
!!vectorSpace && vectorSpace.limit > 0 && vectorSpace.size >= vectorSpace.limit
@@ -1,15 +1,26 @@
import { waitFor } from '@testing-library/react'
import { render } from '@/test/console/render'
import { consoleQuery } from '@/service/client'
import { createConsoleQueryClient, renderWithConsoleQuery } from '@/test/console/query-data'
import CreateFromPipeline from '../index'
let mockDatasetPermissionKeys = ['dataset.acl.use']
const mockRouterReplace = vi.fn()
const mockPlan = {
usage: { vectorSpace: 50 },
total: { vectorSpace: 100 },
type: 'professional',
}
const render = (ui: React.ReactElement) => {
const queryClient = createConsoleQueryClient()
queryClient.setQueryData(consoleQuery.features.vectorSpace.get.queryOptions().queryKey, {
size: mockPlan.usage.vectorSpace,
limit: mockPlan.total.vectorSpace,
})
return renderWithConsoleQuery(ui, { queryClient })
}
let mockDatasetPermissionKeys = ['dataset.acl.use']
const mockRouterReplace = vi.fn()
vi.mock('@/context/provider-context', () => ({
useProviderContextSelector: (
selector: (state: { plan: typeof mockPlan; enableBilling: boolean }) => unknown,
@@ -83,13 +94,6 @@ vi.mock('@/next/navigation', () => ({
}),
}))
vi.mock('@/service/use-billing', () => ({
useCurrentPlanVectorSpace: () => ({
data: { size: 50, limit: 100 },
isFetching: false,
}),
}))
vi.mock('@/service/use-common', () => ({
useFileUploadConfig: () => ({
data: { file_size_limit: 15, batch_count_limit: 5 },
@@ -4,6 +4,7 @@ import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-so
import type { Node } from '@/app/components/workflow/types'
import type { FileIndexingEstimateResponse } from '@/models/datasets'
import type { InitialDocumentDetail } from '@/models/pipeline'
import { useQuery } from '@tanstack/react-query'
import { useBoolean } from 'ahooks'
import { useAtomValue } from 'jotai'
import { useCallback, useEffect, useMemo, useState } from 'react'
@@ -19,7 +20,7 @@ import {
import { useProviderContextSelector } from '@/context/provider-context'
import { DatasourceType } from '@/models/pipeline'
import { useRouter } from '@/next/navigation'
import { useCurrentPlanVectorSpace } from '@/service/use-billing'
import { consoleQuery } from '@/service/client'
import { useFileUploadConfig } from '@/service/use-common'
import { usePublishedPipelineInfo } from '@/service/use-pipeline'
import { getDatasetACLCapabilities } from '@/utils/permission'
@@ -117,8 +118,9 @@ const CreateFormPipeline = () => {
onlineDocuments.length > 0 ||
websitePages.length > 0 ||
selectedFileIds.length > 0)
const { data: vectorSpace, isFetching: isFetchingVectorSpacePlan } =
useCurrentPlanVectorSpace(shouldCheckVectorSpace)
const { data: vectorSpace, isFetching: isFetchingVectorSpacePlan } = useQuery(
consoleQuery.features.vectorSpace.get.queryOptions({ enabled: shouldCheckVectorSpace }),
)
const isCheckingVectorSpace = shouldCheckVectorSpace && !vectorSpace && isFetchingVectorSpacePlan
const isVectorSpaceFull =
!!vectorSpace && vectorSpace.limit > 0 && vectorSpace.size >= vectorSpace.limit
-9
View File
@@ -1,19 +1,10 @@
import type { CurrentPlanInfoBackend, SubscriptionUrlsBackend } from '@/app/components/billing/type'
import { get } from './base'
export type CurrentPlanVectorSpaceBackend = {
size: number
limit: number
}
export const fetchCurrentPlanInfo = () => {
return get<CurrentPlanInfoBackend>('/features')
}
export const fetchCurrentPlanVectorSpace = () => {
return get<CurrentPlanVectorSpaceBackend>('/features/vector-space')
}
export const fetchSubscriptionUrls = (plan: string, interval: string) => {
return get<SubscriptionUrlsBackend>(`/billing/subscription?plan=${plan}&interval=${interval}`)
}
-11
View File
@@ -1,8 +1,5 @@
import { useMutation, useQuery } from '@tanstack/react-query'
import { consoleClient, consoleQuery } from '@/service/client'
import { fetchCurrentPlanVectorSpace } from './billing'
const currentPlanVectorSpaceQueryKey = ['billing', 'current-plan-vector-space'] as const
export const useBindPartnerStackInfo = () => {
return useMutation({
@@ -25,11 +22,3 @@ export const useBillingUrl = (enabled: boolean) => {
},
})
}
export const useCurrentPlanVectorSpace = (enabled = true) => {
return useQuery({
queryKey: currentPlanVectorSpaceQueryKey,
queryFn: () => fetchCurrentPlanVectorSpace(),
enabled,
})
}