From ccfe06e6ccf9ccee232fa60d5e83e1fa64cf219b 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 | 10 +++++++---
 test.js  | 18 +++++++++++++++++-
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/index.js b/index.js
index 09cb140..1371a89 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,11 @@ 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;
 };
diff --git a/test.js b/test.js
index 2e876d3..9159197 100644
--- a/test.js
+++ b/test.js
@@ -5,8 +5,24 @@ import isJpg from 'is-jpg';
 import isPng from 'is-png';
 import pify from 'pify';
 import PNG from 'png-js';
+import puppeteer from 'puppeteer';
 import server from './fixtures/server';
-import m from '.';
+import screenshotStream from '.';
+
+let browser;
+let m;
+
+test.before(async () => {
+	browser = await puppeteer.launch();
+	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', {