PYTHON-3064 Add typings to test package (#844)

Author: blink1073

.github/workflows/test-python.yml

@@ -46,4 +46,5 @@ 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 pymongo
+          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 test

bson/son.py

@@ -28,7 +28,7 @@
 # This is essentially the same as re._pattern_type
 RE_TYPE: Type[Pattern[Any]] = type(re.compile(""))
 
-_Key = TypeVar("_Key", bound=str)
+_Key = TypeVar("_Key")
 _Value = TypeVar("_Value")
 _T = TypeVar("_T")
 

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,12 @@ ignore_missing_imports = True
 [mypy-snappy.*]
 ignore_missing_imports = True
 
+[mypy-test.*]
+allow_redefinition = true
+allow_untyped_globals = 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[float] = 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/srv_resolver.py

@@ -39,7 +39,7 @@ def maybe_decode(text):
 def _resolve(*args, **kwargs):
     if hasattr(resolver, 'resolve'):
         # dnspython >= 2
-        return resolver.resolve(*args, **kwargs)  # type: ignore
+        return resolver.resolve(*args, **kwargs)
     # dnspython 1.X
     return resolver.query(*args, **kwargs)
 

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:
@@ -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
@@ -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',