Allow customizing Chrome page captures, sockets, and connections

README.md

@@ -425,6 +425,39 @@ Drivers receive `SourceType::Html` or `SourceType::Url` and handle every output
 
 `$options->format()` returns `OutputFormat::Png`, `OutputFormat::Jpeg`, or `OutputFormat::Pdf`. Return an `ImageResult` for images or a `PdfResult` for PDFs. Each driver call receives an independent snapshot of both option objects.
 
+### Extending Chrome
+
+Extend `ChromeDriver` to customize Chrome while keeping its capture and cleanup behavior. Register the subclass through `Screenshot::extend()` like any other driver.
+
+The driver provides these protected methods:
+
+- `createBrowserFactory()` supplies the factory used to start Chrome.
+- `getBrowserOptions(BrowserOptions $options)` prepares launch options, such as flags and the browser profile.
+- `getConnectionOptions(?BrowserOptions $options = null)` prepares options used for startup, shared-browser reconnection, and shutdown. Shared captures also apply its HTTP headers to the connection before creating a page.
+- `createPageCapture(Browser $browser, CaptureOptions $options, bool $shared = false)` supplies the `PageCapture` instance. Subclass `PageCapture` to override `configurePage()`, `loadPage()`, or `newCapture()`.
+- `createSocket(string $key)` supplies the `BrowserSocket` used for both shared captures and `close()`.
+
+For example, apply custom headers to both dedicated and shared captures:
+
+```php
+use DirectoryTree\Screenshot\BrowserOptions;
+use DirectoryTree\Screenshot\Drivers\ChromeDriver;
+
+class CustomChromeDriver extends ChromeDriver
+{
+    protected function getConnectionOptions(?BrowserOptions $options = null): array
+    {
+        $options = parent::getConnectionOptions($options);
+
+        $options['headers']['X-Renderer'] = 'invoices';
+
+        return $options;
+    }
+}
+```
+
+`getBrowserOptions()` includes `getConnectionOptions()` by default. Override `getConnectionOptions()` for settings that must also apply when reusing a browser; launch-only overrides do not run on reconnection. During `close()`, no capture options are available, so the base method retains the driver's configured connection timeout.
+
 ## Testing
 
 Use the facade fake to test your application without starting Chrome:

src/Drivers/ChromeDriver.php

@@ -38,7 +38,7 @@ public function capture(string $source, SourceType $type, CaptureOptions $option
             return $socket->run(
                 fn () => $this->createBrowser($options->browser()),
                 fn (Browser $browser) => $this->capturePage($browser, $source, $type, $options, shared: true),
-                array_merge($this->browserOptions, ['sendSyncDefaultTimeout' => $options->browser()->getBrowserTimeout()]),
+                $this->getConnectionOptions($options->browser()),
             );
         }
 
@@ -53,11 +53,15 @@ protected function capturePage(Browser $browser, string $source, SourceType $typ
         $failed = true;
 
         try {
-            if ($shared && isset($this->browserOptions['headers'])) {
-                $browser->getConnection()->setConnectionHttpHeaders($this->browserOptions['headers']);
+            if ($shared) {
+                $connectionOptions = $this->getConnectionOptions($options->browser());
+
+                if (isset($connectionOptions['headers'])) {
+                    $browser->getConnection()->setConnectionHttpHeaders($connectionOptions['headers']);
+                }
             }
 
-            $result = (new PageCapture($browser, $options, $shared))->capture($source, $type);
+            $result = $this->createPageCapture($browser, $options, $shared)->capture($source, $type);
 
             $failed = false;
 
@@ -75,6 +79,14 @@ protected function capturePage(Browser $browser, string $source, SourceType $typ
         }
     }
 
+    /**
+     * Create the capture used to prepare and render a browser page.
+     */
+    protected function createPageCapture(Browser $browser, CaptureOptions $options, bool $shared = false): PageCapture
+    {
+        return new PageCapture($browser, $options, $shared);
+    }
+
     /**
      * Start a dedicated browser or a persistent shared browser.
      */
@@ -100,19 +112,28 @@ protected function getBrowserOptions(BrowserOptions $options): array
     {
         $socket = $this->socket($options);
 
-        $browserOptions = $this->browserOptions;
+        $browserOptions = $this->getConnectionOptions($options);
 
         if ($socket && empty($browserOptions['userDataDir'])) {
             $browserOptions['userDataDir'] = $socket->path().'.profile';
         }
 
         return array_merge($browserOptions, [
             'windowSize' => [$options->getWidth(), $options->getHeight()],
-            'sendSyncDefaultTimeout' => $options->getBrowserTimeout(),
             'keepAlive' => $socket !== null,
         ]);
     }
 
+    /**
+     * Prepare connection settings, using capture timeouts when available.
+     */
+    protected function getConnectionOptions(?BrowserOptions $options = null): array
+    {
+        return array_merge($this->browserOptions, $options ? [
+            'sendSyncDefaultTimeout' => $options->getBrowserTimeout(),
+        ] : []);
+    }
+
     /**
      * Resolve the shared socket selected for this capture.
      */
@@ -122,6 +143,14 @@ protected function socket(BrowserOptions $options): ?BrowserSocket
             return null;
         }
 
+        return $this->createSocket($key);
+    }
+
+    /**
+     * Create the socket used to capture with or close a shared browser.
+     */
+    protected function createSocket(string $key): BrowserSocket
+    {
         return BrowserSocket::named($key, $this->socketDirectory, $this->lockTimeout);
     }
 
@@ -130,6 +159,6 @@ protected function socket(BrowserOptions $options): ?BrowserSocket
      */
     public function close(string $key): void
     {
-        BrowserSocket::named($key, $this->socketDirectory, $this->lockTimeout)->close($this->browserOptions);
+        $this->createSocket($key)->close($this->getConnectionOptions());
     }
 }

tests/Unit/ChromeDriverTest.php

@@ -2,11 +2,15 @@
 
 use DirectoryTree\Screenshot\BrowserOptions;
 use DirectoryTree\Screenshot\CaptureOptions;
+use DirectoryTree\Screenshot\Drivers\Chrome\PageCapture;
 use DirectoryTree\Screenshot\Drivers\ChromeDriver;
 use DirectoryTree\Screenshot\Enums\SourceType;
 use DirectoryTree\Screenshot\ImageOptions;
+use DirectoryTree\Screenshot\ImageResult;
 use DirectoryTree\Screenshot\Support\BrowserSocket;
 use HeadlessChromium\Browser;
+use HeadlessChromium\Browser\ProcessAwareBrowser;
+use HeadlessChromium\BrowserFactory;
 use HeadlessChromium\Clip;
 use HeadlessChromium\Communication\Connection;
 use HeadlessChromium\Communication\Message;
@@ -290,3 +294,109 @@ protected function socket(BrowserOptions $options): ?BrowserSocket
         fn () => (new ChromeDriver)->capture('invoice', $type, new CaptureOptions(new BrowserOptions, new ImageOptions)),
     )->toThrow(TypeError::class);
 })->with(['html', 'url', 'invalid']);
+
+test('driver factories and connection options apply throughout the browser lifecycle', function ($mode) {
+    $shared = $mode !== 'dedicated';
+    $browser = Mockery::mock(ProcessAwareBrowser::class);
+    $factory = Mockery::mock(BrowserFactory::class);
+    $socket = Mockery::mock(BrowserSocket::class);
+    $pageCapture = Mockery::mock(PageCapture::class);
+    $result = new ImageResult(imageFixture());
+    $connectionOptions = [
+        'customFlags' => ['--disable-gpu'],
+        'headers' => ['X-Capture' => 'custom'],
+        'sendSyncDefaultTimeout' => 4321,
+        'connectionDelay' => 0.01,
+    ];
+
+    if ($mode === 'reused') {
+        $factory->shouldNotReceive('createBrowser');
+    } else {
+        $launchOptions = $connectionOptions;
+
+        if ($shared) {
+            $socket->shouldReceive('path')->once()->andReturn('/tmp/custom-browser.socket');
+            $launchOptions['userDataDir'] = '/tmp/custom-browser.socket.profile';
+        }
+
+        $factory->shouldReceive('createBrowser')->once()->with(array_merge($launchOptions, [
+            'windowSize' => [800, 600],
+            'keepAlive' => $shared,
+        ]))->andReturn($browser);
+    }
+
+    if ($shared) {
+        $socket->shouldReceive('run')->once()
+            ->with(Mockery::type('callable'), Mockery::type('callable'), $connectionOptions)
+            ->andReturnUsing(fn ($create, $capture) => $capture($mode === 'reused' ? $browser : $create()));
+
+        $socket->shouldReceive('close')->once()->with($connectionOptions);
+
+        $connection = Mockery::mock(Connection::class);
+        $browser->shouldReceive('getConnection')->once()->andReturn($connection);
+        $connection->shouldReceive('setConnectionHttpHeaders')->once()->with(['X-Capture' => 'custom']);
+        $browser->shouldNotReceive('close');
+    } else {
+        $browser->shouldReceive('close')->once();
+    }
+
+    $pageCapture->shouldReceive('capture')->once()->with('hello', SourceType::Html)->andReturn($result);
+
+    $driver = new class($factory, $socket, $pageCapture, $browser, $shared) extends ChromeDriver
+    {
+        public function __construct(
+            protected BrowserFactory $factory,
+            protected BrowserSocket $browserSocket,
+            protected PageCapture $pageCapture,
+            protected Browser $browser,
+            protected bool $shared,
+        ) {
+            parent::__construct(browserOptions: [
+                'customFlags' => ['--disable-gpu'],
+                'headers' => ['X-Capture' => 'original'],
+                'sendSyncDefaultTimeout' => 999,
+            ]);
+        }
+
+        protected function createBrowserFactory(): BrowserFactory
+        {
+            return $this->factory;
+        }
+
+        protected function createSocket(string $key): BrowserSocket
+        {
+            expect($key)->toBe('custom');
+
+            return $this->browserSocket;
+        }
+
+        protected function createPageCapture(Browser $browser, CaptureOptions $options, bool $shared = false): PageCapture
+        {
+            expect($browser)->toBe($this->browser);
+            expect($options->browser()->getWidth())->toBe(800);
+            expect($shared)->toBe($this->shared);
+
+            return $this->pageCapture;
+        }
+
+        protected function getConnectionOptions(?BrowserOptions $options = null): array
+        {
+            return array_merge(parent::getConnectionOptions($options), [
+                'headers' => ['X-Capture' => 'custom'],
+                'sendSyncDefaultTimeout' => 4321,
+                'connectionDelay' => 0.01,
+            ]);
+        }
+    };
+
+    $options = new CaptureOptions(
+        (new BrowserOptions)->setSize(800, 600)->setTimeouts(2000, 1234)->setSocket($shared ? 'custom' : null),
+        new ImageOptions,
+    );
+
+    expect($driver->capture('hello', SourceType::Html, $options))->toBe($result);
+
+    if ($shared) {
+        $driver->close('custom');
+    }
+})->with(['dedicated', 'shared', 'reused']);