PYTHON-3064 Add typings to test package

.github/workflows/test-python.yml

@@ -46,4 +46,4 @@ jobs:
           pip install -e ".[zstd, srv]"
       - name: Run mypy
         run: |
-          mypy --install-types --non-interactive bson gridfs tools
+          mypy --install-types --non-interactive bson gridfs tools test

mypy.ini

@@ -11,6 +11,9 @@ warn_unused_configs = true
 warn_unused_ignores = true
 warn_redundant_casts = true
 
+[mypy-gevent.*]
+ignore_missing_imports = True
+
 [mypy-kerberos.*]
 ignore_missing_imports = True
 
@@ -29,5 +32,11 @@ ignore_missing_imports = True
 [mypy-snappy.*]
 ignore_missing_imports = True
 
+[mypy-test.*]
+allow_redefinition = true
+
 [mypy-winkerberos.*]
 ignore_missing_imports = True
+
+[mypy-xmlrunner.*]
+ignore_missing_imports = True

pymongo/socket_checker.py

@@ -16,9 +16,8 @@
 
 import errno
 import select
-import socket
 import sys
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 # PYTHON-2320: Jython does not fully support poll on SSL sockets,
 # https://bugs.jython.org/issue2900
@@ -43,7 +42,7 @@ def __init__(self) -> None:
         else:
             self._poller = None
 
-    def select(self, sock: Any, read: bool = False, write: bool = False, timeout: int = 0) -> bool:
+    def select(self, sock: Any, read: bool = False, write: bool = False, timeout: Optional[Union[float, int]] = 0) -> bool:
         """Select for reads or writes with a timeout in seconds (or None).
 
         Returns True if the socket is readable/writable, False on timeout.

pymongo/typings.py

@@ -14,7 +14,7 @@
 
 """Type aliases used by PyMongo"""
 from typing import (TYPE_CHECKING, Any, Dict, List, Mapping, MutableMapping, Optional,
-                    Tuple, Type, TypeVar, Union)
+                    Sequence, Tuple, Type, TypeVar, Union)
 
 if TYPE_CHECKING:
     from bson.raw_bson import RawBSONDocument
@@ -25,5 +25,5 @@
 _Address = Tuple[str, Optional[int]]
 _CollationIn = Union[Mapping[str, Any], "Collation"]
 _DocumentIn = Union[MutableMapping[str, Any], "RawBSONDocument"]
-_Pipeline = List[Mapping[str, Any]]
+_Pipeline = Sequence[Mapping[str, Any]]
 _DocumentType = TypeVar('_DocumentType', Mapping[str, Any], MutableMapping[str, Any], Dict[str, Any])

test/__init__.py

@@ -40,6 +40,7 @@
 
 from contextlib import contextmanager
 from functools import wraps
+from typing import Dict, no_type_check
 from unittest import SkipTest
 
 import pymongo
@@ -48,7 +49,9 @@
 from bson.son import SON
 from pymongo import common, message
 from pymongo.common import partition_node
+from pymongo.database import Database
 from pymongo.hello import HelloCompat
+from pymongo.mongo_client import MongoClient
 from pymongo.server_api import ServerApi
 from pymongo.ssl_support import HAVE_SSL, _ssl
 from pymongo.uri_parser import parse_uri
@@ -86,7 +89,7 @@
                             os.path.join(CERT_PATH, 'client.pem'))
 CA_PEM = os.environ.get('CA_PEM', os.path.join(CERT_PATH, 'ca.pem'))
 
-TLS_OPTIONS = dict(tls=True)
+TLS_OPTIONS: Dict = dict(tls=True)
 if CLIENT_PEM:
     TLS_OPTIONS['tlsCertificateKeyFile'] = CLIENT_PEM
 if CA_PEM:
@@ -102,13 +105,13 @@
     # Remove after PYTHON-2712
     from pymongo import pool
     pool._MOCK_SERVICE_ID = True
-    res = parse_uri(SINGLE_MONGOS_LB_URI)
+    res = parse_uri(SINGLE_MONGOS_LB_URI or "")
     host, port = res['nodelist'][0]
     db_user = res['username'] or db_user
     db_pwd = res['password'] or db_pwd
 elif TEST_SERVERLESS:
     TEST_LOADBALANCER = True
-    res = parse_uri(SINGLE_MONGOS_LB_URI)
+    res = parse_uri(SINGLE_MONGOS_LB_URI or "")
     host, port = res['nodelist'][0]
     db_user = res['username'] or db_user
     db_pwd = res['password'] or db_pwd
@@ -184,6 +187,7 @@ def enable(self):
     def __enter__(self):
         self.enable()
 
+    @no_type_check
     def disable(self):
         common.HEARTBEAT_FREQUENCY = self.old_heartbeat_frequency
         common.MIN_HEARTBEAT_INTERVAL = self.old_min_heartbeat_interval
@@ -224,6 +228,8 @@ def _all_users(db):
 
 
 class ClientContext(object):
+    client: MongoClient
+
     MULTI_MONGOS_LB_URI = MULTI_MONGOS_LB_URI
 
     def __init__(self):
@@ -247,9 +253,9 @@ def __init__(self):
         self.tls = False
         self.tlsCertificateKeyFile = False
         self.server_is_resolvable = is_server_resolvable()
-        self.default_client_options = {}
+        self.default_client_options: Dict = {}
         self.sessions_enabled = False
-        self.client = None
+        self.client = None  # type: ignore
         self.conn_lock = threading.Lock()
         self.is_data_lake = False
         self.load_balancer = TEST_LOADBALANCER
@@ -340,6 +346,7 @@ def _init_client(self):
                 try:
                     self.cmd_line = self.client.admin.command('getCmdLineOpts')
                 except pymongo.errors.OperationFailure as e:
+                    assert e.details is not None
                     msg = e.details.get('errmsg', '')
                     if e.code == 13 or 'unauthorized' in msg or 'login' in msg:
                         # Unauthorized.
@@ -418,6 +425,7 @@ def _init_client(self):
             else:
                 self.server_parameters = self.client.admin.command(
                     'getParameter', '*')
+                assert self.cmd_line is not None
                 if 'enableTestCommands=1' in self.cmd_line['argv']:
                     self.test_commands_enabled = True
                 elif 'parsed' in self.cmd_line:
@@ -436,7 +444,8 @@ def _init_client(self):
                 self.mongoses.append(address)
                 if not self.serverless:
                     # Check for another mongos on the next port.
-                    next_address = address[0], address[1] + 1
+                    assert address is not None
+                    next_address = address[0], address[1]  + 1
                     mongos_client = self._connect(
                         *next_address, **self.default_client_options)
                     if mongos_client:
@@ -479,7 +488,7 @@ 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
@@ -496,6 +505,7 @@ def _check_user_provided(self):
         try:
             return db_user in _all_users(client.admin)
         except pymongo.errors.OperationFailure as e:
+            assert e.details is not None
             msg = e.details.get('errmsg', '')
             if e.code == 18 or 'auth fails' in msg:
                 # Auth failed.
@@ -505,6 +515,7 @@ def _check_user_provided(self):
 
     def _server_started_with_auth(self):
         # MongoDB >= 2.0
+        assert self.cmd_line is not None
         if 'parsed' in self.cmd_line:
             parsed = self.cmd_line['parsed']
             # MongoDB >= 2.6
@@ -525,6 +536,7 @@ def _server_started_with_ipv6(self):
         if not socket.has_ipv6:
             return False
 
+        assert self.cmd_line is not None
         if 'parsed' in self.cmd_line:
             if not self.cmd_line['parsed'].get('net', {}).get('ipv6'):
                 return False
@@ -642,7 +654,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
@@ -932,6 +944,9 @@ def fail_point(self, command_args):
 
 class IntegrationTest(PyMongoTestCase):
     """Base class for TestCases that need a connection to MongoDB to pass."""
+    client: MongoClient
+    db: Database
+    credentials: Dict[str, str]
 
     @classmethod
     @client_context.require_connection
@@ -1073,7 +1088,7 @@ def run(self, test):
 
 
 if HAVE_XML:
-    class PymongoXMLTestRunner(XMLTestRunner):
+    class PymongoXMLTestRunner(XMLTestRunner):  # type: ignore[misc]
         def run(self, test):
             setup()
             result = super(PymongoXMLTestRunner, self).run(test)

test/auth_aws/test_auth_aws.py

@@ -26,6 +26,7 @@
 
 
 class TestAuthAWS(unittest.TestCase):
+    uri: str
 
     @classmethod
     def setUpClass(cls):

test/mockupdb/test_cursor_namespace.py

@@ -21,6 +21,9 @@
 
 
 class TestCursorNamespace(unittest.TestCase):
+    server: MockupDB
+    client: MongoClient
+
     @classmethod
     def setUpClass(cls):
         cls.server = MockupDB(auto_ismaster={'maxWireVersion': 6})
@@ -69,6 +72,9 @@ def op():
 
 
 class TestKillCursorsNamespace(unittest.TestCase):
+    server: MockupDB
+    client: MongoClient
+
     @classmethod
     def setUpClass(cls):
         cls.server = MockupDB(auto_ismaster={'maxWireVersion': 6})

test/mockupdb/test_getmore_sharded.py

@@ -27,7 +27,7 @@ def test_getmore_sharded(self):
         servers = [MockupDB(), MockupDB()]
 
         # Collect queries to either server in one queue.
-        q = Queue()
+        q: Queue = Queue()
         for server in servers:
             server.subscribe(q.put)
             server.autoresponds('ismaster', ismaster=True, msg='isdbgrid',

test/mockupdb/test_handshake.py

@@ -48,19 +48,19 @@ def respond(r):
                                         ServerApiVersion.V1))}
     client = MongoClient("mongodb://"+primary.address_string,
                          appname='my app', # For _check_handshake_data()
-                         **dict([k_map.get((k, v), (k, v)) for k, v
+                         **dict([k_map.get((k, v), (k, v)) for k, v  # type: ignore[arg-type]
                                  in kwargs.items()]))
-    
+
     self.addCleanup(client.close)
 
     # We have an autoresponder luckily, so no need for `go()`.
     assert client.db.command(hello)
 
     # We do this checking here rather than in the autoresponder `respond()`
     # because it runs in another Python thread so there are some funky things
-    # with error handling within that thread, and we want to be able to use 
+    # with error handling within that thread, and we want to be able to use
     # self.assertRaises().
-    self.handshake_req.assert_matches(protocol(hello, **kwargs))
+    self.handshake_req.assert_matches(protocol(hello, **kwargs))  # type: ignore[attr-defined]
     _check_handshake_data(self.handshake_req)
 
 

test/mockupdb/test_mixed_version_sharded.py

@@ -30,7 +30,7 @@ def setup_server(self, upgrade):
         self.mongos_old, self.mongos_new = MockupDB(), MockupDB()
 
         # Collect queries to either server in one queue.
-        self.q = Queue()
+        self.q: Queue = Queue()
         for server in self.mongos_old, self.mongos_new:
             server.subscribe(self.q.put)
             server.autoresponds('getlasterror')
@@ -59,7 +59,7 @@ def create_mixed_version_sharded_test(upgrade):
     def test(self):
         self.setup_server(upgrade)
         start = time.time()
-        servers_used = set()
+        servers_used: set = set()
         while len(servers_used) < 2:
             go(upgrade.function, self.client)
             request = self.q.get(timeout=1)

test/mockupdb/test_op_msg.py

@@ -233,6 +233,8 @@
 
 
 class TestOpMsg(unittest.TestCase):
+    server: MockupDB
+    client: MongoClient
 
     @classmethod
     def setUpClass(cls):

test/mockupdb/test_op_msg_read_preference.py

@@ -14,6 +14,7 @@
 
 import copy
 import itertools
+from typing import Any
 
 from mockupdb import MockupDB, going, CommandBase
 from pymongo import MongoClient, ReadPreference
@@ -27,6 +28,8 @@
 
 class OpMsgReadPrefBase(unittest.TestCase):
     single_mongod = False
+    primary: MockupDB
+    secondary: MockupDB
 
     @classmethod
     def setUpClass(cls):
@@ -142,7 +145,7 @@ def test(self):
                                     tag_sets=None)
 
         client = self.setup_client(read_preference=pref)
-
+        expected_pref: Any
         if operation.op_type == 'always-use-secondary':
             expected_server = self.secondary
             expected_pref = ReadPreference.SECONDARY

test/mockupdb/test_query_read_pref_sharded.py

@@ -50,7 +50,7 @@ def test_query_and_read_mode_sharded_op_msg(self):
             for pref in read_prefs:
                 collection = client.db.get_collection('test',
                                                       read_preference=pref)
-                cursor = collection.find(query.copy())
+                cursor = collection.find(query.copy())  # type: ignore[attr-defined]
                 with going(next, cursor):
                     request = server.receives()
                     # Command is not nested in $query.