From f31a09bbf3e343298d12e54fffe73fa66fe9b174 Mon Sep 17 00:00:00 2001
From: Kevin Martensson <kevinmartensson@gmail.com>
Date: Sun, 24 Sep 2017 23:59:03 +0200
Subject: [PATCH] Add `keepAlive` and `browser` options

---
 index.js | 12 +++++++++---
 test.js  | 18 +++++++++++++++++-
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/index.js b/index.js
index 09cb140..74132e3 100644
--- a/index.js
+++ b/index.js
@@ -34,7 +34,7 @@ module.exports = async (url, opts) => {
 
 	const uri = isUrl(url) ? url : fileUrl(url);
 	const {
-		cookies, crop, format, headers, height, hide, password, scale,
+		cookies, crop, format, headers, height, hide, keepAlive, password, scale,
 		script, selector, timeout, transparent, userAgent, username, width
 	} = opts;
 
@@ -52,7 +52,7 @@ module.exports = async (url, opts) => {
 		opts.omitBackground = true;
 	}
 
-	const browser = await puppeteer.launch();
+	const browser = opts.browser || await puppeteer.launch();
 	const page = await browser.newPage();
 	const viewport = {
 		height,
@@ -96,7 +96,13 @@ module.exports = async (url, opts) => {
 	}
 
 	const buf = await page.screenshot(opts);
-	await browser.close();
+	await page.close();
+
+	if (keepAlive !== true) {
+		await browser.close();
+	}
 
 	return buf;
 };
+
+module.exports.startBrowser = puppeteer.launch;
diff --git a/test.js b/test.js
index 2e876d3..d4a3ab7 100644
--- a/test.js
+++ b/test.js
@@ -6,7 +6,23 @@ import isPng from 'is-png';
 import pify from 'pify';
 import PNG from 'png-js';
 import server from './fixtures/server';
-import m from '.';
+import screenshotStream, {startBrowser} from '.';
+
+let browser;
+let m;
+
+test.before(async () => {
+	browser = await startBrowser();
+
+	m = (url, opts) => screenshotStream(url, Object.assign({}, opts, {
+		browser,
+		keepAlive: true
+	}));
+});
+
+test.after(async () => {
+	await browser.close();
+});
 
 test('generate screenshot', async t => {
 	t.true(isPng(await m('http://yeoman.io', {