Return 404 for calculator paths whose first segment is not a country

app/src/docs/ARCHITECTURE-COUNTRY-ID.md

@@ -16,12 +16,13 @@ Every country-specific calculator route sits under the country segment, for exam
 
 **Location**: `calculator-app/src/app/[countryId]/layout.tsx`
 
-The entry point for every country route. It:
+The entry point for every country route. It is a server component that checks `countryId` with `isCountryId` (`src/libs/countries.ts`) and calls `notFound()` when it is not a supported country. Any other first path segment, such as `/llms.txt` or `/US`, is therefore an HTTP 404 rendered by `calculator-app/src/app/not-found.tsx`, not a 200 calculator page.
 
-1. Checks `countryId` against `countryIds`, and replaces the route with `/` when it is not supported.
-2. Provides the country through `CountryProvider` (`src/contexts/CountryContext.tsx`).
-3. Provides `NavigationProvider` and `LocationProvider`, which expose Next.js navigation and the current location to shared components through `NavigationContext` and `LocationContext`.
-4. Renders the calculator's provider tree, `CalculatorProviders` (`calculator-app/src/app/[countryId]/providers.tsx`).
+For a supported country it renders `CountryShell` (`calculator-app/src/app/[countryId]/CountryShell.tsx`), a client component that:
+
+1. Provides the country through `CountryProvider` (`src/contexts/CountryContext.tsx`).
+2. Provides `NavigationProvider` and `LocationProvider`, which expose Next.js navigation and the current location to shared components through `NavigationContext` and `LocationContext`.
+3. Renders the calculator's provider tree, `CalculatorProviders` (`calculator-app/src/app/[countryId]/providers.tsx`).
 
 ## useCurrentCountry
 
@@ -54,7 +55,9 @@ When adding a country-dependent feature:
 
 ## Related files
 
-- `calculator-app/src/app/[countryId]/layout.tsx`: validation and context providers
+- `calculator-app/src/app/[countryId]/layout.tsx`: validation (404 for unsupported countries)
+- `calculator-app/src/app/[countryId]/CountryShell.tsx`: context providers
+- `calculator-app/src/app/not-found.tsx`: the 404 page
 - `calculator-app/src/app/[countryId]/providers.tsx`: provider tree and `MetadataBootstrap`
 - `src/contexts/CountryContext.tsx`: `CountryProvider` and `useCountryContext`
 - `src/hooks/useCurrentCountry.ts`: hook that reads the country

app/src/libs/countries.ts

@@ -2,6 +2,14 @@ export const countryIds = ['us', 'uk', 'ca', 'ng', 'il'] as const;
 
 export type CountryId = (typeof countryIds)[number];
 
+/**
+ * Whether a value, such as the first segment of a URL path, is a supported
+ * country ID. Exact match only: `'US'` and `' us'` are not countries.
+ */
+export function isCountryId(value: unknown): value is CountryId {
+  return (countryIds as readonly unknown[]).includes(value);
+}
+
 /**
  * Default country used as fallback when country hasn't loaded from URL yet.
  * This is necessary for initial state before useCurrentCountry() hook provides the actual country.

app/src/tests/unit/calculator/CountryLayout.test.tsx

@@ -1,11 +1,24 @@
+import { notFound } from 'next/navigation';
+import * as fs from 'fs';
+import * as path from 'path';
+// Not @test-utils: its render wraps a CountryProvider('us') of its own, which
+// would hide a layout that failed to provide the country.
 import { render, screen } from '@testing-library/react';
+import fc from 'fast-check';
 import { beforeEach, describe, expect, test, vi } from 'vitest';
 import { useCurrentCountry } from '@/hooks/useCurrentCountry';
+import { countryIds } from '@/libs/countries';
 import CountryLayout from '../../../../../calculator-app/src/app/[countryId]/layout';
 
-const replace = vi.fn();
+// Next.js's notFound() throws to stop rendering; the mock throws a marker so
+// a 404 can be told apart from any other failure.
+const { NOT_FOUND } = vi.hoisted(() => ({ NOT_FOUND: new Error('notFound()') }));
+
 vi.mock('next/navigation', () => ({
-  useRouter: () => ({ push: vi.fn(), replace, back: vi.fn(), prefetch: vi.fn() }),
+  notFound: vi.fn(() => {
+    throw NOT_FOUND;
+  }),
+  useRouter: () => ({ push: vi.fn(), replace: vi.fn(), back: vi.fn(), prefetch: vi.fn() }),
   usePathname: () => '/',
   useSearchParams: () => new URLSearchParams(),
 }));
@@ -18,40 +31,109 @@ vi.mock('../../../../../calculator-app/src/app/[countryId]/providers', () => ({
   ),
 }));
 
-/** Route params as Next.js passes them, already resolved so `use()` reads them synchronously. */
-function resolvedParams(countryId: string) {
-  const value = { countryId };
-  return Object.assign(Promise.resolve(value), { status: 'fulfilled', value });
-}
-
 function CountryProbe() {
   return <p>Country: {useCurrentCountry()}</p>;
 }
 
+/** The layout as Next.js calls it for a request whose first path segment is `countryId`. */
+function renderLayout(countryId: string) {
+  return CountryLayout({ params: Promise.resolve({ countryId }), children: <CountryProbe /> });
+}
+
+/** 'rendered' or 'not found'; any other failure is rethrown. */
+async function outcome(countryId: string) {
+  try {
+    await renderLayout(countryId);
+    return 'rendered';
+  } catch (error) {
+    if (error === NOT_FOUND) {
+      return 'not found';
+    }
+    throw error;
+  }
+}
+
 describe('calculator country layout', () => {
   beforeEach(() => {
-    replace.mockClear();
+    vi.mocked(notFound).mockClear();
   });
 
-  test('given a supported country then the page renders with that country', () => {
-    render(
-      <CountryLayout params={resolvedParams('uk')}>
-        <CountryProbe />
-      </CountryLayout>
+  test.each(countryIds)(
+    'given the supported country %s then the calculator renders with that country',
+    async (countryId) => {
+      render(await renderLayout(countryId));
+
+      expect(screen.getByTestId('calculator-providers')).toContainElement(
+        screen.getByText(`Country: ${countryId}`)
+      );
+      expect(notFound).not.toHaveBeenCalled();
+    }
+  );
+
+  test.each([
+    // Files crawlers and browsers ask the calculator host for.
+    'llms.txt',
+    'llms-full.txt',
+    'favicon.ico',
+    'sitemap.xml',
+    'wp-login.php',
+    // Countries the calculator does not serve, or near misses of ones it does.
+    'xx',
+    'gb',
+    'usa',
+    'US',
+    'Uk',
+    ' us',
+    'us ',
+    '',
+  ])('given the first path segment %j then the layout is a 404', async (segment) => {
+    await expect(renderLayout(segment)).rejects.toBe(NOT_FOUND);
+    expect(notFound).toHaveBeenCalledTimes(1);
+  });
+
+  test('given any first path segment then the layout is a 404 exactly when it is not a supported country', async () => {
+    const country = fc.constantFrom(...countryIds);
+    const segment = fc.oneof(
+      country,
+      fc.string(),
+      fc.string({ unit: 'binary' }),
+      // A supported country with something around it, or in another case.
+      fc
+        .tuple(fc.string(), country, fc.string())
+        .map(([before, id, after]) => `${before}${id}${after}`),
+      country.map((id) => id.toUpperCase())
     );
-    expect(screen.getByTestId('calculator-providers')).toContainElement(
-      screen.getByText('Country: uk')
+
+    await fc.assert(
+      fc.asyncProperty(segment, async (countryId) => {
+        vi.mocked(notFound).mockClear();
+        const supported = (countryIds as readonly string[]).includes(countryId);
+
+        expect(await outcome(countryId)).toBe(supported ? 'rendered' : 'not found');
+        expect(notFound).toHaveBeenCalledTimes(supported ? 0 : 1);
+      })
     );
-    expect(replace).not.toHaveBeenCalled();
   });
+});
+
+describe('calculator paths outside the country segment', () => {
+  const calculatorApp = path.resolve(__dirname, '../../../../../calculator-app');
 
-  test('given an unsupported country then nothing renders and the route is replaced with /', () => {
-    const { container } = render(
-      <CountryLayout params={resolvedParams('xx')}>
-        <CountryProbe />
-      </CountryLayout>
+  // Next.js serves public files and static route folders before the
+  // [countryId] route, so an entry named exactly after a country would
+  // shadow that country's calculator.
+  test('given the calculator public folder then it is app/public', () => {
+    expect(fs.realpathSync(path.join(calculatorApp, 'public'))).toBe(
+      fs.realpathSync(path.join(calculatorApp, '../app/public'))
     );
-    expect(container).toBeEmptyDOMElement();
-    expect(replace).toHaveBeenCalledWith('/');
   });
+
+  test.each(['public', 'src/app'])(
+    'given the top level of calculator-app/%s then no entry is named after a country',
+    (folder) => {
+      const names = fs.readdirSync(path.join(calculatorApp, folder));
+
+      expect(names.filter((name) => (countryIds as readonly string[]).includes(name))).toEqual([]);
+    }
+  );
 });

app/src/tests/unit/calculator/NotFound.test.tsx

@@ -0,0 +1,22 @@
+import { render, screen } from '@test-utils';
+import { describe, expect, test } from 'vitest';
+import NotFound, { metadata } from '../../../../../calculator-app/src/app/not-found';
+
+describe('calculator 404 page', () => {
+  test('given a path the calculator does not serve then the page says it could not be found', () => {
+    render(<NotFound />);
+
+    expect(screen.getByText('404')).toBeInTheDocument();
+    expect(screen.getByText('This page could not be found.')).toBeInTheDocument();
+  });
+
+  test('given the 404 page then it links back to the calculator root', () => {
+    render(<NotFound />);
+
+    expect(screen.getByRole('link', { name: 'Go to the calculator' })).toHaveAttribute('href', '/');
+  });
+
+  test('given the 404 page then its title says the page was not found', () => {
+    expect(metadata.title).toBe('Page not found | PolicyEngine calculator');
+  });
+});

app/src/tests/unit/libs/countries.test.ts

@@ -0,0 +1,53 @@
+import fc from 'fast-check';
+import { describe, expect, test } from 'vitest';
+import { countryIds, isCountryId } from '@/libs/countries';
+
+describe('isCountryId', () => {
+  test.each(countryIds)('given the supported country %s then it is a country ID', (countryId) => {
+    expect(isCountryId(countryId)).toBe(true);
+  });
+
+  test.each([undefined, null, 0, {}, ['us'], 'US', 'usa', 'llms.txt', ''])(
+    'given %j then it is not a country ID',
+    (value) => {
+      expect(isCountryId(value)).toBe(false);
+    }
+  );
+
+  test('given any value then it is a country ID only when it is one of countryIds', () => {
+    fc.assert(
+      fc.property(fc.anything(), (value) => {
+        if (isCountryId(value)) {
+          expect(countryIds).toContain(value);
+        }
+      })
+    );
+  });
+
+  // Every country ID is two lower-case letters, so padding one never yields another.
+  test('given every supported country then it is two lower-case letters', () => {
+    for (const countryId of countryIds) {
+      expect(countryId).toMatch(/^[a-z]{2}$/);
+    }
+  });
+
+  test('given a supported country with anything added around it then it is not a country ID', () => {
+    fc.assert(
+      fc.property(
+        fc.constantFrom(...countryIds),
+        fc.string({ unit: 'binary' }),
+        fc.string({ unit: 'binary' }),
+        (countryId, before, after) => {
+          fc.pre(before.length + after.length > 0);
+          expect(isCountryId(`${before}${countryId}${after}`)).toBe(false);
+        }
+      )
+    );
+  });
+
+  test('given a supported country in upper case then it is not a country ID', () => {
+    for (const countryId of countryIds) {
+      expect(isCountryId(countryId.toUpperCase())).toBe(false);
+    }
+  });
+});

calculator-app/src/app/[countryId]/CountryShell.tsx

@@ -0,0 +1,54 @@
+"use client";
+
+import { useMemo } from "react";
+import { usePathname, useRouter, useSearchParams } from "next/navigation";
+import { CalculatorProviders } from "./providers";
+import { CountryProvider } from "@/contexts/CountryContext";
+import { LocationProvider } from "@/contexts/LocationContext";
+import { NavigationProvider } from "@/contexts/NavigationContext";
+import type { CountryId } from "@/libs/countries";
+
+/**
+ * Provides CountryContext, NavigationContext, and LocationContext to the
+ * shared components in app/src. The server layout has already rejected
+ * unsupported countries, so `countryId` is always a supported one.
+ */
+export function CountryShell({
+  countryId,
+  children,
+}: {
+  countryId: CountryId;
+  children: React.ReactNode;
+}) {
+  const router = useRouter();
+  const pathname = usePathname();
+  const searchParams = useSearchParams();
+
+  const navValue = useMemo(
+    () => ({
+      push: (path: string) => router.push(path),
+      replace: (path: string) => router.replace(path),
+      back: () => router.back(),
+      prefetch: (path: string) => router.prefetch(path),
+    }),
+    [router],
+  );
+
+  const locationValue = useMemo(
+    () => ({
+      pathname,
+      search: searchParams.toString() ? `?${searchParams.toString()}` : "",
+    }),
+    [pathname, searchParams],
+  );
+
+  return (
+    <CountryProvider value={countryId}>
+      <NavigationProvider value={navValue}>
+        <LocationProvider value={locationValue}>
+          <CalculatorProviders>{children}</CalculatorProviders>
+        </LocationProvider>
+      </NavigationProvider>
+    </CountryProvider>
+  );
+}