tests: add type checking and fix type bugs

.pre-commit-config.yaml

@@ -45,7 +45,7 @@ repos:
     rev: v1.0.0
     hooks:
       - id: mypy
-        exclude: ^setup.py|^tests/|^vendor/|^astacus/proto/
+        exclude: ^setup.py|^vendor/|^astacus/proto/
         additional_dependencies:
           - types-PyYAML>=6.0.12.2
           - types-requests>=2.28.11.5

Makefile

@@ -78,6 +78,10 @@ yapf: black
 black:
 	pre-commit run black --all-files
 
+.PHONY: mypy
+mypy:
+	pre-commit run mypy --all-files
+
 .PHONY: reformat
 reformat: isort black
 

astacus/common/cassandra/schema.py

@@ -10,7 +10,7 @@
 from .utils import is_system_keyspace
 from astacus.common.utils import AstacusModel
 from cassandra import metadata as cm
-from typing import Any, Dict, Iterator, List, Set
+from typing import Any, Dict, Iterator, List, Mapping, Set
 
 import hashlib
 import itertools
@@ -236,7 +236,7 @@ def _extract_dcs(metadata: cm.KeyspaceMetadata) -> Dict[str, str]:
 
 
 class CassandraKeyspace(CassandraNamed):
-    network_topology_strategy_dcs: Dict[str, str]  # not [str, int] because of transient replication
+    network_topology_strategy_dcs: Mapping[str, str]  # not [str, int] because of transient replication
     durable_writes: bool
 
     aggregates: List[CassandraAggregate]

astacus/coordinator/plugins/clickhouse/steps.py

@@ -151,7 +151,7 @@ class RetrieveDatabasesAndTablesStep(Step[DatabasesAndTables]):
     node), and relies on that to query only the first server of the cluster.
     """
 
-    clients: List[ClickHouseClient]
+    clients: Sequence[ClickHouseClient]
 
     async def run_step(self, cluster: Cluster, context: StepsContext) -> DatabasesAndTables:
         clickhouse_client = self.clients[0]
@@ -279,7 +279,7 @@ class RemoveFrozenTablesStep(Step[None]):
     When the system unfreeze flag is enabled, clears frozen parts from all disks in a single go.
     """
 
-    clients: List[ClickHouseClient]
+    clients: Sequence[ClickHouseClient]
     freeze_name: str
     unfreeze_timeout: float
 
@@ -293,7 +293,7 @@ async def run_step(self, cluster: Cluster, context: StepsContext) -> None:
 
 @dataclasses.dataclass
 class FreezeUnfreezeTablesStepBase(Step[None]):
-    clients: List[ClickHouseClient]
+    clients: Sequence[ClickHouseClient]
     freeze_name: str
     freeze_unfreeze_timeout: float
 
@@ -429,7 +429,7 @@ class RestoreReplicatedDatabasesStep(Step[None]):
     After this step, all tables will be empty.
     """
 
-    clients: List[ClickHouseClient]
+    clients: Sequence[ClickHouseClient]
     replicated_databases_zookeeper_path: str
     replicated_database_settings: ReplicatedDatabaseSettings
     drop_databases_timeout: float
@@ -615,7 +615,7 @@ class RestoreReplicaStep(Step[None]):
     """
 
     zookeeper_client: ZooKeeperClient
-    clients: List[ClickHouseClient]
+    clients: Sequence[ClickHouseClient]
     disks: Disks
     restart_timeout: float
     max_concurrent_restart_per_node: int
@@ -699,7 +699,7 @@ class AttachMergeTreePartsStep(Step[None]):
     details.
     """
 
-    clients: List[ClickHouseClient]
+    clients: Sequence[ClickHouseClient]
     disks: Disks
     attach_timeout: float
     max_concurrent_attach_per_node: int
@@ -735,7 +735,7 @@ class SyncTableReplicasStep(Step[None]):
     are all exchanged between all nodes.
     """
 
-    clients: List[ClickHouseClient]
+    clients: Sequence[ClickHouseClient]
     sync_timeout: float
     max_concurrent_sync_per_node: int
 

tests/integration/conftest.py

@@ -68,7 +68,7 @@ async def run_process_and_wait_for_pattern(
     fail_pattern: Optional[str] = None,
     env: Mapping[str, str] = MappingProxyType({}),
     timeout: float = 10.0,
-) -> AsyncIterator[asyncio.subprocess.Process]:
+) -> AsyncIterator[subprocess.Popen[bytes]]:
     # This stringification is a workaround for a bug in pydev (pydev_monkey.py:111)
     str_args = [str(arg) for arg in args]
     pattern_found = asyncio.Event()
@@ -102,7 +102,7 @@ def read_logs() -> None:
 
 @dataclasses.dataclass
 class Service:
-    process: asyncio.subprocess.Process
+    process: subprocess.Popen[bytes] | asyncio.subprocess.Process
     data_dir: Path
     port: int
     host: str = "localhost"
@@ -161,6 +161,8 @@ async def create_zookeeper(ports: Ports) -> AsyncIterator[Service]:
     java_path = await get_command_path("java")
     if java_path is None:
         pytest.skip("java installation not found")
+        # newer versions of mypy should be able to infer that this is unreachable
+        assert False
     port = ports.allocate()
     with tempfile.TemporaryDirectory(prefix=f"zookeeper_{port}_") as data_dir_str:
         data_dir = Path(data_dir_str)
@@ -177,6 +179,7 @@ async def create_zookeeper(ports: Ports) -> AsyncIterator[Service]:
         command = get_zookeeper_command(java_path=java_path, data_dir=data_dir, port=port)
         if command is None:
             pytest.skip("zookeeper installation not found")
+            assert False
         async with contextlib.AsyncExitStack() as stack:
             max_attempts = 10
             for attempt in range(max_attempts):
@@ -190,9 +193,9 @@ async def create_zookeeper(ports: Ports) -> AsyncIterator[Service]:
                             timeout=20.0,
                         )
                     )
+                    yield Service(process=process, port=port, data_dir=data_dir)
                     break
                 except FailPatternFoundError:
                     if attempt + 1 == max_attempts:
                         raise
                     await asyncio.sleep(2.0)
-            yield Service(process=process, port=port, data_dir=data_dir)

tests/integration/coordinator/plugins/clickhouse/conftest.py

@@ -38,6 +38,7 @@
 import pytest
 import rohmu
 import secrets
+import subprocess
 import sys
 import tempfile
 import urllib.parse
@@ -126,7 +127,7 @@ class MinioBucket:
 
 @dataclasses.dataclass(frozen=True)
 class MinioService:
-    process: asyncio.subprocess.Process
+    process: subprocess.Popen[bytes]
     data_dir: Path
     host: str
     server_port: int
@@ -217,7 +218,15 @@ async def create_minio_service(ports: Ports) -> AsyncIterator[MinioService]:
             "MINIO_ROOT_USER": root_user,
             "MINIO_ROOT_PASSWORD": root_password,
         }
-        command = ["/usr/bin/minio", "server", data_dir, "--address", server_netloc, "--console-address", console_netloc]
+        command: List[Union[str, Path]] = [
+            "/usr/bin/minio",
+            "server",
+            data_dir,
+            "--address",
+            server_netloc,
+            "--console-address",
+            console_netloc,
+        ]
         async with run_process_and_wait_for_pattern(
             args=command,
             cwd=data_dir,

tests/integration/coordinator/plugins/clickhouse/test_client.py

@@ -5,6 +5,7 @@
 from .conftest import ClickHouseCommand, create_clickhouse_service, get_clickhouse_client
 from astacus.coordinator.plugins.clickhouse.client import ClickHouseClientQueryError
 from tests.integration.conftest import Ports, Service
+from typing import cast, Sequence
 
 import pytest
 import time
@@ -18,7 +19,7 @@
 @pytest.mark.asyncio
 async def test_client_execute(clickhouse: Service) -> None:
     client = get_clickhouse_client(clickhouse)
-    response = await client.execute(b"SHOW DATABASES")
+    response = cast(Sequence[list[str]], await client.execute(b"SHOW DATABASES"))
     assert sorted(list(response)) == [["INFORMATION_SCHEMA"], ["default"], ["information_schema"], ["system"]]
 
 

tests/integration/coordinator/plugins/clickhouse/test_plugin.py

@@ -278,7 +278,7 @@ async def test_restores_table_with_nested_fields(restored_cluster: List[ClickHou
     assert response == [[123, [4], [5]]]
 
 
-async def check_object_storage_data(cluster: List[ClickHouseClient]) -> None:
+async def check_object_storage_data(cluster: Sequence[ClickHouseClient]) -> None:
     s1_data = [[123, "foo"], [456, "bar"]]
     s2_data = [[789, "baz"]]
     cluster_data = [s1_data, s1_data, s2_data]

tests/integration/coordinator/plugins/clickhouse/test_steps.py

@@ -9,6 +9,7 @@
 from astacus.coordinator.plugins.clickhouse.steps import RetrieveDatabasesAndTablesStep
 from base64 import b64decode
 from tests.integration.conftest import create_zookeeper, Ports
+from typing import cast, Sequence
 from uuid import UUID
 
 import pytest
@@ -50,11 +51,17 @@ async def test_retrieve_tables(ports: Ports, clickhouse_command: ClickHouseComma
             step = RetrieveDatabasesAndTablesStep(clients=[client])
             context = StepsContext()
             databases, tables = await step.run_step(Cluster(nodes=[]), context=context)
-            database_uuid_lines = list(await client.execute(b"SELECT base64Encode(name),uuid FROM system.databases"))
+            database_uuid_lines = cast(
+                Sequence[tuple[str, str]],
+                await client.execute(b"SELECT base64Encode(name),uuid FROM system.databases"),
+            )
             database_uuids = {
                 b64decode(database_name): UUID(database_uuid) for database_name, database_uuid in database_uuid_lines
             }
-            table_uuid_lines = list(await client.execute("SELECT uuid FROM system.tables where name = 'tablé_1'".encode()))
+            table_uuid_lines: Sequence[Sequence[str]] = cast(
+                list[tuple[str]],
+                list(await client.execute("SELECT uuid FROM system.tables where name = 'tablé_1'".encode())),
+            )
             table_uuid = UUID(table_uuid_lines[0][0])
 
             assert databases == [

tests/integration/coordinator/plugins/flink/test_steps.py

@@ -6,6 +6,7 @@
 from astacus.coordinator.plugins.flink.manifest import FlinkManifest
 from astacus.coordinator.plugins.flink.steps import FlinkManifestStep, RestoreDataStep, RetrieveDataStep
 from astacus.coordinator.plugins.zookeeper import KazooZooKeeperClient
+from unittest.mock import Mock
 from uuid import uuid4
 
 import pytest
@@ -35,6 +36,6 @@ async def test_restore_data(zookeeper_client: KazooZooKeeperClient):
     manifest = FlinkManifest(data=data)
     context = StepsContext()
     context.set_result(FlinkManifestStep, manifest)
-    await RestoreDataStep(zookeeper_client, ["catalog", "flink"]).run_step(cluster=None, context=context)
-    res = await RetrieveDataStep(zookeeper_client, ["catalog", "flink"]).run_step(cluster=None, context=None)
+    await RestoreDataStep(zookeeper_client, ["catalog", "flink"]).run_step(cluster=Mock(), context=context)
+    res = await RetrieveDataStep(zookeeper_client, ["catalog", "flink"]).run_step(cluster=Mock(), context=Mock())
     assert res == data

tests/system/conftest.py

@@ -188,4 +188,5 @@ def astacus_run(
 
 
 def astacus_ls(astacus: TestNode) -> List[str]:
+    assert astacus.root_path
     return sorted(str(x.relative_to(astacus.root_path)) for x in astacus.root_path.glob("**/*"))

tests/system/test_config_reload.py

@@ -27,6 +27,7 @@ def test_reload_config(tmpdir, rootdir: str, astacus1: TestNode, astacus2: TestN
     check_without_status = astacus_run(rootdir, astacus1, "check-reload", check=True, capture_output=True)
     assert check_without_status.returncode == 0
     assert "Configuration does not need to be reloaded" in check_without_status.stdout.decode()
+    assert astacus1.root_path
     # Write some data to backup, including files that don't match the reloaded glob
     (astacus1.root_path / "saved.foo").write_text("dont_care")
     (astacus1.root_path / "ignored.bar").write_text("dont_care")

tests/unit/common/cassandra/test_client.py

@@ -4,12 +4,13 @@
 """
 
 from astacus.common.cassandra.config import CassandraClientConfiguration
+from pytest_mock import MockFixture
 
 import astacus.common.cassandra.client as client_module
 import pytest
 
 
-def test_cassandra_session(mocker):
+def test_cassandra_session(mocker: MockFixture) -> None:
     ccluster = mocker.MagicMock()
     csession = mocker.MagicMock()
     session = client_module.CassandraSession(cluster=ccluster, session=csession)
@@ -32,7 +33,7 @@ def test_cassandra_session(mocker):
     assert len(csession.execute.mock_calls) == 2
 
 
-def create_client(mocker, ssl=False):
+def create_client(mocker: MockFixture, ssl: bool = False) -> client_module.CassandraClient:
     mocker.patch.object(client_module, "Cluster")
     mocker.patch.object(client_module, "WhiteListRoundRobinPolicy")
 
@@ -47,14 +48,14 @@ def create_client(mocker, ssl=False):
 
 
 @pytest.mark.parametrize("ssl", [False, True])
-def test_cassandra_client(mocker, ssl):
+def test_cassandra_client(mocker: MockFixture, ssl: bool) -> None:
     client: client_module.CassandraClient = create_client(mocker, ssl=ssl)
     with client.connect() as session:
         assert isinstance(session, client_module.CassandraSession)
 
 
 @pytest.mark.asyncio
-async def test_cassandra_client_run(mocker):
+async def test_cassandra_client_run(mocker: MockFixture):
     client = create_client(mocker)
 
     def test_fun(cas):

tests/unit/common/test_op.py

@@ -7,6 +7,7 @@
 
 """
 from astacus.common import op
+from astacus.common.exceptions import ExpiredOperationException
 from astacus.common.statsd import StatsClient
 from starlette.background import BackgroundTasks
 from starlette.datastructures import URL
@@ -35,7 +36,7 @@ def set_status_fail(self):
         (None, op.Op.Status.done, None),
         # If operation throws ExpiredOperationException, op status
         # should stay running as it may point to the next operation
-        (op.ExpiredOperationException, op.Op.Status.running, None),
+        (ExpiredOperationException, op.Op.Status.running, None),
         # If operation throws 'something else', it should fail the op status
         (AssertionError, op.Op.Status.fail, AssertionError),
     ],

tests/unit/common/test_storage.py

@@ -11,7 +11,7 @@
 from astacus.common import exceptions
 from astacus.common.cachingjsonstorage import CachingJsonStorage
 from astacus.common.rohmustorage import RohmuConfig, RohmuStorage
-from astacus.common.storage import FileStorage, JsonStorage
+from astacus.common.storage import FileStorage, Json, JsonStorage
 from contextlib import nullcontext as does_not_raise
 from pathlib import Path
 from rohmu.object_storage import google
@@ -24,7 +24,7 @@
 TEXT_HEXDIGEST_DATA = b"data" * 15
 
 TEST_JSON = "jsonblob"
-TEST_JSON_DATA = {"foo": 7, "array": [1, 2, 3], "true": True}
+TEST_JSON_DATA: Json = {"foo": 7, "array": [1, 2, 3], "true": True}
 
 
 def create_storage(*, tmpdir, engine, **kw):

tests/unit/coordinator/plugins/cassandra/test_backup_steps.py

@@ -30,7 +30,7 @@ class RetrieveTestCase:
     duplicate_address: bool = False
 
     # Output
-    expected_error: Optional[Exception] = None
+    expected_error: Optional[type[Exception]] = None
 
     def __str__(self):
         return self.name

tests/unit/coordinator/plugins/cassandra/test_plugin.py

@@ -7,7 +7,7 @@
 from astacus.coordinator.plugins.base import StepFailedError
 from astacus.coordinator.plugins.cassandra import plugin
 from astacus.coordinator.plugins.cassandra.model import CassandraConfigurationNode
-from tests.unit.node.test_node_cassandra import CassandraTestConfig
+from tests.unit.conftest import CassandraTestConfig
 from types import SimpleNamespace
 
 import pytest

tests/unit/coordinator/plugins/clickhouse/test_config.py

@@ -3,16 +3,16 @@
 See LICENSE for details
 """
 from astacus.coordinator.plugins.clickhouse.client import HttpClickHouseClient
-from astacus.coordinator.plugins.clickhouse.config import ClickHouseConfiguration, ClickHouseNode
-from astacus.coordinator.plugins.clickhouse.plugin import get_clickhouse_clients, get_zookeeper_client
-from astacus.coordinator.plugins.zookeeper import KazooZooKeeperClient, KazooZooKeeperConnection
+from astacus.coordinator.plugins.clickhouse.config import ClickHouseConfiguration, ClickHouseNode, get_clickhouse_clients
+from astacus.coordinator.plugins.zookeeper import KazooZooKeeperClient, KazooZooKeeperConnection, ZooKeeperUser
 from astacus.coordinator.plugins.zookeeper_config import (
+    get_zookeeper_client,
     ZooKeeperConfiguration,
     ZooKeeperConfigurationUser,
     ZooKeeperNode,
-    ZooKeeperUser,
 )
 from kazoo.client import KazooClient
+from pydantic import SecretStr
 from typing import cast, List
 
 import pytest
@@ -31,7 +31,7 @@ def test_get_zookeeper_client() -> None:
 def test_get_authenticated_zookeeper_client() -> None:
     configuration = ZooKeeperConfiguration(
         nodes=[ZooKeeperNode(host="::1", port=5556)],
-        user=ZooKeeperConfigurationUser(username="local-user", password="secret"),
+        user=ZooKeeperConfigurationUser(username="local-user", password=SecretStr("secret")),
     )
     client = get_zookeeper_client(configuration)
     assert client is not None and isinstance(client, KazooZooKeeperClient)