diff --git a/test/test_cursor.py b/test/test_cursor.py
index 08859277a..5fe605336 100644
--- a/test/test_cursor.py
+++ b/test/test_cursor.py
@@ -1213,26 +1213,23 @@ self.assertFalse(c2.alive)
         if server_started_with_auth(self.db.connection):
             raise SkipTest("SERVER-4754 - This test uses profiling.")
 
-        def run_with_profiling(func):
-            self.db.set_profiling_level(OFF)
-            self.db.system.profile.drop()
-            self.db.set_profiling_level(ALL)
-            func()
-            self.db.set_profiling_level(OFF)
-
-        def find():
-            list(self.db.test.find().comment('foo'))
-            op = self.db.system.profile.find({'ns': 'pymongo_test.test',
-                                              'op': 'query',
-                                              'query.$comment': 'foo'})
-            self.assertEqual(op.count(), 1)
-
-        run_with_profiling(find)
-
         # MongoDB 3.1.5 changed the ns for commands.
         regex = {'$regex': 'pymongo_test.(\$cmd|test)'}
 
-        def count():
+        if version.at_least(self.db.connection, (3, 1, 8, -1)):
+            query_key = "query.comment"
+        else:
+            query_key = "query.$comment"
+
+        self.client.drop_database(self.db)
+        self.db.set_profiling_level(ALL)
+        try:
+            list(self.db.test.find().comment('foo'))
+            op = self.db.system.profile.find({'ns': 'pymongo_test.test',
+                                              'op': 'query',
+                                              query_key: 'foo'})
+            self.assertEqual(op.count(), 1)
+
             self.db.test.find().comment('foo').count()
             op = self.db.system.profile.find({'ns': regex,
                                               'op': 'command',
@@ -1240,25 +1237,21 @@ self.assertFalse(c2.alive)
                                               'command.$comment': 'foo'})
             self.assertEqual(op.count(), 1)
 
-        run_with_profiling(count)
-
-        def distinct():
             self.db.test.find().comment('foo').distinct('type')
             op = self.db.system.profile.find({'ns': regex,
                                               'op': 'command',
                                               'command.distinct': 'test',
                                               'command.$comment': 'foo'})
             self.assertEqual(op.count(), 1)
-
-        run_with_profiling(distinct)
+        finally:
+            self.db.set_profiling_level(OFF)
+            self.db.system.profile.drop()
 
         self.db.test.insert([{}, {}])
         cursor = self.db.test.find()
         cursor.next()
         self.assertRaises(InvalidOperation, cursor.comment, 'hello')
 
-        self.db.system.profile.drop()
-
     def test_cursor_transfer(self):
 
         # This is just a test, don't try this at home...
diff --git a/test/test_database.py b/test/test_database.py
index 1de22c817..cf6983b89 100644
--- a/test/test_database.py
+++ b/test/test_database.py
@@ -237,6 +237,7 @@ class TestDatabase(unittest.TestCase):
             raise SkipTest('profile is not supported by mongos')
         db = self.client.pymongo_test
 
+        db.system.profile.drop()
         db.set_profiling_level(ALL)
         db.test.find_one()
         db.set_profiling_level(OFF)