fix tests more

Author: juliusgeo

test/test_mypy.py

@@ -339,12 +339,92 @@ def test_typeddict_document_type(self) -> None:
 
     @only_type_check
     def test_typeddict_document_type_insertion(self) -> None:
-        coll: Collection[Movie] = self.db.test
-        mov = Movie(name="THX-1138", year=1971)
+        client: MongoClient[Movie] = MongoClient()
+        coll = client.test.test
+        mov = Movie(_id=ObjectId(), name="THX-1138", year=1971)
         coll.insert_one(mov)
         coll.insert_many([mov])
         coll.bulk_write([InsertOne(mov)])
 
+    # This should work the same as the test above, but this time using NotRequired to allow
+    # automatic insertion of the _id field by insert_one.
+    @only_type_check
+    def test_typeddict_document_type_normal(self) -> None:
+        client: MongoClient[Movie] = MongoClient()
+        coll: Collection[Movie] = client.test.test
+        mov = Movie(_id=ObjectId(), name="THX-1138", year=1971)
+        coll.insert_many([mov])
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        # This should fail because the output is a Movie.
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId
+        coll.drop()
+        coll.insert_one(mov)
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId
+        coll.drop()
+        coll.bulk_write([InsertOne(mov)])
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId
+        coll.drop()
+
+    @only_type_check
+    def test_typeddict_document_type_not_required(self) -> None:
+        client: MongoClient[ImplicitMovie] = MongoClient()
+        coll: Collection[ImplicitMovie] = client.test.test
+        mov = ImplicitMovie(name="THX-1138", year=1971)
+        coll.insert_many([mov])
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        # This should fail because the output is a Movie.
+        assert out["foo"]  # type:ignore[typeddict-item]
+        # This should fail because _id is not included in our TypedDict definition.
+        assert type(out["_id"]) == ObjectId  # type:ignore[typeddict-item]
+        coll.insert_one(mov)
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId
+        coll.bulk_write([InsertOne(mov)])
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId
+
+        assert out is not None
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId
+        coll.drop()
+
+    @only_type_check
+    def test_typeddict_document_type_no_id(self) -> None:
+        client: MongoClient[MovieWithoutId] = MongoClient()
+        coll: Collection[MovieWithoutId] = client.test.test
+        mov = MovieWithoutId(name="THX-1138", year=1971)
+        coll.insert_many([mov])
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId  # type:ignore[typeddict-item]
+        coll.drop()
+        coll.insert_one(mov)
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId  # type:ignore[typeddict-item]
+        coll.drop()
+        coll.bulk_write([InsertOne(mov)])
+        out = coll.find_one({"name": "THX-1138"})
+        assert out is not None
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert type(out["_id"]) == ObjectId  # type:ignore[typeddict-item]
+        coll.drop()
+
     @only_type_check
     def test_raw_bson_document_type(self) -> None:
         client = MongoClient(document_class=RawBSONDocument)