From f347407c5b71093dece0312d357ad1d5f48de910 Mon Sep 17 00:00:00 2001
From: Noah Levitt <nlevitt@archive.org>
Date: Mon, 17 Sep 2018 13:24:59 -0700
Subject: [PATCH] making tests work better and pass

but not in python 2.7. mock.MagicMock is not an iterator there
apparently :(
---
 .travis.yml             |  3 ++-
 setup.py                |  9 +++++++++
 tests/test_rethinker.py | 33 ++++++++++++++++++++-------------
 3 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index ab5e580..cc6dfec 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -16,6 +16,7 @@ matrix:
   allow_failures:
   - python: nightly
   - python: 3.7-dev
+  - python: 2.7
 
 services:
 - docker
@@ -25,7 +26,7 @@ before_install:
 - docker run -d --publish=28015:28015 rethinkdb
 
 install:
-- pip install . pytest
+- pip install .[test]
 
 script:
 - py.test -v tests
diff --git a/setup.py b/setup.py
index 56a44b4..eb412a5 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,13 @@
 import setuptools
 import codecs
 
+test_deps = ['pytest']
+
+try:
+    import unittest.mock
+except:
+    test_deps.append('mock')
+
 setuptools.setup(
     name='doublethink',
     version='0.2.0.dev88',
@@ -10,8 +17,10 @@ setuptools.setup(
         'Programming Language :: Python :: 3.4',
         'Programming Language :: Python :: 3.5',
         'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
     ],
     install_requires=['rethinkdb'],
+    extras_require={'test': test_deps},
     url='https://github.com/internetarchive/doublethink',
     author='Noah Levitt',
     author_email='nlevitt@archive.org',
diff --git a/tests/test_rethinker.py b/tests/test_rethinker.py
index 13865ae..4c7f1b4 100644
--- a/tests/test_rethinker.py
+++ b/tests/test_rethinker.py
@@ -24,7 +24,10 @@ import gc
 import pytest
 import rethinkdb as r
 import datetime
-from unittest import mock
+try:
+    from unittest import mock
+except:
+    import mock
 
 logging.basicConfig(stream=sys.stderr, level=logging.INFO,
         format="%(asctime)s %(process)d %(levelname)s %(threadName)s %(name)s.%(funcName)s(%(filename)s:%(lineno)d) %(message)s")
@@ -129,6 +132,9 @@ def test_utcnow():
 
     ## XXX what else can we test without jumping through hoops?
 
+class SpecificException(Exception):
+    pass
+
 class MockRethinker(doublethink.Rethinker):
     def __init__(self, *args, **kwargs):
         self.m = mock.MagicMock()
@@ -141,13 +147,13 @@ class MockRethinker(doublethink.Rethinker):
                             'Cannot perform read: The primary replica '
                             "isn't connected... THIS IS A TEST!")
                 else:
-                    e = Exception
+                    e = SpecificException
 
-                count = 0
+                # dict because: https://stackoverflow.com/questions/3190706/nonlocal-keyword-in-python-2-x
+                count = {'value': 0}
                 def run(*args, **kwargs):
-                    nonlocal count
-                    count += 1
-                    if count <= 2:
+                    count['value'] += 1
+                    if count['value'] <= 2:
                         raise e
                     else:
                         return mock.MagicMock()
@@ -160,20 +166,21 @@ class MockRethinker(doublethink.Rethinker):
                             'Cannot perform read: The primary replica '
                             "isn't connected... THIS IS A TEST!")
                 else:
-                    e = Exception
+                    e = SpecificException
 
                 def run(*args, **kwargs):
                     mmm = mock.MagicMock()
-                    count = 0
+                    # dict because: https://stackoverflow.com/questions/3190706/nonlocal-keyword-in-python-2-x
+                    count = {'value': 0}
                     def next_(*args, **kwargs):
-                        nonlocal count
-                        count += 1
-                        if count <= 2:
+                        count['value'] += 1
+                        if count['value'] <= 2:
                             raise e
                         else:
                             return mock.MagicMock()
                     mmm.__iter__ = lambda *args, **kwargs: mmm
                     mmm.__next__ = next_
+                    mmm.next = next_
                     return mmm
 
                 mm.run = run
@@ -191,14 +198,14 @@ class MockRethinker(doublethink.Rethinker):
 def test_error_handling():
     rr = MockRethinker(db='my_db')
 
-    with pytest.raises(Exception):
+    with pytest.raises(SpecificException):
         rr.table('err_running_query').run()
 
     # should not raise exception
     rr.table('recoverable_err_running_query').run()
 
     it = rr.table('err_in_iterator').run() # no exception yet
-    with pytest.raises(Exception):
+    with pytest.raises(SpecificException):
         next(it)  # exception here
 
     it = rr.table('recoverable_err_in_iterator').run() # no exception yet