Refactor variant categorization in StructurePreview to utilize categorizeVariants function for improved clarity and maintainability

packages/gitbook/src/components/SpaceLayout/categorizeVariants.ts

@@ -2,10 +2,15 @@ import { languages } from '@/intl/translations';
 import type { GitBookSiteContext } from '@/lib/context';
 import { getSiteSpaceLanguages, normalizeLanguage } from '@/lib/sites';
 
+type SiteSpaceVariantsContext = Pick<
+    GitBookSiteContext,
+    'locale' | 'siteSpace' | 'siteSpaces' | 'visibleSiteSpaces'
+>;
+
 /**
  * Categorize the variants of the space into generic and translation variants.
  */
-export function categorizeVariants(context: GitBookSiteContext) {
+export function categorizeVariants(context: SiteSpaceVariantsContext) {
     const { siteSpace } = context;
 
     // By default, variants only include visible spaces.

packages/gitbook/src/components/StructurePreview/StructurePreview.tsx

@@ -121,6 +121,9 @@ function StructurePreviewHeader(props: { snapshot: StructurePreviewSnapshot }) {
             (sections.list.length > 1 ||
                 sections.list.some((section) => section.object === 'site-section-group'))
     );
+    const translationSiteSpace =
+        variants.translations.find((space) => space.id === snapshot.siteSpace.id) ??
+        snapshot.siteSpace;
 
     return (
         <HeaderLayout
@@ -186,7 +189,7 @@ function StructurePreviewHeader(props: { snapshot: StructurePreviewSnapshot }) {
                         {!withSections && variants.translations.length > 1 ? (
                             <StructurePreviewTranslationsDropdown
                                 snapshot={snapshot}
-                                siteSpace={snapshot.siteSpace}
+                                siteSpace={translationSiteSpace}
                                 siteSpaces={variants.translations}
                                 className="flex! site-header:theme-bold:text-header-link hover:site-header:theme-bold:bg-header-link/3 focus-visible:site-header:theme-bold:bg-header-link/3 aria-expanded:site-header:theme-bold:bg-header-link/5"
                             />
@@ -201,7 +204,7 @@ function StructurePreviewHeader(props: { snapshot: StructurePreviewSnapshot }) {
                         {variants.translations.length > 1 ? (
                             <StructurePreviewTranslationsDropdown
                                 snapshot={snapshot}
-                                siteSpace={snapshot.siteSpace}
+                                siteSpace={translationSiteSpace}
                                 siteSpaces={variants.translations}
                                 className="my-1.5 ml-2 self-start"
                             />

packages/gitbook/src/components/StructurePreview/state.test.ts

@@ -1,8 +1,10 @@
 import { describe, expect, it } from 'bun:test';
 
+import { languages } from '@/intl/translations';
 import { defaultCustomization } from '@/lib/utils';
+import { TranslationLanguage } from '@gitbook/api';
 
-import { encodePreviewSiteSections, isStructurePreviewMessage } from './state';
+import { encodePreviewSiteSections, getPreviewVariants, isStructurePreviewMessage } from './state';
 import type { StructurePreviewSnapshot } from './types';
 
 function createSnapshot(
@@ -98,4 +100,61 @@ describe('structure preview state', () => {
         expect(encoded?.current.url).toBe('#');
         expect(encoded?.list[0]?.object).toBe('site-section-group');
     });
+
+    it('categorizes translation variants like the rendered header', () => {
+        const currentSiteSpace = {
+            id: 'v15-it',
+            title: 'v15',
+            path: '',
+            default: false,
+            hidden: false,
+            urls: {},
+            space: {
+                id: 'space-v15-it',
+                revision: 'revision-v15-it',
+                language: TranslationLanguage.It,
+            },
+        };
+        const siteSpaces = [
+            { id: 'v20-en', title: 'v20', language: TranslationLanguage.En },
+            { id: 'v20-fr', title: 'v20', language: TranslationLanguage.Fr },
+            { id: 'v20-it', title: 'v20', language: TranslationLanguage.It },
+            { id: 'v15-en', title: 'v15', language: TranslationLanguage.En },
+            { id: 'v15-fr', title: 'v15', language: TranslationLanguage.Fr },
+            currentSiteSpace,
+        ].map((siteSpace) =>
+            'space' in siteSpace
+                ? siteSpace
+                : {
+                      id: siteSpace.id,
+                      title: siteSpace.title,
+                      path: '',
+                      default: false,
+                      hidden: false,
+                      urls: {},
+                      space: {
+                          id: `space-${siteSpace.id}`,
+                          revision: `revision-${siteSpace.id}`,
+                          language: siteSpace.language,
+                      },
+                  }
+        );
+        const snapshot = createSnapshot({
+            locale: TranslationLanguage.It,
+            siteSpace: currentSiteSpace,
+            siteSpaces,
+            visibleSiteSpaces: siteSpaces,
+        } as unknown as Partial<StructurePreviewSnapshot>);
+
+        const variants = getPreviewVariants(snapshot);
+
+        expect(variants.generic.map((space) => space.id)).toEqual(['v20-it', 'v15-it']);
+        expect(
+            variants.translations.map((space) => ({ id: space.id, title: space.title }))
+        ).toEqual([
+            { id: 'v15-en', title: languages.en.language },
+            { id: 'v15-fr', title: languages.fr.language },
+            { id: 'v15-it', title: languages.it.language },
+        ]);
+    });
 });

packages/gitbook/src/components/StructurePreview/state.ts

@@ -13,6 +13,7 @@ import type {
     ClientSiteSectionGroup,
     ClientSiteSections,
 } from '@/components/SiteSections/encodeClientSiteSections';
+import { categorizeVariants } from '@/components/SpaceLayout/categorizeVariants';
 import { getLocalizedDescription, getLocalizedTitle } from '@/lib/sites';
 
 import type { StructurePreviewMessage, StructurePreviewSnapshot } from './types';
@@ -108,28 +109,7 @@ function encodePreviewSection(
 }
 
 export function getPreviewVariants(snapshot: StructurePreviewSnapshot) {
-    const siteSpaces =
-        snapshot.visibleSiteSpaces.length > 0 ? snapshot.visibleSiteSpaces : snapshot.siteSpaces;
-    const currentLanguage = snapshot.locale ?? snapshot.siteSpace.space.language ?? 'en';
-    const languages = [...new Set(siteSpaces.map((space) => space.space.language ?? 'en'))];
-    const isMultiLanguage = languages.length > 1;
-
-    return {
-        generic: isMultiLanguage
-            ? siteSpaces.filter(
-                  (space) =>
-                      space.id === snapshot.siteSpace.id ||
-                      (space.space.language ?? 'en') === currentLanguage
-              )
-            : siteSpaces,
-        translations: isMultiLanguage
-            ? siteSpaces.filter(
-                  (space) =>
-                      space.id === snapshot.siteSpace.id ||
-                      (space.space.language ?? 'en') !== currentLanguage
-              )
-            : [],
-    };
+    return categorizeVariants(snapshot);
 }
 
 export function getHeaderSocialAccounts(customization: SiteCustomizationSettings) {