From cc2e83f07d1fbfe5c3e79e4e69d3ed670e114ead Mon Sep 17 00:00:00 2001
From: Rodolphe Marques <marques.rodolphe@gmail.com>
Date: Thu, 1 Jun 2017 14:30:20 +0200
Subject: [PATCH] Added unit test for rethinkdb connection timeout

- Cosmetic changes
---
 bigchaindb/backend/rethinkdb/connection.py | 26 ++++++----------------
 tests/backend/rethinkdb/test_connection.py | 19 +++++++++++++++-
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/bigchaindb/backend/rethinkdb/connection.py b/bigchaindb/backend/rethinkdb/connection.py
index 3ad779e2..3b2459ac 100644
--- a/bigchaindb/backend/rethinkdb/connection.py
+++ b/bigchaindb/backend/rethinkdb/connection.py
@@ -1,7 +1,4 @@
 import rethinkdb as r
-import time
-import bigchaindb
-
 from bigchaindb.backend.connection import Connection
 from bigchaindb.backend.exceptions import ConnectionError, OperationError
 
@@ -40,20 +37,11 @@ class RethinkDBConnection(Connection):
             :exc:`rethinkdb.ReqlDriverError`: After
                 :attr:`~.RethinkDBConnection.max_tries`.
         """
-        # NOTE: fix for timeout not initialized on rethinkdb connection
-        #       https://github.com/bigchaindb/bigchaindb/issues/1337
 
-        connected = False
-        dbconf = bigchaindb.config['database']
-        timeout = dbconf['connection_timeout']
-        end_time = time.time()*1000 + timeout
-        while not connected:
-            try:
-                rconn = r.connect(host=self.host, port=self.port, db=self.dbname)
-                connected = True
-            except (r.ReqlDriverError, r.ReqlTimeoutError) as exc:
-                if str(exc) == 'mock' or time.time()*1000 > end_time:
-                    raise ConnectionError from exc
-                time.sleep(timeout/1000)
-                pass
-        return rconn
+        try:
+            return r.connect(host=self.host,
+                             port=self.port,
+                             db=self.dbname,
+                             timeout=self.connection_timeout)
+        except (r.ReqlDriverError, r.ReqlTimeoutError) as exc:
+            raise ConnectionError from exc
diff --git a/tests/backend/rethinkdb/test_connection.py b/tests/backend/rethinkdb/test_connection.py
index df393716..c9703a6b 100644
--- a/tests/backend/rethinkdb/test_connection.py
+++ b/tests/backend/rethinkdb/test_connection.py
@@ -137,11 +137,28 @@ def test_changefeed_reconnects_when_connection_lost(monkeypatch):
 
 @patch('rethinkdb.connect')
 def test_connection_happens_one_time_if_successful(mock_connect):
+    import bigchaindb
     from bigchaindb.backend import connect
 
+    timeout = bigchaindb.config['database']['connection_timeout']
     query = r.expr('1')
     conn = connect('rethinkdb', 'localhost', 1337, 'whatev')
     conn.run(query)
     mock_connect.assert_called_once_with(host='localhost',
                                          port=1337,
-                                         db='whatev')
+                                         db='whatev',
+                                         timeout=timeout)
+
+
+@patch('rethinkdb.connect', side_effect=r.ReqlTimeoutError())
+def test_connection_timeout(mock_connect):
+    from bigchaindb.backend import connect
+    from bigchaindb.backend.exceptions import ConnectionError
+
+    query = r.expr('1')
+    conn = connect()
+
+    # connection should raise a ConnectionError after 3 tries
+    with pytest.raises(ConnectionError):
+        conn.run(query)
+    assert mock_connect.call_count == 3