From df99ae56b09733be8e9723cd9f8163d66f149c69 Mon Sep 17 00:00:00 2001
From: Christopher Chedeau <vjeuxx@gmail.com>
Date: Fri, 13 Jan 2017 10:39:07 -0800
Subject: [PATCH] Add test for shebang and move to index.js (#170)

When looking into adding a test, I realized that the logic was inside of bin/prettier.js and therefore only applying to the cli. Moving it to index.js and adding a test so that it's more robust :)
---
 bin/prettier.js                               |  9 +--
 index.js                                      | 69 +++++++++++--------
 .../shebang/__snapshots__/jsfmt.spec.js.snap  |  8 +++
 tests/shebang/jsfmt.spec.js                   |  1 +
 tests/shebang/shebang.js                      |  2 +
 5 files changed, 54 insertions(+), 35 deletions(-)
 create mode 100644 tests/shebang/__snapshots__/jsfmt.spec.js.snap
 create mode 100644 tests/shebang/jsfmt.spec.js
 create mode 100644 tests/shebang/shebang.js

diff --git a/bin/prettier.js b/bin/prettier.js
index f034c7e3c..a5b6229f0 100755
--- a/bin/prettier.js
+++ b/bin/prettier.js
@@ -33,13 +33,6 @@ if (!filenames.length && !stdin) {
   process.exit(1);
 }
 
-function formatWithShebang(input) {
-  const index = input.indexOf("\n");
-  const shebang = input.slice(0, index + 1);
-  const programInput = input.slice(index + 1);
-  return shebang + format(programInput);
-}
-
 function format(input) {
   return jscodefmt.format(input, {
     printWidth: argv["print-width"],
@@ -77,7 +70,7 @@ if (stdin) {
 
       let output;
       try {
-        output = input.startsWith("#!") ? formatWithShebang(input) : format(input);
+        output = format(input);
       } catch (e) {
         process.exitCode = 2;
         console.error(e);
diff --git a/index.js b/index.js
index 082babfb6..e37fcdf40 100644
--- a/index.js
+++ b/index.js
@@ -23,34 +23,49 @@ var babylonOptions = {
   ]
 };
 
+function format(text, opts) {
+  opts = opts || {};
+  let ast;
+
+  if(opts.useFlowParser) {
+    ast = flowParser.parse(text);
+    if(ast.errors.length > 0) {
+      let msg = ast.errors[0].message + " on line " + ast.errors[0].loc.start.line
+      if(opts.filename) {
+        msg += " in file " + opts.filename;
+      }
+      throw new Error(msg);
+    }
+  }
+  else {
+    ast = babylon.parse(text, babylonOptions);
+  }
+
+  // Interleave comment nodes
+  if(ast.comments) {
+    comments.attach(ast.comments, ast, text);
+    ast.comments = [];
+  }
+  ast.tokens = [];
+  opts.originalText = text;
+
+  const printer = new Printer(opts);
+  return printer.printGenerically(ast).code;
+}
+
+function formatWithShebang(text, opts) {
+  if (!text.startsWith("#!")) {
+    return format(text, opts);
+  }
+
+  const index = text.indexOf("\n");
+  const shebang = text.slice(0, index + 1);
+  const programText = text.slice(index + 1);
+  return shebang + format(programText, opts);
+}
+
 module.exports = {
   format: function(text, opts) {
-    opts = opts || {};
-    let ast;
-
-    if(opts.useFlowParser) {
-      ast = flowParser.parse(text);
-      if(ast.errors.length > 0) {
-        let msg = ast.errors[0].message + " on line " + ast.errors[0].loc.start.line
-        if(opts.filename) {
-          msg += " in file " + opts.filename;
-        }
-        throw new Error(msg);
-      }
-    }
-    else {
-      ast = babylon.parse(text, babylonOptions);
-    }
-
-    // Interleave comment nodes
-    if(ast.comments) {
-      comments.attach(ast.comments, ast, text);
-      ast.comments = [];
-    }
-    ast.tokens = [];
-    opts.originalText = text;
-
-    const printer = new Printer(opts);
-    return printer.printGenerically(ast).code;
+    return formatWithShebang(text, opts);
   }
 };
diff --git a/tests/shebang/__snapshots__/jsfmt.spec.js.snap b/tests/shebang/__snapshots__/jsfmt.spec.js.snap
new file mode 100644
index 000000000..ab63719bc
--- /dev/null
+++ b/tests/shebang/__snapshots__/jsfmt.spec.js.snap
@@ -0,0 +1,8 @@
+exports[`test shebang.js 1`] = `
+"#!/usr/bin/env node
+function a() {}
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#!/usr/bin/env node
+function a() {}
+"
+`;
diff --git a/tests/shebang/jsfmt.spec.js b/tests/shebang/jsfmt.spec.js
new file mode 100644
index 000000000..989047bcc
--- /dev/null
+++ b/tests/shebang/jsfmt.spec.js
@@ -0,0 +1 @@
+run_spec(__dirname);
diff --git a/tests/shebang/shebang.js b/tests/shebang/shebang.js
new file mode 100644
index 000000000..552d6b4c7
--- /dev/null
+++ b/tests/shebang/shebang.js
@@ -0,0 +1,2 @@
+#!/usr/bin/env node
+function a() {}