Fix: 404.html load correctly on preview (#9907)

* Move vite 404 middleware

* Add custom 404.html rendering test for preview routing

* add a changest

* add TODO comment

.changeset/tidy-deers-double.md

@@ -0,0 +1,5 @@
+---
+"astro": patch
+---
+
+Load 404.html on all non-existent paths on astro preview.

packages/astro/src/core/preview/vite-plugin-astro-preview.ts

@@ -12,6 +12,18 @@ const HAS_FILE_EXTENSION_REGEXP = /^.*\.[^\\]+$/;
 export function vitePluginAstroPreview(settings: AstroSettings): Plugin {
 	const { base, outDir, trailingSlash } = settings.config;
 
+	function handle404(req: IncomingMessage, res: ServerResponse) {
+		const errorPagePath = fileURLToPath(outDir + '/404.html');
+		if (fs.existsSync(errorPagePath)) {
+			res.statusCode = 404;
+			res.setHeader('Content-Type', 'text/html;charset=utf-8');
+			res.end(fs.readFileSync(errorPagePath));
+		} else {
+			res.statusCode = 404;
+			res.end(notFoundTemplate(req.url!, 'Not Found'));
+		}
+	}
+
 	return {
 		name: 'astro:preview',
 		apply: 'serve',
@@ -48,6 +60,14 @@ export function vitePluginAstroPreview(settings: AstroSettings): Plugin {
 					}
 				}
 
+				// TODO: look into why the replacement needs to happen here
+				for (const middleware of server.middlewares.stack) {
+					// This hardcoded name will not break between Vite versions
+					if ((middleware.handle as Connect.HandleFunction).name === 'vite404Middleware') {
+						middleware.handle = handle404;
+					}
+				}
+
 				next();
 			});
 
@@ -77,25 +97,6 @@ export function vitePluginAstroPreview(settings: AstroSettings): Plugin {
 
 					next();
 				});
-
-				// Vite has its own 404 middleware, we replace it with ours instead.
-				for (const middleware of server.middlewares.stack) {
-					// This hardcoded name will not break between Vite versions
-					if ((middleware.handle as Connect.HandleFunction).name === 'vite404Middleware') {
-						// Fallback to 404 page if it exists
-						middleware.handle = (req: IncomingMessage, res: ServerResponse) => {
-							const errorPagePath = fileURLToPath(outDir + '/404.html');
-							if (fs.existsSync(errorPagePath)) {
-								res.statusCode = 404;
-								res.setHeader('Content-Type', 'text/html;charset=utf-8');
-								res.end(fs.readFileSync(errorPagePath));
-							} else {
-								res.statusCode = 404;
-								res.end(notFoundTemplate(req.url!, 'Not Found'));
-							}
-						};
-					}
-				}
 			};
 		},
 	};

packages/astro/test/preview-routing.test.js

@@ -1,4 +1,5 @@
 import { expect } from 'chai';
+import * as cheerio from 'cheerio';
 import { loadFixture } from './test-utils.js';
 
 describe('Preview Routing', function () {
@@ -182,6 +183,41 @@ describe('Preview Routing', function () {
 				expect(response.status).to.equal(404);
 			});
 		});
+
+		describe('Load custom 404.html', () => {
+			/** @type {import('./test-utils').Fixture} */
+			let fixture;
+			/** @type {import('./test-utils').PreviewServer} */
+			let previewServer;
+
+			let $;
+
+			before(async () => {
+				fixture = await loadFixture({
+					root: './fixtures/custom-404-html/',
+					server: {
+						port: 4003,
+					},
+				});
+				await fixture.build();
+				previewServer = await fixture.preview();
+			});
+
+			after(async () => {
+				await previewServer.stop();
+			});
+
+			it('renders custom 404 for /a', async () => {
+				const res = await fixture.fetch('/a');
+				expect(res.status).to.equal(404);
+
+				const html = await res.text();
+				$ = cheerio.load(html);
+
+				expect($('h1').text()).to.equal('Page not found');
+				expect($('p').text()).to.equal('This 404 is a static HTML file.');
+			});
+		});
 	});
 
 	describe('build format: file', () => {
@@ -201,7 +237,7 @@ describe('Preview Routing', function () {
 					},
 					trailingSlash: 'never',
 					server: {
-						port: 4003,
+						port: 4004,
 					},
 				});
 				await fixture.build();
@@ -261,7 +297,7 @@ describe('Preview Routing', function () {
 					},
 					trailingSlash: 'always',
 					server: {
-						port: 4004,
+						port: 4005,
 					},
 				});
 				await fixture.build();
@@ -324,7 +360,7 @@ describe('Preview Routing', function () {
 					},
 					trailingSlash: 'ignore',
 					server: {
-						port: 4005,
+						port: 4006,
 					},
 				});
 				await fixture.build();
@@ -387,7 +423,7 @@ describe('Preview Routing', function () {
 					},
 					trailingSlash: 'ignore',
 					server: {
-						port: 4006,
+						port: 4007,
 					},
 				});
 				await fixture.build();
@@ -423,5 +459,43 @@ describe('Preview Routing', function () {
 				expect(response.status).to.equal(404);
 			});
 		});
+
+		describe('Load custom 404.html', () => {
+			/** @type {import('./test-utils').Fixture} */
+			let fixture;
+			/** @type {import('./test-utils').PreviewServer} */
+			let previewServer;
+
+			let $;
+
+			before(async () => {
+				fixture = await loadFixture({
+					root: './fixtures/custom-404-html/',
+					build: {
+						format: 'file',
+					},
+					server: {
+						port: 4008,
+					},
+				});
+				await fixture.build();
+				previewServer = await fixture.preview();
+			});
+
+			after(async () => {
+				await previewServer.stop();
+			});
+
+			it('renders custom 404 for /a', async () => {
+				const res = await fixture.fetch('/a');
+				expect(res.status).to.equal(404);
+
+				const html = await res.text();
+				$ = cheerio.load(html);
+
+				expect($('h1').text()).to.equal('Page not found');
+				expect($('p').text()).to.equal('This 404 is a static HTML file.');
+			});
+		});
 	});
 });