diff --git a/frontend/src/pages/Upload.test.tsx b/frontend/src/pages/Upload.test.tsx new file mode 100644 index 00000000..c9c52fec --- /dev/null +++ b/frontend/src/pages/Upload.test.tsx @@ -0,0 +1,75 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot } from "react-dom/client"; +import { MemoryRouter } from "react-router-dom"; +import { expect, test, vi } from "vitest"; + +import Upload from "./Upload"; + +const { fetchWithHeaders, toast } = vi.hoisted(() => ({ + fetchWithHeaders: vi.fn(), + toast: vi.fn(), +})); + +vi.mock("@/contexts/ApiContext", () => ({ + useApi: () => ({ baseUrl: "http://api", fetchWithHeaders }), +})); +vi.mock("@/hooks/use-toast", () => ({ useToast: () => ({ toast }) })); + +test.each([false, true])("opens the uploaded Hub dataset (private: %s)", async (isPrivate) => { + Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + const datasetInfo = { dataset_repo_id: "local/sushi2", single_task: "pick", num_episodes: 2 }; + let completeUpload!: (response: Response) => void; + const pendingUpload = new Promise((resolve) => { + completeUpload = resolve; + }); + fetchWithHeaders.mockReset().mockImplementation(async (url: string) => { + if (url.endsWith("/dataset-info")) { + return Response.json({ ...datasetInfo, success: true }); + } + if (url.endsWith("/upload-dataset")) { + return pendingUpload; + } + throw new Error(`Unexpected request: ${url}`); + }); + const open = vi.spyOn(window, "open").mockReturnValue(null); + const host = document.createElement("div"); + document.body.append(host); + const root = createRoot(host); + const button = (text: string) => + [...host.querySelectorAll("button")].find((item) => item.textContent?.trim() === text); + + try { + await act(async () => { + root.render( + + + , + ); + }); + if (isPrivate) { + await act(async () => host.querySelector("#private")?.click()); + } + const upload = button("Upload to HuggingFace Hub"); + expect(upload).toBeDefined(); + await act(async () => upload?.click()); + const [, options] = fetchWithHeaders.mock.calls.find(([url]) => url.endsWith("/upload-dataset"))!; + expect(JSON.parse(options.body)).toMatchObject({ dataset_repo_id: "local/sushi2", private: isPrivate }); + await act(async () => { + host.querySelector("#private")?.click(); + completeUpload(Response.json({ success: true, hub_repo_id: "alice/sushi2" })); + }); + + const view = button("View on HuggingFace Hub"); + expect(view).toBeDefined(); + await act(async () => view?.click()); + const target = new URL(String(open.mock.calls[0]?.[0])); + expect(target.pathname).toBe(isPrivate ? "/login" : "/spaces/lerobot/visualize_dataset"); + const viewer = isPrivate ? new URL(target.searchParams.get("next")!, target.origin) : target; + expect(viewer.searchParams.get("path")).toBe("/alice/sushi2"); + } finally { + act(() => root.unmount()); + host.remove(); + open.mockRestore(); + } +}); diff --git a/frontend/src/pages/Upload.tsx b/frontend/src/pages/Upload.tsx index d74e5734..96b34295 100644 --- a/frontend/src/pages/Upload.tsx +++ b/frontend/src/pages/Upload.tsx @@ -60,8 +60,9 @@ const Upload = () => { const [tagsInput, setTagsInput] = useState(uploadConfig.tags.join(", ")); const [isUploading, setIsUploading] = useState(false); - const [uploadSuccess, setUploadSuccess] = useState(false); + const [uploadedDataset, setUploadedDataset] = useState<{ repoId: string; private: boolean } | null>(null); const [showDeleteConfirm, setShowDeleteConfirm] = useState(false); + const uploadSuccess = uploadedDataset !== null; // Load actual dataset information from backend React.useEffect(() => { @@ -122,12 +123,12 @@ const Upload = () => { loadDatasetInfo(); }, [initialDatasetInfo, navigate, toast]); - const openInHubViewer = (repoId: string) => { + const openInHubViewer = (repoId: string, requiresLogin = true) => { const spacePath = `/spaces/lerobot/visualize_dataset?path=${encodeURIComponent(`/${repoId}`)}`; - // The user owns/manages the dataset (it appears under their hub - // listing), so login-redirect always works whether public or - // private. Avoids passing `private` through navigation state. - const target = `https://huggingface.co/login?next=${encodeURIComponent(spacePath)}`; + // Use the login redirect when the dataset is private or its visibility is unknown. + const target = requiresLogin + ? `https://huggingface.co/login?next=${encodeURIComponent(spacePath)}` + : `https://huggingface.co${spacePath}`; window.open(target, "_blank", "noopener,noreferrer"); }; @@ -148,6 +149,7 @@ const Upload = () => { const handleUploadToHub = async () => { if (!datasetInfo) return; + const isPrivate = uploadConfig.private; setIsUploading(true); try { // Parse tags from input @@ -161,17 +163,18 @@ const Upload = () => { body: JSON.stringify({ dataset_repo_id: datasetInfo.dataset_repo_id, tags, - private: uploadConfig.private, + private: isPrivate, }), }); const data = await response.json(); if (response.ok && data.success) { - setUploadSuccess(true); + const hubRepoId = data.hub_repo_id ?? datasetInfo.dataset_repo_id; + setUploadedDataset({ repoId: hubRepoId, private: isPrivate }); toast({ title: "Upload Successful!", - description: `Dataset ${datasetInfo.dataset_repo_id} has been uploaded to HuggingFace Hub.`, + description: `Dataset ${hubRepoId} has been uploaded to HuggingFace Hub.`, }); } else { const fallback = "Failed to upload dataset to HuggingFace Hub."; @@ -267,7 +270,7 @@ const Upload = () => { {/* Success State */} - {uploadSuccess && ( + {uploadedDataset !== null && (
@@ -281,15 +284,7 @@ const Upload = () => {