Remove duplicate code from the raise type error and type null tests

Author: bruno-firnkes

bindings/python/pymongoarrow/lib.pyx

@@ -416,7 +416,7 @@ cdef void process_raw_bson_stream(const uint8_t * docstream, size_t length, obje
                     elif value_t == BSON_TYPE_NULL:
                         if context.raise_on_type_null:
                             raise PyMongoArrowError(f"Null value for {key}!")
-                        date32_builder.append_null()    
+                        date32_builder.append_null()
                     else:
                         if context.raise_on_type_error:
                             raise PyMongoArrowError(f"Type mismatch! {key} is not a data32")
@@ -1063,4 +1063,4 @@ cdef class BinaryBuilder(_ArrayBuilderBase):
         return pyarrow_wrap_array(out).cast(BinaryType(self._subtype))
 
     cdef shared_ptr[CBinaryBuilder] unwrap(self):
-        return self.builder
+        return self.builder

bindings/python/test/test_bson.py

@@ -263,10 +263,10 @@ def test_nested(self):
         self._run_test(docs, as_dict)
 
 
-class TestTypeErrorBsonToArrowConversion(TestBsonToArrowConversionBase):
+class TestSimpleBsonToArrowConversion(TestBsonToArrowConversionBase):
     def setUp(self):
         self.schema = Schema(
-            {"_id": ObjectId, "data": int64(), "title": string()}, raise_on_type_error=True
+            {"_id": ObjectId, "data": int64(), "title": string()},
         )
         self.context = PyMongoArrowContext.from_schema(self.schema)
 
@@ -286,12 +286,12 @@ def test_simple(self):
 
         self._run_test(docs, as_dict)
 
-    def test_with_missing_data(self):
+    def test_missing_data(self):
         ids = [ObjectId() for i in range(4)]
         docs = [
             {"_id": ids[0], "data": 10, "title": "a"},
             {"_id": ids[1], "data": 20},
-            {"_id": ids[2]}, 
+            {"_id": ids[2]},
             {"_id": ids[3], "data": 40},
             {"foo": 1},
             {},
@@ -304,28 +304,14 @@ def test_with_missing_data(self):
 
         self._run_test(docs, as_dict)
 
-    def test_with_missmatch_data_and_rising(self):
+    def test_mismatch_data(self):
         ids = [ObjectId() for i in range(4)]
         docs = [
             {"_id": ids[0], "data": 10, "title": "ä"},
             {"_id": ids[1], "data": 20, "title": "b"},
             {"_id": ids[2], "data": "30", "title": "č"},
             {"_id": ids[3], "data": 40, "title": "ê"},
         ]
-
-        payload = type(self)._generate_payload(docs)
-
-        with pytest.raises(PyMongoArrowError, match="Type mismatch! b'data' is not an int64"):
-            process_bson_stream(payload, self.context)
-
-    def test_with_null_data(self):
-        ids = [ObjectId() for i in range(4)]
-        docs = [
-            {"_id": ids[0], "data": 10, "title": "ä"},
-            {"_id": ids[1], "data": 20, "title": "b"},
-            {"_id": ids[2], "data": None, "title": "č"},
-            {"_id": ids[3], "data": 40, "title": "ê"},
-        ]
         as_dict = {
             "_id": ids,
             "data": [10, 20, None, 40],
@@ -334,53 +320,12 @@ def test_with_null_data(self):
 
         self._run_test(docs, as_dict)
 
-class TestTypeNullBsonToArrowConversion(TestBsonToArrowConversionBase):
-    def setUp(self):
-        self.schema = Schema(
-            {"_id": ObjectId, "data": int64(), "title": string()}, raise_on_type_null=True
-        )
-        self.context = PyMongoArrowContext.from_schema(self.schema)
-
-    def test_simple(self):
+    def test_null_data(self):
         ids = [ObjectId() for i in range(4)]
         docs = [
             {"_id": ids[0], "data": 10, "title": "ä"},
             {"_id": ids[1], "data": 20, "title": "b"},
-            {"_id": ids[2], "data": 30, "title": "č"},
-            {"_id": ids[3], "data": 40, "title": "ê"},
-        ]
-        as_dict = {
-            "_id": ids,
-            "data": [10, 20, 30, 40],
-            "title": ["ä", "b", "č", "ê"],
-        }
-
-        self._run_test(docs, as_dict)
-
-    def test_with_missing_data(self):
-        ids = [ObjectId() for i in range(4)]
-        docs = [
-            {"_id": ids[0], "data": 10, "title": "a"},
-            {"_id": ids[1], "data": 20},
-            {"_id": ids[2]}, 
-            {"_id": ids[3], "data": 40},
-            {"foo": 1},
-            {},
-        ]
-        as_dict = {
-            "_id": ids + [None, None],
-            "data": [10, 20, None, 40, None, None],
-            "title": ["a", None, None, None, None, None],
-        }
-
-        self._run_test(docs, as_dict)
-
-    def test_with_missmatch_data(self):
-        ids = [ObjectId() for i in range(4)]
-        docs = [
-            {"_id": ids[0], "data": 10, "title": "ä"},
-            {"_id": ids[1], "data": 20, "title": "b"},
-            {"_id": ids[2], "data": "30", "title": "č"},
+            {"_id": ids[2], "data": None, "title": "č"},
             {"_id": ids[3], "data": 40, "title": "ê"},
         ]
         as_dict = {
@@ -391,62 +336,15 @@ def test_with_missmatch_data(self):
 
         self._run_test(docs, as_dict)
 
-    def test_with_null_data_and_rising(self):
-        ids = [ObjectId() for i in range(4)]
-        docs = [
-            {"_id": ids[0], "data": 10, "title": "ä"},
-            {"_id": ids[1], "data": 20, "title": "b"},
-            {"_id": ids[2], "data": None, "title": "č"},
-            {"_id": ids[3], "data": 40, "title": "ê"},
-        ]
-
-        payload = type(self)._generate_payload(docs)
 
-        with pytest.raises(PyMongoArrowError, match="Null value for b'data'!"):
-            process_bson_stream(payload, self.context)
-
-class TestTypeNullAndTypeErrorBsonToArrowConversion(TestBsonToArrowConversionBase):
+class TestTypeErrorBsonToArrowConversion(TestSimpleBsonToArrowConversion):
     def setUp(self):
         self.schema = Schema(
-            {"_id": ObjectId, "data": int64(), "title": string()}, raise_on_type_null=True, raise_on_type_error=True
+            {"_id": ObjectId, "data": int64(), "title": string()}, raise_on_type_error=True
         )
         self.context = PyMongoArrowContext.from_schema(self.schema)
 
-    def test_simple(self):
-        ids = [ObjectId() for i in range(4)]
-        docs = [
-            {"_id": ids[0], "data": 10, "title": "ä"},
-            {"_id": ids[1], "data": 20, "title": "b"},
-            {"_id": ids[2], "data": 30, "title": "č"},
-            {"_id": ids[3], "data": 40, "title": "ê"},
-        ]
-        as_dict = {
-            "_id": ids,
-            "data": [10, 20, 30, 40],
-            "title": ["ä", "b", "č", "ê"],
-        }
-
-        self._run_test(docs, as_dict)
-
-    def test_with_missing_data(self):
-        ids = [ObjectId() for i in range(4)]
-        docs = [
-            {"_id": ids[0], "data": 10, "title": "a"},
-            {"_id": ids[1], "data": 20},
-            {"_id": ids[2]}, 
-            {"_id": ids[3], "data": 40},
-            {"foo": 1},
-            {},
-        ]
-        as_dict = {
-            "_id": ids + [None, None],
-            "data": [10, 20, None, 40, None, None],
-            "title": ["a", None, None, None, None, None],
-        }
-
-        self._run_test(docs, as_dict)
-
-    def test_with_missmatch_data_and_rising(self):
+    def test_mismatch_data(self):
         ids = [ObjectId() for i in range(4)]
         docs = [
             {"_id": ids[0], "data": 10, "title": "ä"},
@@ -460,7 +358,15 @@ def test_with_missmatch_data_and_rising(self):
         with pytest.raises(PyMongoArrowError, match="Type mismatch! b'data' is not an int64"):
             process_bson_stream(payload, self.context)
 
-    def test_with_null_data_and_rising(self):
+
+class TestTypeNullBsonToArrowConversion(TestSimpleBsonToArrowConversion):
+    def setUp(self):
+        self.schema = Schema(
+            {"_id": ObjectId, "data": int64(), "title": string()}, raise_on_type_null=True
+        )
+        self.context = PyMongoArrowContext.from_schema(self.schema)
+
+    def test_null_data(self):
         ids = [ObjectId() for i in range(4)]
         docs = [
             {"_id": ids[0], "data": 10, "title": "ä"},
@@ -472,4 +378,16 @@ def test_with_null_data_and_rising(self):
         payload = type(self)._generate_payload(docs)
 
         with pytest.raises(PyMongoArrowError, match="Null value for b'data'!"):
-            process_bson_stream(payload, self.context)
+            process_bson_stream(payload, self.context)
+
+
+class TestTypeNullAndTypeErrorBsonToArrowConversion(
+    TestTypeNullBsonToArrowConversion, TestTypeErrorBsonToArrowConversion
+):
+    def setUp(self):
+        self.schema = Schema(
+            {"_id": ObjectId, "data": int64(), "title": string()},
+            raise_on_type_null=True,
+            raise_on_type_error=True,
+        )
+        self.context = PyMongoArrowContext.from_schema(self.schema)