add threading tests

Author: blink1073

bindings/python/pyproject.toml

@@ -101,9 +101,9 @@ testpaths = ["test"]
 log_cli_level = "INFO"
 faulthandler_timeout = 1500
 xfail_strict = true
-filterwarnings = [
-  "error",
-]
+# filterwarnings = [
+#   "error",
+# ]
 
 [tool.ruff]
 line-length = 100

bindings/python/test/test_arrow.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import concurrent.futures
 import io
 import json
 import tempfile
+import threading
 import unittest
 import unittest.mock as mock
 from datetime import date, datetime
@@ -244,8 +246,8 @@ def test_aggregate_omits_id_if_not_in_schema(self):
     def round_trip(self, data, schema, coll=None):
         if coll is None:
             coll = self.coll
-        self.coll.drop()
-        res = write(self.coll, data)
+        coll.drop()
+        res = write(coll, data)
         self.assertEqual(len(data), res.raw_result["insertedCount"])
         self.assertEqual(data, find_arrow_all(coll, {}, schema=schema))
         return res
@@ -1052,6 +1054,29 @@ def test_empty_embedded_array(self):
         assert pmapatable2.to_pylist()[0] == doc2
         write_table(pmapatable2, io.BytesIO())
 
+    def test_threading(self):
+        schema, data = self._create_data()
+
+        def run_test():
+            client = client_context.get_client(
+                event_listeners=[self.getmore_listener, self.cmd_listener]
+            )
+            name = f"test-{threading.current_thread().name}"
+            coll = client.pymongoarrow_test.get_collection(
+                name, write_concern=WriteConcern(w="majority")
+            )
+            coll.drop()
+            self.round_trip(data, Schema(schema), coll=coll)
+            client.close()
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = []
+            for i in range(5):
+                futures.append(executor.submit(run_test))
+            concurrent.futures.wait(futures)
+            for future in futures:
+                future.result()
+
 
 class TestArrowExplicitApi(ArrowApiTestMixin, unittest.TestCase):
     def run_find(self, *args, **kwargs):

bindings/python/test/test_pandas.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # from datetime import datetime, timedelta
+import concurrent.futures
 import datetime
 import tempfile
+import threading
 import unittest
 import unittest.mock as mock
 import warnings
@@ -266,18 +268,22 @@ def inner(i):
         raw_data["nested"] = [inner(i) for i in range(3)]
         return pd.DataFrame(data=raw_data).astype(schema)
 
-    def test_auto_schema(self):
+    def _check_auto_schema(self, coll):
+        coll.drop()
         data = self._create_nested_data()
-        self.coll.drop()
-        res = write(self.coll, data)
+        res = write(coll, data)
         self.assertEqual(len(data), res.raw_result["insertedCount"])
         for func in [find_pandas_all, aggregate_pandas_all]:
-            out = func(self.coll, {} if func == find_pandas_all else []).drop(columns=["_id"])
+            out = func(coll, {} if func == find_pandas_all else []).drop(columns=["_id"])
             for name in data.columns:
                 val = out[name]
                 if str(val.dtype) == "object":
                     val = val.astype(data[name].dtype)
                 pd.testing.assert_series_equal(data[name], val)
+        coll.drop()
+
+    def test_auto_schema(self):
+        self._check_auto_schema(self.coll)
 
     def test_auto_schema_heterogeneous(self):
         vals = [1, "2", True, 4]
@@ -345,6 +351,26 @@ def test_exclude_none(self):
         col_data = list(self.coll.find({}))
         assert "b" not in col_data[3]
 
+    def test_threading(self):
+        def run_test():
+            client = client_context.get_client(
+                event_listeners=[self.getmore_listener, self.cmd_listener]
+            )
+            name = f"test-{threading.current_thread().name}"
+            coll = client.pymongoarrow_test.get_collection(
+                name, write_concern=WriteConcern(w="majority")
+            )
+            self._check_auto_schema(coll)
+            client.close()
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = []
+            for i in range(5):
+                futures.append(executor.submit(run_test))
+            concurrent.futures.wait(futures)
+            for future in futures:
+                future.result()
+
 
 class TestBSONTypes(PandasTestBase):
     @classmethod

bindings/python/test/test_polars.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import concurrent.futures
+import threading
 import unittest
 import unittest.mock as mock
 import uuid
@@ -65,14 +67,13 @@ def setUpClass(cls):
     def setUp(self):
         """Insert simple use case data."""
         self.coll.drop()
-        self.coll.insert_many(
-            [
-                {"_id": 1, "data": 10},
-                {"_id": 2, "data": 20},
-                {"_id": 3, "data": 30},
-                {"_id": 4},
-            ]
-        )
+        self.data = [
+            {"_id": 1, "data": 10},
+            {"_id": 2, "data": 20},
+            {"_id": 3, "data": 30},
+            {"_id": 4},
+        ]
+        self.coll.insert_many(self.data)
         self.cmd_listener.reset()
         self.getmore_listener.reset()
 
@@ -115,15 +116,17 @@ def test_find_simple(self):
         self.assertEqual(find_cmd.command_name, "find")
         self.assertEqual(find_cmd.command["projection"], {"_id": True, "data": True})
 
-    def test_aggregate_simple(self):
+    def _check_aggregation_simple(self, coll):
         expected = pl.DataFrame(
             data={
                 "_id": pl.Series(values=[1, 2, 3, 4], dtype=pl.Int32),
                 "data": pl.Series(values=[20, 40, 60, None], dtype=pl.Int64),
             }
         )
+        coll.drop()
+        coll.insert_many(self.data)
         projection = {"_id": True, "data": {"$multiply": [2, "$data"]}}
-        table = aggregate_polars_all(self.coll, [{"$project": projection}], schema=self.schema)
+        table = aggregate_polars_all(coll, [{"$project": projection}], schema=self.schema)
         self.assertTrue(table.equals(expected))
 
         agg_cmd = self.cmd_listener.results["started"][-1]
@@ -132,6 +135,9 @@ def test_aggregate_simple(self):
         self.assertEqual(agg_cmd.command["pipeline"][0]["$project"], projection)
         self.assertEqual(agg_cmd.command["pipeline"][1]["$project"], {"_id": True, "data": True})
 
+    def test_aggregate_simple(self):
+        self._check_aggregation_simple(self.coll)
+
     @mock.patch.object(Collection, "insert_many", side_effect=Collection.insert_many, autospec=True)
     def test_write_batching(self, mock):
         data = pl.DataFrame(data={"_id": pl.Series(values=range(100040), dtype=pl.Int64)})
@@ -395,3 +401,23 @@ def test_bson_types(self):
                 assert dfpl["value"].dtype == data_type["ptype"]
             except pl.exceptions.ComputeError:
                 assert isinstance(table["value"].type, pa.ExtensionType)
+
+    def test_threading(self):
+        def run_test():
+            client = client_context.get_client(
+                event_listeners=[self.getmore_listener, self.cmd_listener]
+            )
+            name = f"test-{threading.current_thread().name}"
+            coll = client.pymongoarrow_test.get_collection(
+                name, write_concern=WriteConcern(w="majority")
+            )
+            self._check_aggregation_simple(coll)
+            client.close()
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = []
+            for i in range(5):
+                futures.append(executor.submit(run_test))
+            concurrent.futures.wait(futures)
+            for future in futures:
+                future.result()