test: refactor tests to reduce duplication

Author: nroscino

tests/index.test.ts

@@ -12,10 +12,7 @@ import {Client} from '@modelcontextprotocol/sdk/client/index.js';
 import {StdioClientTransport} from '@modelcontextprotocol/sdk/client/stdio.js';
 import {executablePath} from 'puppeteer';
 
-import {
-  OFF_BY_DEFAULT_CATEGORIES,
-  ToolCategory,
-} from '../src/tools/categories.js';
+import {OFF_BY_DEFAULT_CATEGORIES} from '../src/tools/categories.js';
 import type {ToolDefinition} from '../src/tools/ToolDefinition.js';
 
 describe('e2e', () => {
@@ -80,44 +77,7 @@ describe('e2e', () => {
       async client => {
         const {tools} = await client.listTools();
         const exposedNames = tools.map(t => t.name).sort();
-        const files = fs.readdirSync('build/src/tools');
-        const definedNames = [];
-        for (const file of files) {
-          if (
-            file === 'ToolDefinition.js' ||
-            file === 'tools.js' ||
-            file === 'slim'
-          ) {
-            continue;
-          }
-          const fileTools = await import(`../src/tools/${file}`);
-          for (const maybeTool of Object.values<unknown>(fileTools)) {
-            if (typeof maybeTool === 'function') {
-              const tool = (maybeTool as (val: boolean) => ToolDefinition)(
-                false,
-              );
-              if (tool && typeof tool === 'object' && 'name' in tool) {
-                if (tool.annotations?.conditions) {
-                  continue;
-                }
-                definedNames.push(tool.name);
-              }
-              continue;
-            }
-            if (
-              typeof maybeTool === 'object' &&
-              maybeTool !== null &&
-              'name' in maybeTool
-            ) {
-              const tool = maybeTool as ToolDefinition;
-              if (tool.annotations?.conditions) {
-                continue;
-              }
-              definedNames.push(tool.name);
-            }
-          }
-        }
-
+        const definedNames = await getToolsWithFilteredCategories();
         definedNames.sort();
         assert.deepStrictEqual(exposedNames, definedNames);
       },
@@ -129,54 +89,9 @@ describe('e2e', () => {
     await withClient(async client => {
       const {tools} = await client.listTools();
       const exposedNames = tools.map(t => t.name).sort();
-      const files = fs.readdirSync('build/src/tools');
-      const definedNames = [];
-      for (const file of files) {
-        if (
-          file === 'ToolDefinition.js' ||
-          file === 'tools.js' ||
-          file === 'slim'
-        ) {
-          continue;
-        }
-        const fileTools = await import(`../src/tools/${file}`);
-        for (const maybeTool of Object.values<unknown>(fileTools)) {
-          if (typeof maybeTool === 'function') {
-            const tool = (maybeTool as (val: boolean) => ToolDefinition)(false);
-            if (tool && typeof tool === 'object' && 'name' in tool) {
-              if (tool.annotations?.conditions) {
-                continue;
-              }
-              if (
-                tool.annotations?.category &&
-                tool.annotations?.category === ToolCategory.EXTENSIONS
-              ) {
-                continue;
-              }
-              definedNames.push(tool.name);
-            }
-            continue;
-          }
-          if (
-            typeof maybeTool === 'object' &&
-            maybeTool !== null &&
-            'name' in maybeTool
-          ) {
-            const tool = maybeTool as ToolDefinition;
-            if (tool.annotations?.conditions) {
-              continue;
-            }
-            if (
-              tool.annotations?.category &&
-              tool.annotations?.category === ToolCategory.EXTENSIONS
-            ) {
-              continue;
-            }
-            definedNames.push(tool.name);
-          }
-        }
-      }
-
+      const definedNames = await getToolsWithFilteredCategories(
+        OFF_BY_DEFAULT_CATEGORIES,
+      );
       definedNames.sort();
       assert.deepStrictEqual(exposedNames, definedNames);
     });
@@ -245,3 +160,56 @@ describe('e2e', () => {
     );
   });
 });
+
+async function getToolsWithFilteredCategories(
+  filterOutCategories: string[] = [],
+): Promise<string[]> {
+  const files = fs.readdirSync('build/src/tools');
+  const definedNames = [];
+  for (const file of files) {
+    if (
+      file === 'ToolDefinition.js' ||
+      file === 'tools.js' ||
+      file === 'slim'
+    ) {
+      continue;
+    }
+    const fileTools = await import(`../src/tools/${file}`);
+    for (const maybeTool of Object.values<unknown>(fileTools)) {
+      if (typeof maybeTool === 'function') {
+        const tool = (maybeTool as (val: boolean) => ToolDefinition)(false);
+        if (tool && typeof tool === 'object' && 'name' in tool) {
+          if (tool.annotations?.conditions) {
+            continue;
+          }
+          if (
+            tool.annotations?.category &&
+            filterOutCategories.includes(tool.annotations?.category)
+          ) {
+            continue;
+          }
+          definedNames.push(tool.name);
+        }
+        continue;
+      }
+      if (
+        typeof maybeTool === 'object' &&
+        maybeTool !== null &&
+        'name' in maybeTool
+      ) {
+        const tool = maybeTool as ToolDefinition;
+        if (tool.annotations?.conditions) {
+          continue;
+        }
+        if (
+          tool.annotations?.category &&
+          filterOutCategories.includes(tool.annotations?.category)
+        ) {
+          continue;
+        }
+        definedNames.push(tool.name);
+      }
+    }
+  }
+  return definedNames;
+}