Convert test.test_gridfs_bucket to async

Author: sleepyStick

test/test_gridfs_bucket.py

@@ -16,12 +16,14 @@
 """Tests for the gridfs package."""
 from __future__ import annotations
 
+import asyncio
 import datetime
 import itertools
 import sys
 import threading
 import time
 from io import BytesIO
+from test.helpers import ConcurrentRunner
 from unittest.mock import patch
 
 sys.path[0:0] = [""]
@@ -44,10 +46,12 @@
 from pymongo.read_preferences import ReadPreference
 from pymongo.synchronous.mongo_client import MongoClient
 
+_IS_SYNC = True
 
-class JustWrite(threading.Thread):
+
+class JustWrite(ConcurrentRunner):
     def __init__(self, gfs, num):
-        threading.Thread.__init__(self)
+        super().__init__()
         self.gfs = gfs
         self.num = num
         self.daemon = True
@@ -59,9 +63,9 @@ def run(self):
             file.close()
 
 
-class JustRead(threading.Thread):
+class JustRead(ConcurrentRunner):
     def __init__(self, gfs, num, results):
-        threading.Thread.__init__(self)
+        super().__init__()
         self.gfs = gfs
         self.num = num
         self.results = results
@@ -89,12 +93,13 @@ def setUp(self):
 
     def test_basic(self):
         oid = self.fs.upload_from_stream("test_filename", b"hello world")
-        self.assertEqual(b"hello world", self.fs.open_download_stream(oid).read())
+        self.assertEqual(b"hello world", (self.fs.open_download_stream(oid)).read())
         self.assertEqual(1, self.db.fs.files.count_documents({}))
         self.assertEqual(1, self.db.fs.chunks.count_documents({}))
 
         self.fs.delete(oid)
-        self.assertRaises(NoFile, self.fs.open_download_stream, oid)
+        with self.assertRaises(NoFile):
+            self.fs.open_download_stream(oid)
         self.assertEqual(0, self.db.fs.files.count_documents({}))
         self.assertEqual(0, self.db.fs.chunks.count_documents({}))
 
@@ -111,7 +116,7 @@ def test_multi_chunk_delete(self):
 
     def test_empty_file(self):
         oid = self.fs.upload_from_stream("test_filename", b"")
-        self.assertEqual(b"", self.fs.open_download_stream(oid).read())
+        self.assertEqual(b"", (self.fs.open_download_stream(oid)).read())
         self.assertEqual(1, self.db.fs.files.count_documents({}))
         self.assertEqual(0, self.db.fs.chunks.count_documents({}))
 
@@ -128,10 +133,12 @@ def test_corrupt_chunk(self):
         self.db.fs.chunks.update_one({"files_id": files_id}, {"$set": {"data": Binary(b"foo", 0)}})
         try:
             out = self.fs.open_download_stream(files_id)
-            self.assertRaises(CorruptGridFile, out.read)
+            with self.assertRaises(CorruptGridFile):
+                out.read()
 
             out = self.fs.open_download_stream(files_id)
-            self.assertRaises(CorruptGridFile, out.readline)
+            with self.assertRaises(CorruptGridFile):
+                out.readline()
         finally:
             self.fs.delete(files_id)
 
@@ -146,13 +153,13 @@ def test_upload_ensures_index(self):
         self.assertTrue(
             any(
                 info.get("key") == [("files_id", 1), ("n", 1)]
-                for info in chunks.index_information().values()
+                for info in (chunks.index_information()).values()
             )
         )
         self.assertTrue(
             any(
                 info.get("key") == [("filename", 1), ("uploadDate", 1)]
-                for info in files.index_information().values()
+                for info in (files.index_information()).values()
             )
         )
 
@@ -174,33 +181,35 @@ def test_ensure_index_shell_compat(self):
             self.assertTrue(
                 any(
                     info.get("key") == [("filename", 1), ("uploadDate", 1)]
-                    for info in files.index_information().values()
+                    for info in (files.index_information()).values()
                 )
             )
             files.drop()
 
     def test_alt_collection(self):
         oid = self.alt.upload_from_stream("test_filename", b"hello world")
-        self.assertEqual(b"hello world", self.alt.open_download_stream(oid).read())
+        self.assertEqual(b"hello world", (self.alt.open_download_stream(oid)).read())
         self.assertEqual(1, self.db.alt.files.count_documents({}))
         self.assertEqual(1, self.db.alt.chunks.count_documents({}))
 
         self.alt.delete(oid)
-        self.assertRaises(NoFile, self.alt.open_download_stream, oid)
+        with self.assertRaises(NoFile):
+            self.alt.open_download_stream(oid)
         self.assertEqual(0, self.db.alt.files.count_documents({}))
         self.assertEqual(0, self.db.alt.chunks.count_documents({}))
 
-        self.assertRaises(NoFile, self.alt.open_download_stream, "foo")
+        with self.assertRaises(NoFile):
+            self.alt.open_download_stream("foo")
         self.alt.upload_from_stream("foo", b"hello world")
-        self.assertEqual(b"hello world", self.alt.open_download_stream_by_name("foo").read())
+        self.assertEqual(b"hello world", (self.alt.open_download_stream_by_name("foo")).read())
 
         self.alt.upload_from_stream("mike", b"")
         self.alt.upload_from_stream("test", b"foo")
         self.alt.upload_from_stream("hello world", b"")
 
         self.assertEqual(
             {"mike", "test", "hello world", "foo"},
-            {k["filename"] for k in list(self.db.alt.files.find())},
+            {k["filename"] for k in self.db.alt.files.find().to_list()},
         )
 
     def test_threaded_reads(self):
@@ -240,13 +249,14 @@ def test_get_last_version(self):
         two = two._id
         three = self.fs.upload_from_stream("test", b"baz")
 
-        self.assertEqual(b"baz", self.fs.open_download_stream_by_name("test").read())
+        self.assertEqual(b"baz", (self.fs.open_download_stream_by_name("test")).read())
         self.fs.delete(three)
-        self.assertEqual(b"bar", self.fs.open_download_stream_by_name("test").read())
+        self.assertEqual(b"bar", (self.fs.open_download_stream_by_name("test")).read())
         self.fs.delete(two)
-        self.assertEqual(b"foo", self.fs.open_download_stream_by_name("test").read())
+        self.assertEqual(b"foo", (self.fs.open_download_stream_by_name("test")).read())
         self.fs.delete(one)
-        self.assertRaises(NoFile, self.fs.open_download_stream_by_name, "test")
+        with self.assertRaises(NoFile):
+            self.fs.open_download_stream_by_name("test")
 
     def test_get_version(self):
         self.fs.upload_from_stream("test", b"foo")
@@ -256,28 +266,30 @@ def test_get_version(self):
         self.fs.upload_from_stream("test", b"baz")
         time.sleep(0.01)
 
-        self.assertEqual(b"foo", self.fs.open_download_stream_by_name("test", revision=0).read())
-        self.assertEqual(b"bar", self.fs.open_download_stream_by_name("test", revision=1).read())
-        self.assertEqual(b"baz", self.fs.open_download_stream_by_name("test", revision=2).read())
+        self.assertEqual(b"foo", (self.fs.open_download_stream_by_name("test", revision=0)).read())
+        self.assertEqual(b"bar", (self.fs.open_download_stream_by_name("test", revision=1)).read())
+        self.assertEqual(b"baz", (self.fs.open_download_stream_by_name("test", revision=2)).read())
 
-        self.assertEqual(b"baz", self.fs.open_download_stream_by_name("test", revision=-1).read())
-        self.assertEqual(b"bar", self.fs.open_download_stream_by_name("test", revision=-2).read())
-        self.assertEqual(b"foo", self.fs.open_download_stream_by_name("test", revision=-3).read())
+        self.assertEqual(b"baz", (self.fs.open_download_stream_by_name("test", revision=-1)).read())
+        self.assertEqual(b"bar", (self.fs.open_download_stream_by_name("test", revision=-2)).read())
+        self.assertEqual(b"foo", (self.fs.open_download_stream_by_name("test", revision=-3)).read())
 
-        self.assertRaises(NoFile, self.fs.open_download_stream_by_name, "test", revision=3)
-        self.assertRaises(NoFile, self.fs.open_download_stream_by_name, "test", revision=-4)
+        with self.assertRaises(NoFile):
+            self.fs.open_download_stream_by_name("test", revision=3)
+        with self.assertRaises(NoFile):
+            self.fs.open_download_stream_by_name("test", revision=-4)
 
     def test_upload_from_stream(self):
         oid = self.fs.upload_from_stream("test_file", BytesIO(b"hello world"), chunk_size_bytes=1)
         self.assertEqual(11, self.db.fs.chunks.count_documents({}))
-        self.assertEqual(b"hello world", self.fs.open_download_stream(oid).read())
+        self.assertEqual(b"hello world", (self.fs.open_download_stream(oid)).read())
 
     def test_upload_from_stream_with_id(self):
         oid = ObjectId()
         self.fs.upload_from_stream_with_id(
             oid, "test_file_custom_id", BytesIO(b"custom id"), chunk_size_bytes=1
         )
-        self.assertEqual(b"custom id", self.fs.open_download_stream(oid).read())
+        self.assertEqual(b"custom id", (self.fs.open_download_stream(oid)).read())
 
     @patch("gridfs.synchronous.grid_file._UPLOAD_BUFFER_CHUNKS", 3)
     @client_context.require_failCommand_fail_point
@@ -316,14 +328,14 @@ def test_open_upload_stream(self):
         gin = self.fs.open_upload_stream("from_stream")
         gin.write(b"from stream")
         gin.close()
-        self.assertEqual(b"from stream", self.fs.open_download_stream(gin._id).read())
+        self.assertEqual(b"from stream", (self.fs.open_download_stream(gin._id)).read())
 
     def test_open_upload_stream_with_id(self):
         oid = ObjectId()
         gin = self.fs.open_upload_stream_with_id(oid, "from_stream_custom_id")
         gin.write(b"from stream with custom id")
         gin.close()
-        self.assertEqual(b"from stream with custom id", self.fs.open_download_stream(oid).read())
+        self.assertEqual(b"from stream with custom id", (self.fs.open_download_stream(oid)).read())
 
     def test_missing_length_iter(self):
         # Test fix that guards against PHP-237
@@ -345,12 +357,12 @@ def test_gridfs_lazy_connect(self):
         client = self.single_client("badhost", connect=False, serverSelectionTimeoutMS=0)
         cdb = client.db
         gfs = gridfs.GridFSBucket(cdb)
-        self.assertRaises(ServerSelectionTimeoutError, gfs.delete, 0)
+        with self.assertRaises(ServerSelectionTimeoutError):
+            gfs.delete(0)
 
         gfs = gridfs.GridFSBucket(cdb)
-        self.assertRaises(
-            ServerSelectionTimeoutError, gfs.upload_from_stream, "test", b""
-        )  # Still no connection.
+        with self.assertRaises(ServerSelectionTimeoutError):
+            gfs.upload_from_stream("test", b"")  # Still no connection.
 
     def test_gridfs_find(self):
         self.fs.upload_from_stream("two", b"test2")
@@ -366,14 +378,15 @@ def test_gridfs_find(self):
         cursor = self.fs.find(
             {}, no_cursor_timeout=False, sort=[("uploadDate", -1)], skip=1, limit=2
         )
-        gout = next(cursor)
+        gout = cursor.next()
         self.assertEqual(b"test1", gout.read())
         cursor.rewind()
-        gout = next(cursor)
+        gout = cursor.next()
         self.assertEqual(b"test1", gout.read())
-        gout = next(cursor)
+        gout = cursor.next()
         self.assertEqual(b"test2+", gout.read())
-        self.assertRaises(StopIteration, cursor.__next__)
+        with self.assertRaises(StopIteration):
+            cursor.next()
         cursor.close()
         self.assertRaises(TypeError, self.fs.find, {}, {"_id": True})
 
@@ -383,20 +396,21 @@ def test_grid_in_non_int_chunksize(self):
         self.fs.upload_from_stream("f", data)
         self.db.fs.files.update_one({"filename": "f"}, {"$set": {"chunkSize": 100.0}})
 
-        self.assertEqual(data, self.fs.open_download_stream_by_name("f").read())
+        self.assertEqual(data, (self.fs.open_download_stream_by_name("f")).read())
 
     def test_unacknowledged(self):
         # w=0 is prohibited.
         with self.assertRaises(ConfigurationError):
-            gridfs.GridFSBucket(self.rs_or_single_client(w=0).pymongo_test)
+            gridfs.GridFSBucket((self.rs_or_single_client(w=0)).pymongo_test)
 
     def test_rename(self):
         _id = self.fs.upload_from_stream("first_name", b"testing")
-        self.assertEqual(b"testing", self.fs.open_download_stream_by_name("first_name").read())
+        self.assertEqual(b"testing", (self.fs.open_download_stream_by_name("first_name")).read())
 
         self.fs.rename(_id, "second_name")
-        self.assertRaises(NoFile, self.fs.open_download_stream_by_name, "first_name")
-        self.assertEqual(b"testing", self.fs.open_download_stream_by_name("second_name").read())
+        with self.assertRaises(NoFile):
+            self.fs.open_download_stream_by_name("first_name")
+        self.assertEqual(b"testing", (self.fs.open_download_stream_by_name("second_name")).read())
 
     @patch("gridfs.synchronous.grid_file._UPLOAD_BUFFER_SIZE", 5)
     def test_abort(self):
@@ -407,7 +421,8 @@ def test_abort(self):
         self.assertEqual(3, self.db.fs.chunks.count_documents({"files_id": gin._id}))
         gin.abort()
         self.assertTrue(gin.closed)
-        self.assertRaises(ValueError, gin.write, b"test4")
+        with self.assertRaises(ValueError):
+            gin.write(b"test4")
         self.assertEqual(0, self.db.fs.chunks.count_documents({"files_id": gin._id}))
 
     def test_download_to_stream(self):
@@ -490,7 +505,7 @@ def test_gridfs_replica_set(self):
 
         gfs = gridfs.GridFSBucket(rsc.gfsbucketreplica, "gfsbucketreplicatest")
         oid = gfs.upload_from_stream("test_filename", b"foo")
-        content = gfs.open_download_stream(oid).read()
+        content = (gfs.open_download_stream(oid)).read()
         self.assertEqual(b"foo", content)
 
     def test_gridfs_secondary(self):
@@ -504,7 +519,8 @@ def test_gridfs_secondary(self):
         gfs = gridfs.GridFSBucket(secondary_connection.gfsbucketreplica, "gfsbucketsecondarytest")
 
         # This won't detect secondary, raises error
-        self.assertRaises(NotPrimaryError, gfs.upload_from_stream, "test_filename", b"foo")
+        with self.assertRaises(NotPrimaryError):
+            gfs.upload_from_stream("test_filename", b"foo")
 
     def test_gridfs_secondary_lazy(self):
         # Should detect it's connected to secondary and not attempt to
@@ -518,8 +534,10 @@ def test_gridfs_secondary_lazy(self):
         gfs = gridfs.GridFSBucket(client.gfsbucketreplica, "gfsbucketsecondarylazytest")
 
         # Connects, doesn't create index.
-        self.assertRaises(NoFile, gfs.open_download_stream_by_name, "test_filename")
-        self.assertRaises(NotPrimaryError, gfs.upload_from_stream, "test_filename", b"data")
+        with self.assertRaises(NoFile):
+            gfs.open_download_stream_by_name("test_filename")
+        with self.assertRaises(NotPrimaryError):
+            gfs.upload_from_stream("test_filename", b"data")
 
 
 if __name__ == "__main__":

tools/synchro.py

@@ -216,6 +216,7 @@ def async_only_test(f: str) -> bool:
     "test_examples.py",
     "test_grid_file.py",
     "test_gridfs.py",
+    "test_gridfs_bucket.py",
     "test_gridfs_spec.py",
     "test_heartbeat_monitoring.py",
     "test_index_management.py",