diff --git a/bigchaindb/backend/mongodb/connection.py b/bigchaindb/backend/mongodb/connection.py index ce0c6113..4bd822b4 100644 --- a/bigchaindb/backend/mongodb/connection.py +++ b/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 diff --git a/tests/backend/mongodb/test_connection.py b/tests/backend/mongodb/test_connection.py index feda77ed..be660d0a 100644 --- a/tests/backend/mongodb/test_connection.py +++ b/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)