fix the tests

Author: juliusgeo

.github/workflows/test-python.yml

@@ -67,6 +67,7 @@ jobs:
           # Test overshadowed codec_options.py file
           mypy --install-types --non-interactive bson/codec_options.py
           mypy --install-types --non-interactive --disable-error-code var-annotated --disable-error-code attr-defined --disable-error-code union-attr --disable-error-code assignment --disable-error-code no-redef --disable-error-code index --allow-redefinition --allow-untyped-globals --exclude "test/mypy_fails/*.*" test
+          mypy --install-types --non-interactive test/test_mypy.py
 
   linkcheck:
     name: Check Links

pymongo/client_session.py

@@ -435,7 +435,7 @@ def _max_time_expired_error(exc):
 
 # From the transactions spec, all the retryable writes errors plus
 # WriteConcernFailed.
-_UNKNOWN_COMMIT_ERROR_CODES = _RETRYABLE_ERROR_CODES | frozenset(
+_UNKNOWN_COMMIT_ERROR_CODES: frozenset = _RETRYABLE_ERROR_CODES | frozenset(
     [
         64,  # WriteConcernFailed
         50,  # MaxTimeMSExpired

pymongo/helpers.py

@@ -44,7 +44,7 @@
 # From the SDAM spec, the "not primary" error codes are combined with the
 # "node is recovering" error codes (of which the "node is shutting down"
 # errors are a subset).
-_NOT_PRIMARY_CODES = (
+_NOT_PRIMARY_CODES: frozenset = (
     frozenset(
         [
             10058,  # LegacyNotPrimary <=3.2 "not primary" error code
@@ -58,7 +58,7 @@
     | _SHUTDOWN_CODES
 )
 # From the retryable writes spec.
-_RETRYABLE_ERROR_CODES = _NOT_PRIMARY_CODES | frozenset(
+_RETRYABLE_ERROR_CODES: frozenset = _NOT_PRIMARY_CODES | frozenset(
     [
         7,  # HostNotFound
         6,  # HostUnreachable

pymongo/monitoring.py

@@ -528,7 +528,7 @@ def register(listener: _EventListener) -> None:
 # Note - to avoid bugs from forgetting which if these is all lowercase and
 # which are camelCase, and at the same time avoid having to add a test for
 # every command, use all lowercase here and test against command_name.lower().
-_SENSITIVE_COMMANDS = set(
+_SENSITIVE_COMMANDS: set = set(
     [
         "authenticate",
         "saslstart",

test/__init__.py

@@ -43,11 +43,10 @@
     HAVE_IPADDRESS = True
 except ImportError:
     HAVE_IPADDRESS = False
-
 from contextlib import contextmanager
 from functools import wraps
 from test.version import Version
-from typing import Callable, Dict, Generator, no_type_check
+from typing import Any, Callable, Dict, Generator, no_type_check
 from unittest import SkipTest
 from urllib.parse import quote_plus
 
@@ -331,7 +330,9 @@ def hello(self):
 
     def _connect(self, host, port, **kwargs):
         kwargs.update(self.default_client_options)
-        client = pymongo.MongoClient(host, port, serverSelectionTimeoutMS=5000, **kwargs)
+        client: MongoClient = pymongo.MongoClient(
+            host, port, serverSelectionTimeoutMS=5000, **kwargs
+        )
         try:
             try:
                 client.admin.command(HelloCompat.LEGACY_CMD)  # Can we connect?
@@ -356,7 +357,7 @@ def _init_client(self):
         if self.client is not None:
             # Return early when connected to dataLake as mongohoused does not
             # support the getCmdLineOpts command and is tested without TLS.
-            build_info = self.client.admin.command("buildInfo")
+            build_info: Any = self.client.admin.command("buildInfo")
             if "dataLake" in build_info:
                 self.is_data_lake = True
                 self.auth_enabled = True
@@ -521,14 +522,16 @@ def has_secondaries(self):
     @property
     def storage_engine(self):
         try:
-            return self.server_status.get("storageEngine", {}).get("name")
+            return self.server_status.get("storageEngine", {}).get(
+                "name"
+            )  # type:ignore[union-attr]
         except AttributeError:
             # Raised if self.server_status is None.
             return None
 
     def _check_user_provided(self):
         """Return True if db_user/db_password is already an admin user."""
-        client = pymongo.MongoClient(
+        client: MongoClient = pymongo.MongoClient(
             host,
             port,
             username=db_user,
@@ -694,7 +697,7 @@ def supports_secondary_read_pref(self):
         if self.has_secondaries:
             return True
         if self.is_mongos:
-            shard = self.client.config.shards.find_one()["host"]
+            shard = self.client.config.shards.find_one()["host"]  # type:ignore[index]
             num_members = shard.count(",") + 1
             return num_members > 1
         return False
@@ -1015,12 +1018,12 @@ def fork(
         """
 
         def _print_threads(*args: object) -> None:
-            if _print_threads.called:
+            if _print_threads.called:  # type:ignore[attr-defined]
                 return
-            _print_threads.called = True
+            _print_threads.called = True  # type:ignore[attr-defined]
             print_thread_tracebacks()
 
-        _print_threads.called = False
+        _print_threads.called = False  # type:ignore[attr-defined]
 
         def _target() -> None:
             signal.signal(signal.SIGUSR1, _print_threads)

test/test_mypy.py

@@ -28,7 +28,15 @@
     # Not available in Python 3.7
     class Movie(TypedDict):  # type: ignore[misc]
         _id: ObjectId
-        idiot: ObjectId
+        name: str
+        year: int
+
+    class ImplicitMovie(TypedDict):  # type: ignore[misc]
+        _id: NotRequired[ObjectId]
+        name: str
+        year: int
+
+    class EmptyMovie(TypedDict):  # type: ignore[misc]
         name: str
         year: int
 
@@ -329,6 +337,31 @@ def test_typeddict_document_type_insertion(self) -> None:
         assert out["foo"]  # type:ignore[typeddict-item]
         assert type(out["_id"]) == ObjectId
 
+    # 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_not_required(self) -> None:
+        client: MongoClient[ImplicitMovie] = MongoClient()
+        coll: Collection[ImplicitMovie] = client.test.test
+        insert = coll.insert_one(ImplicitMovie(name="THX-1138", year=1971))
+        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
+
+    @only_type_check
+    def test_typeddict_document_type_empty(self) -> None:
+        client: MongoClient[EmptyMovie] = MongoClient()
+        coll: Collection[EmptyMovie] = client.test.test
+        insert = coll.insert_one(EmptyMovie(name="THX-1138", year=1971))
+        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]
+
     @only_type_check
     def test_raw_bson_document_type(self) -> None:
         client = MongoClient(document_class=RawBSONDocument)

test/utils.py

@@ -29,6 +29,7 @@
 from collections import abc, defaultdict
 from functools import partial
 from test import client_context, db_pwd, db_user
+from typing import Any
 
 from bson import json_util
 from bson.objectid import ObjectId
@@ -601,7 +602,7 @@ def ensure_all_connected(client: MongoClient) -> None:
     Depending on the use-case, the caller may need to clear any event listeners
     that are configured on the client.
     """
-    hello = client.admin.command(HelloCompat.LEGACY_CMD)
+    hello: Any = client.admin.command(HelloCompat.LEGACY_CMD)
     if "setName" not in hello:
         raise ConfigurationError("cluster is not a replica set")
 
@@ -612,7 +613,7 @@ def ensure_all_connected(client: MongoClient) -> None:
     def discover():
         i = 0
         while i < 100 and connected_host_list != target_host_list:
-            hello = client.admin.command(
+            hello: Any = client.admin.command(
                 HelloCompat.LEGACY_CMD, read_preference=ReadPreference.SECONDARY
             )
             connected_host_list.update([hello["me"]])