src/hero.mjs: findHandler(): learn how to find upgrade routes

2 files changed, 73 insertions, 32 deletions

diff --git a/src/hero.mjs b/src/hero.mjs
index 131b796..7483a21 100644
--- a/src/hero.mjs
+++ b/src/hero.mjs
@@ -314,8 +314,8 @@ export const addRoute = (table, methods, path, handlerFn) => {
 	);
 };
 
-export const findStaticHandler = (table, method, segments) => {
-	const handlerFn = u.getIn(table, ["static", method].concat(segments));
+export const findStaticHandler = (table, method, segments, section) => {
+	const handlerFn = u.getIn(table, [section, method].concat(segments));
 	return !handlerFn ? null : { handlerFn, params: {} };
 };
 
@@ -362,12 +362,14 @@ export const findDynamicHandler = (table, method, segments) => {
 	return !tree ? null : firstParamMatch(tree, segments, {});
 };
 
-export const findHandler = (table, method, path) => {
+export const findHandler = (table, method, path, upgrade) => {
 	const segments = pathToSegments(path);
-	return (
-		 findStaticHandler(table, method, segments) ||
-		findDynamicHandler(table, method, segments)
-	);
+	return upgrade ?
+		findStaticHandler(table, method, segments, "websocket") :
+		(
+			 findStaticHandler(table, method, segments, "static") ||
+			findDynamicHandler(table, method, segments)
+		);
 };
 
 export const extractQueryParams = s => {
@@ -392,7 +394,7 @@ export const handleRequest = async (table, reqHandle) => {
 	const { method, url, headers, upgrade, socket } = reqHandle;
 	const [ path, queryParams ] = url.split("?");
 	const { params, handlerFn } = (
-		findHandler(table, method, path) ||
+		findHandler(table, method, path, upgrade) ||
 		make404Handler(table.interceptors)
 	);
 
diff --git a/tests/js/hero.mjs b/tests/js/hero.mjs
index 817f301..7a5738d 100644
--- a/tests/js/hero.mjs
+++ b/tests/js/hero.mjs
@@ -982,6 +982,7 @@ const test_findStaticHandler = async t => {
 		const fn1 = () => {};
 		const fn2 = () => {};
 		const fn3 = () => {};
+		const fn4 = () => {};
 
 		const table = {
 			static: {
@@ -997,39 +998,56 @@ const test_findStaticHandler = async t => {
 					},
 				},
 			},
+			websocket: {
+				GET: {
+					api: {
+						socket: { "": fn4 },
+					},
+				},
+			}
 		};
 
 		assert.deepEqual(
-			findStaticHandler(table, "GET", [ "api", "home", "" ]),
+			findStaticHandler(table, "GET", [ "api", "home", "" ], "static"),
 			{ handlerFn: fn1, params: {} },
 		);
 		assert.deepEqual(
-			findStaticHandler(table, "PUT", [ "api", "home", "" ]),
+			findStaticHandler(table, "PUT", [ "api", "home", "" ], "static"),
 			null,
 		);
 
 		assert.deepEqual(
-			findStaticHandler(table, "GET", [ "api", "settings", "" ]),
+			findStaticHandler(table, "GET", [ "api", "settings", "" ], "static"),
 			{ handlerFn: fn2, params: {} },
 		);
 		assert.deepEqual(
-			findStaticHandler(table, "POST", [ "api", "settings", "" ]),
+			findStaticHandler(table, "POST", [ "api", "settings", "" ], "static"),
 			{ handlerFn: fn3, params: {} },
 		);
 		assert.deepEqual(
-			findStaticHandler(table, "PUT", [ "api", "settings", "" ]),
+			findStaticHandler(table, "PUT", [ "api", "settings", "" ], "static"),
+			null,
+		);
+
+		assert.deepEqual(
+			findStaticHandler(table, "GET", [ "api", "profile", "" ], "static"),
 			null,
 		);
 
 		assert.deepEqual(
-			findStaticHandler(table, "GET", [ "api", "profile", "" ]),
+			findStaticHandler({}, "GET", [ "api", "profile", "" ], "static"),
 			null,
 		);
 
 		assert.deepEqual(
-			findStaticHandler({}, "GET", [ "api", "profile", "" ]),
+			findStaticHandler(table, "GET", [ "api", "socket", "" ], "static"),
 			null,
 		);
+
+		assert.deepEqual(
+			findStaticHandler(table, "GET", [ "api", "socket", "" ], "websocket"),
+			{ handlerFn: fn4, params: {} },
+		);
 	});
 };
 
@@ -1182,14 +1200,16 @@ const test_findDynamicHandler = async t => {
 const test_findHandler = async t => {
 	t.start("findHandler()");
 
-	await t.test("mix of static and dynamic routes", () => {
-		const static1  = () => {};
-		const static2  = () => {};
-		const static3  = () => {};
-		const dynamic1 = () => {};
-		const dynamic2 = () => {};
-		const dynamic3 = () => {};
-		const dynamic4 = () => {};
+	await t.test("mix of static, dynamic and websocket routes", () => {
+		const static1    = () => {};
+		const static2    = () => {};
+		const static3    = () => {};
+		const dynamic1   = () => {};
+		const dynamic2   = () => {};
+		const dynamic3   = () => {};
+		const dynamic4   = () => {};
+		const websocket1 = () => {};
+		const websocket2 = () => {};
 
 		const table = {
 			static: {
@@ -1227,48 +1247,67 @@ const test_findHandler = async t => {
 					},
 				},
 			},
+			websocket: {
+				GET: {
+					user: {
+						"": websocket1,
+						socket: {
+							"": websocket2,
+						},
+					},
+				},
+			},
 		};
 
 		assert.deepEqual(
-			findHandler(table, "GET", "/"),
+			findHandler(table, "GET", "/", false),
 			null,
 		);
 
 		assert.deepEqual(
-			findHandler(table, "GET", "/user/"),
+			findHandler(table, "GET", "/user/", false),
 			{ handlerFn: static1, params: {} },
 		);
+		assert.deepEqual(
+			findHandler(table, "GET", "/user/", true),
+			{ handlerFn: websocket1, params: {} },
+		);
 
 		assert.deepEqual(
-			findHandler(table, "GET", "/pages"),
+			findHandler(table, "GET", "/pages", false),
 			{ handlerFn: static2, params: {} },
 		);
 		assert.deepEqual(
-			findHandler(table, "GET", "/pages/home/"),
+			findHandler(table, "GET", "/pages/home/", false),
 			{ handlerFn: static3, params: {} },
 		);
 
 		assert.deepEqual(
-			findHandler(table, "GET", "/user/some-id"),
+			findHandler(table, "GET", "/user/some-id", false),
 			{ handlerFn: dynamic1, params: { id: "some-id" }},
 		);
 		assert.deepEqual(
-			findHandler(table, "GET", "/user/other-id/info"),
+			findHandler(table, "GET", "/user/other-id/info", false),
 			null,
 		);
 
 		assert.deepEqual(
-			findHandler(table, "PUT", "/user/other-id/info"),
+			findHandler(table, "PUT", "/user/other-id/info", false),
 			{ handlerFn: dynamic3, params: { id: "other-id" }},
 		);
 		assert.deepEqual(
-			findHandler(table, "PUT", "/user/another-id/preferences"),
+			findHandler(table, "PUT", "/user/another-id/preferences", false),
 			{ handlerFn: dynamic4, params: { id: "another-id" }},
 		);
 		assert.deepEqual(
-			findHandler(table, "POST", "/user/another-id/preferences"),
+			findHandler(table, "POST", "/user/another-id/preferences", false),
 			null,
 		);
+
+		assert.deepEqual(
+			findHandler(table, "GET", "/user/socket", true),
+			{ handlerFn: websocket2, params: {} },
+		);
 	});
 };