Merge 4c80a0f into 6225f42

tests/test_auth.py

@@ -250,8 +250,8 @@ def test_Explicit_SCRAM_SHA_1(self):
     def test_InvalidArgs(self):
         conn = self.__get_connection()
         try:
-            yield self.assertFailure(conn[self.db1].authenticate(self.login1, 123), TypeError)
-            yield self.assertFailure(conn[self.db1].authenticate(123, self.password1), TypeError)
+            yield self.assertRaises(TypeError, conn[self.db1].authenticate, self.login1, 123)
+            yield self.assertRaises(TypeError, conn[self.db1].authenticate, 123, self.password1)
         finally:
             yield conn.disconnect()
 

tests/test_cancel.py

@@ -0,0 +1,161 @@
+import struct
+
+from mock import Mock
+from pymongo.uri_parser import parse_uri
+from twisted.internet import defer
+from twisted.internet import reactor
+from twisted.test import proto_helpers
+from twisted.trial import unittest
+from txmongo import MongoProtocol
+from txmongo import Query
+from txmongo.connection import _Connection, ConnectionPool
+from txmongo.protocol import Reply
+
+
+def _delay(time):
+    d = defer.Deferred()
+    reactor.callLater(time, d.callback, None)
+    return d
+
+
+class AssertCallbackNotCalled(object):
+    """
+    Context manager that assures Deferred's callback was not called
+    after it was cancelled. So we can be sure that Deferred's canceller
+    correctly removed it from waiting lists.
+    """
+    def __init__(self, deferred):
+        self.deferred = deferred
+
+    def __enter__(self):
+        self.deferred.callback = Mock()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.deferred.callback.assert_not_called()
+
+
+class TestCancelParts(unittest.TestCase):
+
+    def test_protocol_query(self):
+        tr = proto_helpers.StringTransport()
+        proto = MongoProtocol()
+        proto.makeConnection(tr)
+
+        d = proto.send_QUERY(Query(query={'x': 42}))
+        d.cancel()
+
+        with AssertCallbackNotCalled(d):
+            reply = Reply(response_to=1, documents=[{'x': 42}])
+            reply_bin = struct.pack("<iiiiqii", 1, *reply[2:8]) + b''.join(reply.documents)
+            reply_bin = struct.pack('<i', len(reply_bin) + 4) + reply_bin
+            proto.dataReceived(reply_bin)
+
+        self.failureResultOf(d, defer.CancelledError)
+
+    def test_protocol_connectionReady(self):
+        proto = MongoProtocol()
+        d1 = proto.connectionReady()
+        d2 = proto.connectionReady()
+        d1.cancel()
+
+        with AssertCallbackNotCalled(d1):
+            proto.makeConnection(proto_helpers.StringTransport())
+
+        self.assertTrue(d2.called)
+        self.failureResultOf(d1, defer.CancelledError)
+
+    @defer.inlineCallbacks
+    def test_connection_notifyReady(self):
+        uri = parse_uri("mongodb://localhost:27017/")
+        conn = _Connection(None, uri, 0, 10, 10, 10, 10)
+        d1 = conn.notifyReady()
+        d2 = conn.notifyReady()
+        d1.cancel()
+
+        connector = reactor.connectTCP("localhost", 27017, conn)
+
+        with AssertCallbackNotCalled(d1):
+            yield d2
+
+        self.failureResultOf(d1, defer.CancelledError)
+
+        conn.stopTrying()
+        conn.stopFactory()
+        conn.instance.transport.loseConnection()
+        connector.disconnect()
+
+
+class TestCancelIntegrated(unittest.TestCase):
+
+    def setUp(self):
+        self.conn = ConnectionPool()
+        self.db = self.conn.db
+        self.coll = self.db.coll
+
+    @defer.inlineCallbacks
+    def tearDown(self):
+        yield self.coll.drop()
+        yield self.conn.disconnect()
+
+    @defer.inlineCallbacks
+    def test_integration(self):
+        # Our ConnectionPool is not actually connected yet, so on this
+        # stage operations can be safely cancelled -- they won't be
+        # sent to MongoDB at all. This test checks this.
+
+        d1 = self.coll.insert_one({'x': 1})
+        d2 = self.coll.insert_one({'x': 2})
+        d3 = self.coll.insert_one({'x': 3})
+        d4 = self.coll.insert_one({'x': 4})
+
+        d1.cancel()
+        d3.cancel()
+
+        yield d4
+
+        self.failureResultOf(d1, defer.CancelledError)
+        self.assertTrue(d2.called)
+        self.failureResultOf(d3, defer.CancelledError)
+
+        docs = yield self.coll.distinct('x')
+        self.assertEqual(set(docs), {2, 4})
+
+    @defer.inlineCallbacks
+    def test_remove(self):
+        # Lets test cancellation of some dangerous operation for the peace
+        # of mind. NB: remove can be cancelled only because ConnectionPool
+        # is not connected yet.
+        for i in range(10):
+            self.coll.insert_one({'x': i})
+
+        d1 = self.coll.remove({'x': {"$lt": 3}})
+        d2 = self.coll.remove({'x': {"$gte": 3, "$lt": 6}})
+        d3 = self.coll.remove({'x': {"$gte": 6, "$lt": 9}})
+
+        d2.cancel()
+
+        yield d3
+
+        self.assertTrue(d1.called)
+        self.failureResultOf(d2, defer.CancelledError)
+
+        x = yield self.coll.distinct('x')
+        self.assertEqual(set(x), {3, 4, 5, 9})
+
+    @defer.inlineCallbacks
+    def test_no_way(self):
+        # If ConnectionPool picks already active connection, the query is sent
+        # to MongoDB immediately and there is no way to cancel it
+
+        yield self.coll.count()
+
+        d = self.coll.insert({'x': 42})
+        d.cancel()
+
+        yield _delay(1)
+
+        self.failureResultOf(d, defer.CancelledError)
+
+        cnt = yield self.coll.count()
+        self.assertEqual(cnt, 1)

tests/test_collection.py

@@ -67,17 +67,21 @@ def make_col(base, name):
         self.assertRaises(errors.InvalidName, make_col, self.db.test, "test.")
         self.assertRaises(errors.InvalidName, make_col, self.db.test, "tes..t")
         self.assertRaises(errors.InvalidName, make_col, self.db.test, "tes\x00t")
-        self.assertFailure(self.coll.save("test"), TypeError)
-        self.assertFailure(self.coll.filemd5("test"), ValueError)
-        self.assertFailure(self.db.test.find(spec="test"), TypeError)
-        self.assertFailure(self.db.test.find(fields="test"), TypeError)
-        self.assertFailure(self.db.test.find(skip="test"), TypeError)
-        self.assertFailure(self.db.test.find(limit="test"), TypeError)
-        self.assertFailure(self.db.test.insert([1]), TypeError)
-        self.assertFailure(self.db.test.insert(1), TypeError)
-        self.assertFailure(self.db.test.update(1, 1), TypeError)
-        self.assertFailure(self.db.test.update({}, 1), TypeError)
-        self.assertFailure(self.db.test.update({}, {}, 'a'), TypeError)
+        self.assertRaises(TypeError, self.coll.save, "test")
+        self.assertRaises(ValueError, self.coll.filemd5, "test")
+        #self.assertFailure(self.db.test.find(spec="test"), TypeError)
+        #self.assertFailure(self.db.test.find(fields="test"), TypeError)
+        #self.assertFailure(self.db.test.find(skip="test"), TypeError)
+        #self.assertFailure(self.db.test.find(limit="test"), TypeError)
+        self.assertRaises(TypeError, self.db.test.find, spec="test")
+        self.assertRaises(TypeError, self.db.test.find, fields="test")
+        self.assertRaises(TypeError, self.db.test.find, skip="test")
+        self.assertRaises(TypeError, self.db.test.find, limit="test")
+        self.assertRaises(TypeError, self.db.test.insert, [1])
+        self.assertRaises(TypeError, self.db.test.insert, 1)
+        self.assertRaises(TypeError, self.db.test.update, 1, 1)
+        self.assertRaises(TypeError, self.db.test.update, {}, 1)
+        self.assertRaises(TypeError, self.db.test.update, {}, {}, 'a')
 
         self.assert_(isinstance(self.db.test, Collection))
         self.assertEqual(NotImplemented, self.db.test.__cmp__(7))
@@ -100,8 +104,8 @@ def test_create_index(self):
         db = self.db
         coll = self.coll
 
-        self.assertFailure(coll.create_index(5), TypeError)
-        self.assertFailure(coll.create_index({"hello": 1}), TypeError)
+        self.assertRaises(TypeError, coll.create_index, 5)
+        self.assertRaises(TypeError, coll.create_index, {"hello": 1})
 
         yield coll.insert({'c': 1})  # make sure collection exists.
 
@@ -285,7 +289,7 @@ def test_drop_index(self):
         res = yield self.coll.drop_index(index)
         self.assertEqual(res["ok"], 1)
 
-        self.assertFailure(self.coll.drop_index(123), TypeError)
+        self.assertRaises(TypeError, self.coll.drop_index, 123)
 
 
 class TestRename(unittest.TestCase):

tests/test_connection.py

@@ -64,11 +64,9 @@ def test_Timeout_and_Deadline(self):
         yield self.named_conn.db.coll.insert({'x': 42}, safe=True, timeout=10)
         yield self.named_conn.db.coll.insert({'x': 42}, safe=True, deadline=time()+10)
 
-        d_insert = self.named_conn.db.coll.insert({'x': 42}, safe=True, deadline=time()-10)
-        yield self.assertFailure(d_insert, TimeExceeded)
+        self.assertRaises(TimeExceeded, self.named_conn.db.coll.insert, {'x': 42}, safe=True, deadline=time()-10)
 
-        d_insert = self.named_conn.db.coll.insert({'x': 42}, safe=True, timeout=-10)
-        yield self.assertFailure(d_insert, TimeExceeded)
+        self.assertRaises(TimeExceeded, self.named_conn.db.coll.insert, {'x': 42}, safe=True, timeout=-10)
 
         def patch_deadline(_):
             check_deadline(time()-2)

tests/test_find_and_modify.py

@@ -50,9 +50,8 @@ def test_Update(self):
         self.assertEqual(res["lulz"], 456)
 
     def test_InvalidOptions(self):
-        self.assertFailure(self.coll.find_and_modify(), ValueError)
-        self.assertFailure(self.coll.find_and_modify(update={"$set": {'x': 42}}, remove=True),
-                           ValueError)
+        self.assertRaises(ValueError, self.coll.find_and_modify)
+        self.assertRaises(ValueError, self.coll.find_and_modify, update={"$set": {'x': 42}}, remove=True)
 
     @defer.inlineCallbacks
     def test_Remove(self):

tests/test_objects.py

@@ -55,7 +55,7 @@ def test_MongoObjects(self):
         self.assertEqual(isinstance(mycol3, collection.Collection), True)
         yield mydb.drop_collection("mycol3")
         yield mydb.drop_collection(mycol3)
-        yield self.assertFailure(mydb.drop_collection(None), TypeError)
+        self.assertRaises(TypeError, mydb.drop_collection, None)
         yield conn.disconnect()
 
     @defer.inlineCallbacks
@@ -208,6 +208,36 @@ def test_GridFileObjects(self):
         grid_in_file.test = 1
         yield grid_in_file.write(b"0xDEADBEEF")
         yield grid_in_file.write(b"0xDEADBEEF"*1048576)
+        self.assertTrue("20150101", grid_in_file.upload_date)
+        yield grid_in_file.writelines([b"0xDEADBEEF", b"0xDEADBEEF"])
+        yield grid_in_file.close()
+        if _version.version.major >= 15:
+            with self.assertRaises(AttributeError):
+                grid_in_file.length = 1
+        self.assertEqual(1, grid_in_file.test)
+        if _version.version.major >= 15:
+            with self.assertRaises(AttributeError):
+                _ = grid_in_file.test_none
+        self.assertTrue(grid_in_file.closed)
+        if _version.version.major >= 15:
+            with self.assertRaises(ValueError):
+                yield grid_in_file.write(b"0xDEADBEEF")
+
+        doc = yield db.fs.files.find_one({"_id": grid_in_file._id})
+        grid_out_file = GridOut(db.fs, doc)
+        grid_out_file.seek(0, os.SEEK_END)
+        self.assertEqual(10 * (1 + 1048576 + 2), grid_out_file.tell())
+        yield grid_out_file.close()
+        data = b''
+        for block_dfr in GridOutIterator(grid_out_file, db.fs.chunks):
+            block = yield block_dfr
+            if block:
+                data += block
+            else:
+                break
+        self.assertEqual(data, b"0xDEADBEEF" * (1 + 1048576 + 2))
+
+
         fake_doc = {"_id": "test_id", "length": 1048576, "filename": "test",
                     "upload_date": "20150101"}
         self.assertRaises(TypeError, GridOut, None, None)
@@ -226,20 +256,6 @@ def test_GridFileObjects(self):
         self.assertRaises(IOError, grid_out_file.seek, 0, 4)
         self.assertRaises(IOError, grid_out_file.seek, -1)
         self.assertTrue("'_id': 'test_id'" in repr(grid_out_file))
-        self.assertTrue("20150101", grid_in_file.upload_date)
-        yield grid_in_file.writelines([b"0xDEADBEEF", b"0xDEADBEAF"])
-        yield grid_in_file.close()
-        if _version.version.major >= 15:
-            with self.assertRaises(AttributeError):
-                grid_in_file.length = 1
-        self.assertEqual(1, grid_in_file.test)
-        if _version.version.major >= 15:
-            with self.assertRaises(AttributeError):
-                _ = grid_in_file.test_none
-        self.assertTrue(grid_in_file.closed)
-        if _version.version.major >= 15:
-            with self.assertRaises(ValueError):
-                yield grid_in_file.write(b"0xDEADBEEF")
         yield conn.disconnect()
 
     @defer.inlineCallbacks