diff --git a/pillar/api/nodes/moving.py b/pillar/api/nodes/moving.py index 1496308e..ee18fe6c 100644 --- a/pillar/api/nodes/moving.py +++ b/pillar/api/nodes/moving.py @@ -1,7 +1,7 @@ """Code for moving around nodes.""" import attr -import flask_pymongo.wrappers +import pymongo.database from bson import ObjectId from pillar import attrs_extra @@ -10,7 +10,7 @@ import pillar.api.file_storage.moving @attr.s class NodeMover(object): - db = attr.ib(validator=attr.validators.instance_of(flask_pymongo.wrappers.Database)) + db = attr.ib(validator=attr.validators.instance_of(pymongo.database.Database)) skip_gcs = attr.ib(default=False, validator=attr.validators.instance_of(bool)) _log = attrs_extra.log('%s.NodeMover' % __name__) diff --git a/tests/test_api/test_nodes_moving.py b/tests/test_api/test_nodes_moving.py index 30518e56..230d0195 100644 --- a/tests/test_api/test_nodes_moving.py +++ b/tests/test_api/test_nodes_moving.py @@ -2,7 +2,7 @@ import unittest import mock from bson import ObjectId -import flask_pymongo.wrappers +import pymongo.database # Required to mock this module, otherwise 'pillar.api.file_storage' doesn't have # the attribute 'moving' @@ -13,7 +13,7 @@ class NodeMoverTest(unittest.TestCase): def setUp(self): from pillar.api.nodes import moving - self.db = mock.MagicMock(spec=flask_pymongo.wrappers.Database) + self.db = mock.MagicMock(spec=pymongo.database.Database) self.mover = moving.NodeMover(db=self.db) def test_file_generator(self):