PYTHON-3064 Add typings to test package (mongodb#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',