From 22caa27a7297082b95a16afc797246eb93ffb466 Mon Sep 17 00:00:00 2001
From: Bernie Hackett <bernie@10gen.com>
Date: Tue, 4 Mar 2014 10:35:45 -0800
Subject: [PATCH] Add more tests for maxTimeMS PYTHON-550

---
 test/test_cursor.py   | 23 ++++++++++++++---------
 test/test_database.py | 34 ++++++++++++++++++++++++++++++++--
 2 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/test/test_cursor.py b/test/test_cursor.py
index 824be1cbc..431c9ae94 100644
--- a/test/test_cursor.py
+++ b/test/test_cursor.py
@@ -102,19 +102,24 @@ class TestCursor(unittest.TestCase):
 
         self.assertTrue(coll.find_one(max_time_ms=1000))
 
-        reducer = Code("""function(obj, prev){prev.count++;}""")
-        coll.group(key={"amalia": 1}, condition={}, initial={"count": 0},
-                   reduce=reducer, maxTimeMS=1000)
-
         if "enableTestCommands=1" in get_command_line(self.client)["argv"]:
             self.client.admin.command("configureFailPoint",
                                       "maxTimeAlwaysTimeOut",
                                       mode="alwaysOn")
-            self.assertRaises(ExecutionTimeout,
-                              coll.find_one, max_time_ms=1)
-            self.client.admin.command("configureFailPoint",
-                                      "maxTimeAlwaysTimeOut",
-                                      mode="off")
+            try:
+                cursor = coll.find().max_time_ms(1)
+                try:
+                    cursor.next()
+                except ExecutionTimeout:
+                    pass
+                else:
+                    self.fail("ExecutionTimeout not raised")
+                self.assertRaises(ExecutionTimeout,
+                                  coll.find_one, max_time_ms=1)
+            finally:
+                self.client.admin.command("configureFailPoint",
+                                          "maxTimeAlwaysTimeOut",
+                                          mode="off")
 
     def test_explain(self):
         a = self.db.test.find()
diff --git a/test/test_database.py b/test/test_database.py
index 6f3a7bc9e..b2e0efbb7 100644
--- a/test/test_database.py
+++ b/test/test_database.py
@@ -40,14 +40,15 @@ from pymongo.collection import Collection
 from pymongo.database import Database
 from pymongo.errors import (CollectionInvalid,
                             ConfigurationError,
+                            ExecutionTimeout,
                             InvalidName,
                             OperationFailure)
 from pymongo.son_manipulator import (AutoReference,
                                      NamespaceInjector,
                                      ObjectIdShuffler)
 from test import version
-from test.utils import (is_mongos, server_started_with_auth,
-                        remove_all_users)
+from test.utils import (get_command_line, is_mongos,
+                        remove_all_users, server_started_with_auth)
 from test.test_client import get_client
 
 
@@ -959,6 +960,35 @@ class TestDatabase(unittest.TestCase):
             warnings.resetwarnings()
             warnings.simplefilter("ignore")
 
+    def test_command_max_time_ms(self):
+        if not version.at_least(self.client, (2, 5, 3, -1)):
+            raise SkipTest("MaxTimeMS requires MongoDB >= 2.5.3")
+        if "enableTestCommands=1" not in get_command_line(self.client)["argv"]:
+            raise SkipTest("Test commands must be enabled.")
+
+        self.client.admin.command("configureFailPoint",
+                                  "maxTimeAlwaysTimeOut",
+                                  mode="alwaysOn")
+        try:
+            db = self.client.pymongo_test
+            db.command('count', 'test')
+            self.assertRaises(ExecutionTimeout, db.command,
+                              'count', 'test', maxTimeMS=1)
+            pipeline = [{'$project': {'name': 1, 'count': 1}}]
+            # Database command helper.
+            db.command('aggregate', 'test', pipeline=pipeline)
+            self.assertRaises(ExecutionTimeout, db.command,
+                              'aggregate', 'test',
+                              pipeline=pipeline, maxTimeMS=1)
+            # Collection helper.
+            db.test.aggregate(pipeline=pipeline)
+            self.assertRaises(ExecutionTimeout,
+                              db.test.aggregate, pipeline, maxTimeMS=1)
+        finally:
+            self.client.admin.command("configureFailPoint",
+                                      "maxTimeAlwaysTimeOut",
+                                      mode="off")
+
 
 if __name__ == "__main__":
     unittest.main()