Studio: Allow blocking user from adding demo sites

src/components/content-tab-snapshots.tsx

@@ -5,6 +5,7 @@ import { Icon, check, external } from '@wordpress/icons';
 import { useI18n } from '@wordpress/react-i18n';
 import { PropsWithChildren, useEffect } from 'react';
 import { CLIENT_ID, PROTOCOL_PREFIX, WP_AUTHORIZE_ENDPOINT, SCOPES } from '../constants';
+import { useArchiveErrorMessages } from '../hooks/use-archive-error-messages';
 import { useArchiveSite } from '../hooks/use-archive-site';
 import { useAuth } from '../hooks/use-auth';
 import { useExpirationDate } from '../hooks/use-expiration-date';
@@ -20,7 +21,7 @@ import { CopyTextButton } from './copy-text-button';
 import offlineIcon from './offline-icon';
 import ProgressBar from './progress-bar';
 import { ScreenshotDemoSite } from './screenshot-demo-site';
-import Tooltip from './tooltip';
+import Tooltip, { TooltipProps } from './tooltip';
 
 interface ContentTabSnapshotsProps {
 	selectedSite: SiteDetails;
@@ -52,10 +53,12 @@ function SnapshotRow( {
 	const { url, date, isDeleting } =
 		previousSnapshot && snapshot.isLoading ? previousSnapshot : snapshot;
 	const { countDown, isExpired, dateString } = useExpirationDate( date );
+	const { deleteSnapshot, fetchSnapshotUsage, snapshotCreationBlocked, removeSnapshot } =
+		useSnapshots();
 	const { isUploadingSiteId } = useArchiveSite();
 	const isUploading = isUploadingSiteId( selectedSite.id );
 	const { updateDemoSite, isDemoSiteUpdating } = useUpdateDemoSite();
-	const { removeSnapshot, deleteSnapshot } = useSnapshots();
+	const errorMessages = useArchiveErrorMessages();
 
 	const isOffline = useOffline();
 	const updateDemoSiteOfflineMessage = __(
@@ -64,6 +67,7 @@ function SnapshotRow( {
 	const deleteDemoSiteOfflineMessage = __(
 		'Deleting a demo site requires an internet connection.'
 	);
+	const userBlockedMessage = errorMessages.rest_site_creation_blocked;
 
 	const { progress, setProgress } = useProgressTimer( {
 		paused: ! isDemoSiteUpdating,
@@ -72,6 +76,10 @@ function SnapshotRow( {
 		maxValue: 95,
 	} );
 
+	useEffect( () => {
+		fetchSnapshotUsage();
+	}, [ fetchSnapshotUsage ] );
+
 	useEffect( () => {
 		if ( isDemoSiteUpdating ) {
 			setProgress( 80 );
@@ -160,6 +168,18 @@ function SnapshotRow( {
 			</div>
 		);
 	}
+
+	let tooltipContent: Partial< TooltipProps & { text?: string } > = {};
+	if ( isOffline ) {
+		tooltipContent = {
+			icon: offlineIcon,
+			text: updateDemoSiteOfflineMessage,
+		};
+	} else if ( snapshotCreationBlocked ) {
+		tooltipContent = { text: userBlockedMessage };
+	}
+	const isUpdateDisabled = isOffline || snapshotCreationBlocked;
+
 	return (
 		<div className="self-stretch flex-col px-4 py-3">
 			<div className="flex gap-2 items-center">
@@ -188,17 +208,13 @@ function SnapshotRow( {
 					</div>
 				) : (
 					<>
-						<Tooltip
-							disabled={ ! isOffline }
-							icon={ offlineIcon }
-							text={ updateDemoSiteOfflineMessage }
-						>
+						<Tooltip disabled={ ! isUpdateDisabled } { ...tooltipContent }>
 							<Button
-								aria-description={ isOffline ? updateDemoSiteOfflineMessage : '' }
-								aria-disabled={ isOffline }
+								aria-description={ tooltipContent?.text || '' }
+								aria-disabled={ isUpdateDisabled }
 								variant="primary"
 								onClick={ () => {
-									if ( isOffline ) {
+									if ( isUpdateDisabled ) {
 										return;
 									}
 									handleUpdateDemoSite();
@@ -388,7 +404,8 @@ function AddDemoSiteWithProgress( {
 	const { __, _n } = useI18n();
 	const { archiveSite, isUploadingSiteId, isAnySiteArchiving } = useArchiveSite();
 	const isUploading = isUploadingSiteId( selectedSite.id );
-	const { activeSnapshotCount, snapshotQuota, isLoadingSnapshotUsage } = useSnapshots();
+	const { activeSnapshotCount, snapshotQuota, isLoadingSnapshotUsage, snapshotCreationBlocked } =
+		useSnapshots();
 	const isLimitUsed = activeSnapshotCount >= snapshotQuota;
 	const isOffline = useOffline();
 	const { progress, setProgress } = useProgressTimer( {
@@ -397,14 +414,21 @@ function AddDemoSiteWithProgress( {
 		interval: 1500,
 		maxValue: 95,
 	} );
+	const errorMessages = useArchiveErrorMessages();
+
 	useEffect( () => {
 		if ( isSnapshotLoading ) {
 			setProgress( 80 );
 		}
 	}, [ isSnapshotLoading, setProgress ] );
 
 	const isDisabled =
-		isAnySiteArchiving || isUploading || isLoadingSnapshotUsage || isLimitUsed || isOffline;
+		isAnySiteArchiving ||
+		isUploading ||
+		isLoadingSnapshotUsage ||
+		isLimitUsed ||
+		isOffline ||
+		snapshotCreationBlocked;
 	const siteArchivingMessage = __(
 		'A different demo site is being created. Please wait for it to finish before creating another.'
 	);
@@ -417,6 +441,7 @@ function AddDemoSiteWithProgress( {
 		snapshotQuota
 	);
 	const offlineMessage = __( 'Creating a demo site requires an internet connection.' );
+	const userBlockedMessage = errorMessages.rest_site_creation_blocked;
 
 	let tooltipContent;
 	if ( isOffline ) {
@@ -428,6 +453,8 @@ function AddDemoSiteWithProgress( {
 		tooltipContent = { text: allotmentConsumptionMessage };
 	} else if ( isAnySiteArchiving ) {
 		tooltipContent = { text: siteArchivingMessage };
+	} else if ( snapshotCreationBlocked ) {
+		tooltipContent = { text: userBlockedMessage };
 	}
 
 	return (

src/components/tests/content-tab-snapshots.test.tsx

@@ -101,6 +101,7 @@ describe( 'ContentTabSnapshots', () => {
 			activeSnapshotCount: 1,
 			snapshotQuota: LIMIT_OF_ZIP_SITES_PER_USER,
 			isLoadingSnapshotUsage: false,
+			fetchSnapshotUsage: jest.fn(),
 		} );
 		( useOffline as jest.Mock ).mockReturnValue( false );
 
@@ -152,6 +153,7 @@ describe( 'ContentTabSnapshots', () => {
 					deleted: false,
 				},
 			],
+			fetchSnapshotUsage: jest.fn(),
 		} );
 		render( <ContentTabSnapshots selectedSite={ selectedSite } /> );
 		const createSnapshotButton = screen.getByRole( 'button', { name: 'Add demo site' } );
@@ -176,6 +178,7 @@ describe( 'ContentTabSnapshots', () => {
 					deleted: false,
 				},
 			],
+			fetchSnapshotUsage: jest.fn(),
 		} );
 		render( <ContentTabSnapshots selectedSite={ selectedSite } /> );
 		const createSnapshotButton = screen.getByRole( 'button', { name: 'Add demo site' } );
@@ -203,6 +206,7 @@ describe( 'ContentTabSnapshots', () => {
 					deleted: false,
 				},
 			],
+			fetchSnapshotUsage: jest.fn(),
 		} );
 		mockShowMessageBox.mockResolvedValueOnce( { response: 0, checkboxChecked: false } );
 		render( <ContentTabSnapshots selectedSite={ selectedSite } /> );
@@ -234,6 +238,7 @@ describe( 'ContentTabSnapshots', () => {
 					deleted: false,
 				},
 			],
+			fetchSnapshotUsage: jest.fn(),
 		} );
 		mockShowMessageBox.mockResolvedValueOnce( { response: 1 } );
 		render( <ContentTabSnapshots selectedSite={ selectedSite } /> );
@@ -274,6 +279,7 @@ describe( 'ContentTabSnapshots', () => {
 					deleted: false,
 				},
 			],
+			fetchSnapshotUsage: jest.fn(),
 			deleteSnapshot: mockDeleteSnapshot,
 		} );
 		mockShowMessageBox.mockResolvedValueOnce( { response: 0 } );
@@ -309,6 +315,7 @@ describe( 'ContentTabSnapshots', () => {
 					deleted: false,
 				},
 			],
+			fetchSnapshotUsage: jest.fn(),
 			deleteSnapshot: deleteSnapshotMock,
 		} );
 		mockShowMessageBox.mockResolvedValueOnce( { response: 1 } );
@@ -335,6 +342,7 @@ describe( 'ContentTabSnapshots', () => {
 					deleted: false,
 				},
 			],
+			fetchSnapshotUsage: jest.fn(),
 		} );
 		render( <ContentTabSnapshots selectedSite={ selectedSite } /> );
 
@@ -399,6 +407,7 @@ describe( 'ContentTabSnapshots', () => {
 		( useSnapshots as jest.Mock ).mockReturnValue( {
 			snapshots: [ snapshot ],
 			removeSnapshot,
+			fetchSnapshotUsage: jest.fn(),
 		} );
 
 		const { rerender } = render( <ContentTabSnapshots selectedSite={ selectedSite } /> );
@@ -469,6 +478,18 @@ describe( 'AddDemoSiteWithProgress', () => {
 
 	test( 'Progressbar is present when the second snapshot is being created', async () => {
 		isUploadingSiteId.mockReturnValue( true );
+		( useSnapshots as jest.Mock ).mockReturnValue( {
+			snapshots: [
+				{
+					url: 'fake-site.fake',
+					atomicSiteId: 150,
+					localSiteId: 'site-id-1',
+					date: 1707232820627,
+					deleted: false,
+				},
+			],
+			fetchSnapshotUsage: jest.fn(),
+		} );
 		render( <ContentTabSnapshots selectedSite={ { ...selectedSite, id: 'site-id-1' } } /> );
 		const addDemoSiteButton = screen.queryByRole( 'button', { name: 'Add demo site' } );
 		expect( addDemoSiteButton ).not.toBeInTheDocument();
@@ -477,6 +498,18 @@ describe( 'AddDemoSiteWithProgress', () => {
 
 	test( 'Button is enabled when no snapshots and no other site is being archived', async () => {
 		isUploadingSiteId.mockReturnValue( false );
+		( useSnapshots as jest.Mock ).mockReturnValue( {
+			snapshots: [
+				{
+					url: 'fake-site.fake',
+					atomicSiteId: 150,
+					localSiteId: 'site-id-1',
+					date: 1707232820627,
+					deleted: false,
+				},
+			],
+			fetchSnapshotUsage: jest.fn(),
+		} );
 		render( <ContentTabSnapshots selectedSite={ { ...selectedSite, id: 'site-id-1' } } /> );
 		const addDemoSiteButton = screen.getByRole( 'button', { name: 'Add demo site' } );
 		expect( addDemoSiteButton ).toBeEnabled();
@@ -495,6 +528,18 @@ describe( 'AddDemoSiteWithProgress', () => {
 
 	test( 'Button is disabled when offline', async () => {
 		( useOffline as jest.Mock ).mockReturnValue( true );
+		( useSnapshots as jest.Mock ).mockReturnValue( {
+			snapshots: [
+				{
+					url: 'fake-site.fake',
+					atomicSiteId: 150,
+					localSiteId: 'site-id-1',
+					date: 1707232820627,
+					deleted: false,
+				},
+			],
+			fetchSnapshotUsage: jest.fn(),
+		} );
 		render( <ContentTabSnapshots selectedSite={ { ...selectedSite, id: 'site-id-1' } } /> );
 		const addDemoSiteButton = screen.getByRole( 'button', { name: 'Add demo site' } );
 		expect( addDemoSiteButton ).toHaveAttribute( 'aria-disabled', 'true' );

src/components/tooltip.tsx

@@ -1,7 +1,7 @@
 import { Icon, Popover } from '@wordpress/components';
 import { PropsWithChildren, useState } from 'react';
 
-interface TooltipProps
+export interface TooltipProps
 	extends Pick< React.ComponentProps< typeof Popover >, 'placement' | 'className' > {
 	icon?: JSX.Element;
 	text?: string | JSX.Element;

src/hooks/use-archive-error-messages.ts

@@ -20,6 +20,7 @@ export function useArchiveErrorMessages() {
 				rest_cannot_view: __(
 					"There's been an authentication error. Please log in again before sharing a site."
 				),
+				rest_site_creation_blocked: __( 'Demo sites are not available for your account.' ),
 			} ) as const,
 		[ __ ]
 	);

src/hooks/use-archive-site.ts

@@ -12,7 +12,7 @@ import { SnapshotStatus, SnapshotStatusResponse, useSnapshots } from './use-snap
 
 export function useArchiveSite() {
 	const { uploadingSites, setUploadingSites } = useSiteDetails();
-	const { snapshots, addSnapshot, updateSnapshot } = useSnapshots();
+	const { snapshots, addSnapshot, updateSnapshot, fetchSnapshotUsage } = useSnapshots();
 	const isUploadingSiteId = useCallback(
 		( localSiteId: string ) => uploadingSites[ localSiteId ] || false,
 		[ uploadingSites ]
@@ -70,6 +70,13 @@ export function useArchiveSite() {
 			}
 
 			setUploadingSites( ( _uploadingSites ) => ( { ..._uploadingSites, [ siteId ]: true } ) );
+
+			const snapshotUsage = await fetchSnapshotUsage();
+			if ( snapshotUsage?.site_creation_blocked ) {
+				alert( errorMessages.rest_site_creation_blocked );
+				setUploadingSites( ( _uploadingSites ) => ( { ..._uploadingSites, [ siteId ]: false } ) );
+				return;
+			}
 			const { zipContent, zipPath, exceedsSizeLimit } = await getIpcApi().archiveSite( siteId );
 			if ( exceedsSizeLimit ) {
 				alert(
@@ -133,7 +140,7 @@ export function useArchiveSite() {
 				getIpcApi().removeTemporalFile( zipPath );
 			}
 		},
-		[ __, addSnapshot, client, errorMessages, setUploadingSites ]
+		[ __, addSnapshot, client, errorMessages, fetchSnapshotUsage, setUploadingSites ]
 	);
 	const isAnySiteArchiving = useMemo( () => {
 		const isAnySiteUploading = Object.values( uploadingSites ).some( ( uploading ) => uploading );