Fix tests for Vitest and adjust guard expectations

Author: Martin

bun.lock

@@ -48,9 +48,9 @@
     "build": "bunup",
     "watch": "tsc --watch",
     "inspector": "npx @modelcontextprotocol/inspector dist/index.js",
-    "test": "bun test",
-    "test:watch": "bun test --watch",
-    "test:cov": "bun test --coverage",
+    "test": "bunx vitest run",
+    "test:watch": "bunx vitest watch",
+    "test:cov": "bunx vitest run --coverage",
     "lint": "biome lint .",
     "lint:fix": "biome lint --write .",
     "format": "biome format --write .",

package.json

@@ -29,9 +29,11 @@ vi.mock('../../src/pdf/loader.js', () => ({
 }));
 
 vi.mock('../../src/pdf/extractor.js', async () => {
-  const actual = await vi.importActual<typeof import('../../src/pdf/extractor.js')>('../../src/pdf/extractor.js');
   return {
-    ...actual,
+    buildWarnings: (invalidPages: number[], totalPages: number) => {
+      if (invalidPages.length === 0) return [];
+      return [`Requested page numbers ${invalidPages.join(', ')} exceed total pages (${String(totalPages)}).`];
+    },
     extractPageContent: mockExtractPageContent,
     extractImages: mockExtractImages,
     extractMetadataAndPageCount: mockExtractMetadataAndPageCount,

test/handlers/pageGuards.test.ts

@@ -14,8 +14,10 @@ const mockGetMetadata = vi.fn();
 const mockGetPage = vi.fn();
 const mockGetDocument = vi.fn();
 const mockReadFile = vi.fn();
+const mockStat = vi.fn();
 const mockExtractPageContent = vi.fn();
 let actualExtractPageContent: typeof import('../../src/pdf/extractor.js')['extractPageContent'];
+let extractorModule: typeof import('../../src/pdf/extractor.js');
 
 vi.mock('pdfjs-dist/legacy/build/pdf.mjs', () => ({
   getDocument: mockGetDocument,
@@ -28,19 +30,12 @@ vi.mock('pdfjs-dist/legacy/build/pdf.mjs', () => ({
 vi.mock('node:fs/promises', () => ({
   default: {
     readFile: mockReadFile,
+    stat: mockStat,
   },
   readFile: mockReadFile,
+  stat: mockStat,
 }));
 
-vi.mock('../../src/pdf/extractor.js', async () => {
-  const actual = await vi.importActual<typeof import('../../src/pdf/extractor.js')>('../../src/pdf/extractor.js');
-  actualExtractPageContent = actual.extractPageContent;
-  return {
-    ...actual,
-    extractPageContent: mockExtractPageContent,
-  };
-});
-
 // Dynamically import the handler *once* after mocks are defined
 // Define a more specific type for the handler's return value content
 interface HandlerResultContent {
@@ -51,6 +46,10 @@ let handler: (args: unknown) => Promise<{ content: HandlerResultContent[] }>;
 let readPdfSchema: Schema<unknown>;
 
 beforeAll(async () => {
+  extractorModule = await import('../../src/pdf/extractor.js');
+  actualExtractPageContent = extractorModule.extractPageContent;
+  vi.spyOn(extractorModule, 'extractPageContent').mockImplementation(mockExtractPageContent);
+
   // Import the readPdf tool - the new SDK uses a builder pattern
   const { readPdf } = await import('../../src/handlers/readPdf.js');
   const { readPdfArgsSchema } = await import('../../src/schemas/readPdf.js');
@@ -90,6 +89,7 @@ describe('handleReadPdfFunc Integration Tests', () => {
   beforeEach(() => {
     vi.resetAllMocks();
     mockExtractPageContent.mockImplementation((...args) => actualExtractPageContent(...args));
+    vi.spyOn(extractorModule, 'extractPageContent').mockImplementation(mockExtractPageContent);
     // Reset mocks for pathUtils if we spy on it
     vi.spyOn(pathUtils, 'resolvePath').mockImplementation((p) => p); // Simple mock for resolvePath
 
@@ -164,6 +164,9 @@ describe('handleReadPdfFunc Integration Tests', () => {
             metadata: { 'dc:format': 'application/pdf' },
             num_pages: 3,
             full_text: 'Mock page text 1\n\nMock page text 2\n\nMock page text 3',
+            warnings: [
+              'No pages specified; processed available pages because the document is small. Specify pages or set allow_full_document=true to control full-document requests.',
+            ],
           },
         },
       ],
@@ -316,6 +319,7 @@ describe('handleReadPdfFunc Integration Tests', () => {
     const args = {
       sources: [{ path: 'many-pages.pdf' }],
       include_full_text: true,
+      allow_full_document: true,
     };
 
     const result = await handler(args);
@@ -467,6 +471,9 @@ describe('handleReadPdfFunc Integration Tests', () => {
             metadata: { 'dc:creator': 'URL Author' },
             num_pages: 2,
             full_text: 'URL Mock page text 1\n\nURL Mock page text 2',
+            warnings: [
+              'No pages specified; processed available pages because the document is small. Specify pages or set allow_full_document=true to control full-document requests.',
+            ],
           },
         },
       ],

test/handlers/readPdf.test.ts

@@ -193,8 +193,6 @@ describe('extractor', () => {
     });
 
     it('processes pages in batches to avoid unbounded work on large documents', async () => {
-      vi.useFakeTimers();
-
       let activeRequests = 0;
       let maxConcurrentRequests = 0;
 
@@ -209,28 +207,21 @@ describe('extractor', () => {
           activeRequests++;
           maxConcurrentRequests = Math.max(maxConcurrentRequests, activeRequests);
 
-          return new Promise((resolve) =>
+          return new Promise((resolve) => {
             setTimeout(() => {
               activeRequests--;
               resolve(mockPage);
-            }, 10)
-          );
+            }, 5);
+          });
         }),
       } as unknown as pdfjsLib.PDFDocumentProxy;
 
       const pages = Array.from({ length: 10 }, (_, idx) => idx + 1);
 
-      try {
-        const extractionPromise = extractPageTexts(mockDocument, pages, 'large.pdf');
-
-        await vi.runAllTimersAsync();
-        const result = await extractionPromise;
+      const result = await extractPageTexts(mockDocument, pages, 'large.pdf');
 
-        expect(maxConcurrentRequests).toBeLessThanOrEqual(6);
-        expect(result).toHaveLength(10);
-      } finally {
-        vi.useRealTimers();
-      }
+      expect(maxConcurrentRequests).toBeLessThanOrEqual(6);
+      expect(result).toHaveLength(10);
     });
   });