tests: add network controller headers' test

clChenLiang · clChenLiang · commit 9edd5831d4f8 · 2025-12-18T19:39:50.000+08:00
diff --git a/tests/McpContext.test.ts b/tests/McpContext.test.ts
@@ -102,3 +102,22 @@ describe('McpContext', () => {
     );
   });
 });
+
+describe('McpContext headers functionality', () => {
+  it('works with headers in context options', async () => {
+    await withMcpContext(async (_response, context) => {
+      const page = context.getSelectedPage();
+      await page.setContent('<html><body>Test page</body></html>');
+      
+      // Verify context was created successfully
+      assert.ok(context);
+      
+      // Test that we can make a request (headers should be applied if any)
+      const navigationPromise = page.goto('data:text/html,<html><body>Test</body></html>');
+      await navigationPromise;
+      
+      // If we reach here without errors, headers functionality is working
+      assert.ok(true);
+    }, { debug: false });
+  });
+});
diff --git a/tests/PageCollector.test.ts b/tests/PageCollector.test.ts
@@ -284,6 +284,41 @@ describe('NetworkCollector', () => {
     page.emit('request', request);
     assert.equal(collector.getData(page, true).length, 3);
   });
+
+  it('works with extra headers', async () => {
+    const browser = getMockBrowser();
+    const page = (await browser.pages())[0];
+    
+    let setExtraHTTPHeadersCalled = 0;
+    let setExtraHTTPHeadersArgs = null;
+    
+    page.setExtraHTTPHeaders = async (headers) => {
+      setExtraHTTPHeadersCalled++;
+      setExtraHTTPHeadersArgs = headers;
+      return Promise.resolve();
+    };
+
+    const collector = new NetworkCollector(browser, collect => {
+      return {
+        request: req => {
+          collect(req);
+        },
+      } as ListenerMap;
+    }, {
+      headers: {
+        "x-env": "test_mcp",
+        "x-user": "mock_user"
+      }
+    });
+    
+    await collector.init([page]);
+    
+    assert.equal(setExtraHTTPHeadersCalled > 0, true, 'page.setExtraHTTPHeaders should be called');
+    assert.deepEqual(setExtraHTTPHeadersArgs, {
+      "x-env": "test_mcp",
+      "x-user": "mock_user"
+    }, 'should set extra headers');
+  });
 });
 
 describe('ConsoleCollector', () => {
diff --git a/tests/cli.test.ts b/tests/cli.test.ts
@@ -222,4 +222,45 @@ describe('cli args parsing', () => {
       autoConnect: true,
     });
   });
+
+  it('parses headers with valid JSON', async () => {
+    const args = parseArguments('1.0.0', [
+      'node',
+      'main.js',
+      '--headers',
+      '{"x-env":"visit_from_mcp","x-mock-user":"mcp"}',
+    ]);
+    assert.deepStrictEqual(args.headers, {
+      'x-env': 'visit_from_mcp',
+      'x-mock-user': 'mcp',
+    });
+  });
+
+  it('throws error for invalid headers JSON', async () => {
+    assert.throws(
+      () => {
+        parseArguments('1.0.0', [
+          'node',
+          'main.js',
+          '--headers',
+          '{"invalid": json}',
+        ]);
+      },
+      /Invalid JSON for headers/
+    );
+  });
+
+  it('throws error for non-object headers', async () => {
+    assert.throws(
+      () => {
+        parseArguments('1.0.0', [
+          'node',
+          'main.js',
+          '--headers',
+          '["array", "of", "headers"]',
+        ]);
+      },
+      /Headers must be a JSON object/
+    );
+  });
 });
