Add tests for connection.run

bigchaindb/backend/mongodb/connection.py

@@ -76,6 +76,8 @@ class MongoDBConnection(Connection):
             return query.run(self.conn)
         except pymongo.errors.DuplicateKeyError as exc:
             raise DuplicateKeyError from exc
+        except pymongo.errors.AutoReconnect as exc:
+            raise ConnectionError from exc
         except pymongo.errors.OperationFailure as exc:
             raise OperationError from exc
 

tests/backend/mongodb/test_connection.py

@@ -1,9 +1,9 @@
 from unittest import mock
 
 import pytest
+import pymongo
 from pymongo import MongoClient
 from pymongo.database import Database
-from pymongo.errors import ConnectionFailure, OperationFailure
 
 
 pytestmark = pytest.mark.bdb
@@ -61,7 +61,7 @@ def test_connection_error(mock_sleep, mock_client, mock_init_repl_set):
     # force the driver to throw ConnectionFailure
     # the mock on time.sleep is to prevent the actual sleep when running
     # the tests
-    mock_client.side_effect = ConnectionFailure()
+    mock_client.side_effect = pymongo.errors.ConnectionFailure()
 
     with pytest.raises(ConnectionError):
         conn = connect()
@@ -70,6 +70,30 @@ def test_connection_error(mock_sleep, mock_client, mock_init_repl_set):
     assert mock_client.call_count == 3
 
 
+@mock.patch('bigchaindb.backend.mongodb.connection.initialize_replica_set')
+@mock.patch('pymongo.MongoClient')
+def test_connection_run_errors(mock_client, mock_init_repl_set):
+    from bigchaindb.backend import connect
+    from bigchaindb.backend.exceptions import (DuplicateKeyError,
+                                               OperationError,
+                                               ConnectionError)
+
+    conn = connect()
+    query = mock.Mock()
+
+    query.run.side_effect = pymongo.errors.AutoReconnect('foo')
+    with pytest.raises(ConnectionError):
+        conn.run(query)
+
+    query.run.side_effect = pymongo.errors.DuplicateKeyError('foo')
+    with pytest.raises(DuplicateKeyError):
+        conn.run(query)
+
+    query.run.side_effect = pymongo.errors.OperationFailure('foo')
+    with pytest.raises(OperationError):
+        conn.run(query)
+
+
 def test_check_replica_set_not_enabled(mongodb_connection):
     from bigchaindb.backend.mongodb.connection import _check_replica_set
     from bigchaindb.common.exceptions import ConfigurationError
@@ -145,8 +169,8 @@ def test_initialize_replica_set(mock_cmd_line_opts):
     with mock.patch.object(Database, 'command') as mock_command:
         mock_command.side_effect = [
             mock_cmd_line_opts,
-            OperationFailure(None, details={'codeName': ''})
+            pymongo.errors.OperationFailure(None, details={'codeName': ''})
         ]
 
-        with pytest.raises(OperationFailure):
+        with pytest.raises(pymongo.errors.OperationFailure):
             initialize_replica_set('host', 1337, 1000)