From e0f059378b113dcc40e7de7fb16b4f4a104a74f1 Mon Sep 17 00:00:00 2001
From: Kurt Griffiths <kgriffs@users.noreply.github.com>
Date: Thu, 25 Aug 2016 16:28:32 -0600
Subject: [PATCH] chore(bench.bench): allow benchmarking apps implemented as
 generators (#876)

Change create_bench to support WSGI applications implemented as Python
generators (i.e. that yield strings or bytes).
---
 falcon/bench/bench.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/falcon/bench/bench.py b/falcon/bench/bench.py
index 949916d..f534e2a 100755
--- a/falcon/bench/bench.py
+++ b/falcon/bench/bench.py
@@ -18,9 +18,10 @@
 from __future__ import print_function
 
 import argparse
-from collections import defaultdict
+from collections import defaultdict, deque
 from decimal import Decimal
 import gc
+import inspect
 import random
 import sys
 import timeit
@@ -99,6 +100,11 @@ def profile(name, env, filename=None, verbose=False):
                         sort='tottime', filename=filename)
 
 
+def exhaust(iterator_or_generator):
+    # from https://docs.python.org/dev/library/itertools.html#itertools-recipes
+    deque(iterator_or_generator, maxlen=0)
+
+
 BODY = helpers.rand_string(10240, 10240)  # NOQA
 HEADERS = {'X-Test': 'Funky Chicken'}  # NOQA
 
@@ -114,7 +120,15 @@ def create_bench(name, env):
         if srmock.status != '200 OK':
             raise AssertionError(srmock.status + ' != 200 OK')
 
-    return bench
+    def bench_generator():
+        exhaust(app(env, srmock))
+        if srmock.status != '200 OK':
+            raise AssertionError(srmock.status + ' != 200 OK')
+
+    if inspect.isgeneratorfunction(app):
+        return bench_generator
+    else:
+        return bench
 
 
 def consolidate_datasets(datasets):