From 48ee57cfb2b3d366dd2efb8e32ea538e44762ff0 Mon Sep 17 00:00:00 2001
From: Jeff Dickey <216188+jdxcode@users.noreply.github.com>
Date: Sat, 27 Jan 2018 00:03:16 -0800
Subject: [PATCH] fix: fixed retries

---
 README.md            | 23 +++++++++++++++++++++--
 src/base.ts          | 10 ++++++++--
 src/types.ts         |  3 ++-
 test/retries.test.ts | 17 +++++++++++++++++
 4 files changed, 48 insertions(+), 5 deletions(-)
 create mode 100644 test/retries.test.ts

diff --git a/README.md b/README.md
index 2120edf..7cf1637 100644
--- a/README.md
+++ b/README.md
@@ -24,6 +24,8 @@ extendable utilities for testing
   * [Run](#run)
   * [Add](#add)
   * [Stdout/Stderr Mocking](#stdoutstderr-mocking)
+  * [Done](#done)
+  * [Retries](#retries)
   * [Chai](#chai)
 - [Chaining](#chaining)
 - [Custom Plugins](#custom-plugins)
@@ -287,8 +289,6 @@ Done
 You can get the mocha `done()` callback by passing in a second argument.
 
 ```js
-import {expect, fancy} from 'fancy-test'
-
 describe('calls done', () => {
   fancy
   .it('expects FOO=bar', (_, done) => {
@@ -297,6 +297,25 @@ describe('calls done', () => {
 })
 ```
 
+Retries
+-------
+
+Retry the test n times.
+
+```js
+let count = 3
+
+describe('test retries', () => {
+  fancy
+  .retries(2)
+  .do(() => {
+    count--
+    if (count > 0) throw new Error('x')
+  })
+  .it('retries 3 times')
+})
+```
+
 Chai
 ----
 
diff --git a/src/base.ts b/src/base.ts
index a3bb808..4437f90 100644
--- a/src/base.ts
+++ b/src/base.ts
@@ -21,6 +21,9 @@ const base = <I extends Types.Context>(context: I): Types.Base<I, {}> => {
     }
     if (!arg1) arg1 = context.expectation || 'test'
     async function run(this: mocha.ITestCallbackContext, done?: Types.MochaDone) {
+      // reset error if retrying
+      delete context.error
+      if (context.retries) this.retries(context.retries)
       if (cb) {
         context.chain = [...context.chain, {
           run: async (input: any) => {
@@ -94,12 +97,12 @@ const base = <I extends Types.Context>(context: I): Types.Base<I, {}> => {
         chain: [...context.chain, {finally: (input: any) => cb(input)}]
       })
     },
-    add(key, cb) {
+    add(key, v) {
       return base({
         ...context as any,
         chain: [...context.chain, {
           run: async (ctx: any) => {
-            ctx[key] = await cb(ctx)
+            ctx[key] = await (_.isFunction(v) ? v(ctx) : v)
           }
         }]
       })
@@ -135,3 +138,6 @@ export default base(context)
 .register('only', () => ({
   init: ctx => {ctx.test = it.only}
 }))
+.register('retries', (count: number) => ({
+  init: ctx => ctx.retries = count
+}))
diff --git a/src/types.ts b/src/types.ts
index c60bb10..092f909 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -8,6 +8,7 @@ export interface Context {
   expectation?: string
   chain: Plugin<any>[]
   error?: Error & {code?: string}
+  retries?: number
 }
 
 export interface Plugin<I> {
@@ -36,7 +37,7 @@ export interface It<I> {
 export type Base<I extends Context, T extends Plugins> = {
   it: It<I>
   end: It<I>
-  add<K extends string, O>(key: K, cb: (context: I) => Promise<O> | O): Base<I & {[P in K]: O}, T>
+  add<K extends string, O>(key: K, cb: ((context: I) => Promise<O> | O) | Promise<O> | O): Base<I & {[P in K]: O}, T>
   do<O>(cb: (context: I & O) => any): Base<O & I, T>
   finally(cb: (context: I) => any): Base<I, T>
   register<K extends string, O, A1, A2, A3, A4>(key: K, plugin: (arg1?: A1, arg2?: A2, arg3?: A3, arg4?: A4) => Plugin<O & I>): Base<I, T & {[P in K]: {output: O, a1: A1, a2: A2, a3: A3, a4: A4}}>
diff --git a/test/retries.test.ts b/test/retries.test.ts
new file mode 100644
index 0000000..f43d416
--- /dev/null
+++ b/test/retries.test.ts
@@ -0,0 +1,17 @@
+// tslint:disable no-console
+
+import {fancy} from '../src'
+
+let count = 3
+
+describe('retries', () => {
+  // from readme
+  fancy
+  .retries(2)
+  .do(() => {
+    count--
+    if (count > 0) throw new Error('x')
+  })
+  .it('retries 3 times')
+  // from readme
+})