From 016ca5eba672058d70bf9bff6df64653ac29f3cd Mon Sep 17 00:00:00 2001
From: Jeff Dickey <216188+jdxcode@users.noreply.github.com>
Date: Fri, 19 Jan 2018 23:29:19 -0800
Subject: [PATCH] feat: added testCommand callback

---
 src/command.ts       | 28 +++++++++++++++++++++++-----
 test/command.test.ts |  5 ++++-
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/command.ts b/src/command.ts
index 70fe8b8b..d1dc8edf 100644
--- a/src/command.ts
+++ b/src/command.ts
@@ -4,13 +4,26 @@ import {expect, it, output} from '.'
 
 export interface TestCommandOptions {
   description?: string
-  stdout?: string
-  stderr?: string
+  stdout?: string | boolean
+  stderr?: string | boolean
   exit?: number
   root?: string
 }
 
-export const testCommand = (args: string[], opts: TestCommandOptions) => {
+export type TestCommandCallback<T> = (output: T) => Promise<void> | void
+
+export interface TestCommand {
+  (args: string[], opts: TestCommandOptions & {stdout: true, stderr: true}, fn: TestCommandCallback<{stdout: string, stderr: string}>): void
+  (args: string[], opts: TestCommandOptions & {stdout: true}, fn: TestCommandCallback<{stdout: string}>): void
+  (args: string[], opts: TestCommandOptions & {stderr: true}, fn: TestCommandCallback<{stderr: string}>): void
+  (args: string[], opts: TestCommandOptions, fn?: TestCommandCallback<{}>): void
+}
+
+export const testCommand: TestCommand = (
+  args: string[],
+  opts: TestCommandOptions,
+  fn?: TestCommandCallback<any>
+) => {
   const description = opts.description || args[0]
   let test = it
   if (opts.stdout) test = test.stdout
@@ -25,7 +38,12 @@ export const testCommand = (args: string[], opts: TestCommandOptions) => {
         throw new Error(`Expected exit code to be ${exit} but got ${err['cli-ux'].exit}`)
       }
     }
-    if (opts.stdout) expect(output.stdout).to.equal(opts.stdout)
-    if (opts.stderr) expect(output.stderr).to.equal(opts.stderr)
+    if (typeof opts.stdout === 'string') expect(output.stdout).to.equal(opts.stdout)
+    if (typeof opts.stderr === 'string') expect(output.stderr).to.equal(opts.stderr)
+    if (!fn) return
+    let o: any = {}
+    if (opts.stdout) o.stdout = output.stdout
+    if (opts.stderr) o.stderr = output.stderr
+    await fn(o)
   })
 }
diff --git a/test/command.test.ts b/test/command.test.ts
index 304d16be..f0377dd3 100644
--- a/test/command.test.ts
+++ b/test/command.test.ts
@@ -1,8 +1,11 @@
 import * as path from 'path'
 
-import {describe, testCommand} from '../src'
+import {describe, expect, testCommand} from '../src'
 
 describe('command', () => {
   testCommand(['foo:bar'], {stdout: 'hello world!\n', root: path.join(__dirname, 'fixtures/multi')})
   testCommand(['foo:bar', '--name=foo'], {stdout: 'hello foo!\n', root: path.join(__dirname, 'fixtures/multi')})
+  testCommand(['foo:bar', '--name=foo'], {stdout: true, root: path.join(__dirname, 'fixtures/multi')}, ({stdout}) => {
+    expect(stdout).to.equal('hello foo!\n')
+  })
 })