From b391b29bdc2af059261dc00ecfe13a98cf9f2e54 Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Wed, 11 Dec 2024 16:21:42 -0600 Subject: [PATCH] Improve typing in test_runner/fixtures/httpserver.py (#10103) Signed-off-by: Tristan Partin --- test_runner/fixtures/httpserver.py | 15 ++++++++------- test_runner/regress/test_ddl_forwarding.py | 4 +++- test_runner/regress/test_download_extensions.py | 6 ++++-- .../regress/test_pageserver_metric_collection.py | 6 ++++-- .../regress/test_proxy_metric_collection.py | 6 +++++- test_runner/regress/test_sharding.py | 7 +++++-- test_runner/regress/test_storage_controller.py | 8 +++++--- .../regress/test_threshold_based_eviction.py | 6 +++++- 8 files changed, 39 insertions(+), 19 deletions(-) diff --git a/test_runner/fixtures/httpserver.py b/test_runner/fixtures/httpserver.py index f653fd804c..1f46bb22b2 100644 --- a/test_runner/fixtures/httpserver.py +++ b/test_runner/fixtures/httpserver.py @@ -7,24 +7,25 @@ from pytest_httpserver import HTTPServer if TYPE_CHECKING: from collections.abc import Iterator + from ssl import SSLContext from fixtures.port_distributor import PortDistributor -# TODO: mypy fails with: -# Module "fixtures.neon_fixtures" does not explicitly export attribute "PortDistributor" [attr-defined] -# from fixtures.neon_fixtures import PortDistributor + ListenAddress = tuple[str, int] # compared to the fixtures from pytest_httpserver with same names, these are # always function scoped, so you can check and stop the server in tests. @pytest.fixture(scope="function") -def httpserver_ssl_context(): - return None +def httpserver_ssl_context() -> Iterator[SSLContext | None]: + yield None @pytest.fixture(scope="function") -def make_httpserver(httpserver_listen_address, httpserver_ssl_context) -> Iterator[HTTPServer]: +def make_httpserver( + httpserver_listen_address: ListenAddress, httpserver_ssl_context: SSLContext | None +) -> Iterator[HTTPServer]: host, port = httpserver_listen_address if not host: host = HTTPServer.DEFAULT_LISTEN_HOST @@ -47,6 +48,6 @@ def httpserver(make_httpserver: HTTPServer) -> Iterator[HTTPServer]: @pytest.fixture(scope="function") -def httpserver_listen_address(port_distributor: PortDistributor) -> tuple[str, int]: +def httpserver_listen_address(port_distributor: PortDistributor) -> ListenAddress: port = port_distributor.get_port() return ("localhost", port) diff --git a/test_runner/regress/test_ddl_forwarding.py b/test_runner/regress/test_ddl_forwarding.py index 1c5554c379..de44bbcbc8 100644 --- a/test_runner/regress/test_ddl_forwarding.py +++ b/test_runner/regress/test_ddl_forwarding.py @@ -15,6 +15,8 @@ from werkzeug.wrappers.response import Response if TYPE_CHECKING: from typing import Any, Self + from fixtures.httpserver import ListenAddress + def handle_db(dbs, roles, operation): if operation["op"] == "set": @@ -120,7 +122,7 @@ class DdlForwardingContext: @pytest.fixture(scope="function") def ddl( - httpserver: HTTPServer, vanilla_pg: VanillaPostgres, httpserver_listen_address: tuple[str, int] + httpserver: HTTPServer, vanilla_pg: VanillaPostgres, httpserver_listen_address: ListenAddress ): (host, port) = httpserver_listen_address with DdlForwardingContext(httpserver, vanilla_pg, host, port) as ddl: diff --git a/test_runner/regress/test_download_extensions.py b/test_runner/regress/test_download_extensions.py index b2e19ad713..f18f4e78bd 100644 --- a/test_runner/regress/test_download_extensions.py +++ b/test_runner/regress/test_download_extensions.py @@ -20,6 +20,8 @@ from werkzeug.wrappers.response import Response if TYPE_CHECKING: from typing import Any + from fixtures.httpserver import ListenAddress + # use neon_env_builder_local fixture to override the default neon_env_builder fixture # and use a test-specific pg_install instead of shared one @@ -47,8 +49,8 @@ def neon_env_builder_local( def test_remote_extensions( httpserver: HTTPServer, neon_env_builder_local: NeonEnvBuilder, - httpserver_listen_address, - pg_version, + httpserver_listen_address: ListenAddress, + pg_version: PgVersion, ): # setup mock http server # that expects request for anon.tar.zst diff --git a/test_runner/regress/test_pageserver_metric_collection.py b/test_runner/regress/test_pageserver_metric_collection.py index 5ec8357597..aedfdbd210 100644 --- a/test_runner/regress/test_pageserver_metric_collection.py +++ b/test_runner/regress/test_pageserver_metric_collection.py @@ -27,6 +27,8 @@ from werkzeug.wrappers.response import Response if TYPE_CHECKING: from typing import Any + from fixtures.httpserver import ListenAddress + # TODO: collect all of the env setup *AFTER* removal of RemoteStorageKind.NOOP @@ -34,7 +36,7 @@ if TYPE_CHECKING: def test_metric_collection( httpserver: HTTPServer, neon_env_builder: NeonEnvBuilder, - httpserver_listen_address, + httpserver_listen_address: ListenAddress, ): (host, port) = httpserver_listen_address metric_collection_endpoint = f"http://{host}:{port}/billing/api/v1/usage_events" @@ -195,7 +197,7 @@ def test_metric_collection( def test_metric_collection_cleans_up_tempfile( httpserver: HTTPServer, neon_env_builder: NeonEnvBuilder, - httpserver_listen_address, + httpserver_listen_address: ListenAddress, ): (host, port) = httpserver_listen_address metric_collection_endpoint = f"http://{host}:{port}/billing/api/v1/usage_events" diff --git a/test_runner/regress/test_proxy_metric_collection.py b/test_runner/regress/test_proxy_metric_collection.py index dd63256388..5ff4a99c51 100644 --- a/test_runner/regress/test_proxy_metric_collection.py +++ b/test_runner/regress/test_proxy_metric_collection.py @@ -2,6 +2,7 @@ from __future__ import annotations from collections.abc import Iterator from pathlib import Path +from typing import TYPE_CHECKING import pytest from fixtures.log_helper import log @@ -15,6 +16,9 @@ from pytest_httpserver import HTTPServer from werkzeug.wrappers.request import Request from werkzeug.wrappers.response import Response +if TYPE_CHECKING: + from fixtures.httpserver import ListenAddress + def proxy_metrics_handler(request: Request) -> Response: if request.json is None: @@ -38,7 +42,7 @@ def proxy_metrics_handler(request: Request) -> Response: def proxy_with_metric_collector( port_distributor: PortDistributor, neon_binpath: Path, - httpserver_listen_address, + httpserver_listen_address: ListenAddress, test_output_dir: Path, ) -> Iterator[NeonProxy]: """Neon proxy that routes through link auth and has metric collection enabled.""" diff --git a/test_runner/regress/test_sharding.py b/test_runner/regress/test_sharding.py index 30abf91d3a..743ab0088b 100644 --- a/test_runner/regress/test_sharding.py +++ b/test_runner/regress/test_sharding.py @@ -3,7 +3,7 @@ from __future__ import annotations import os import time from collections import defaultdict -from typing import Any +from typing import TYPE_CHECKING, Any import pytest import requests @@ -27,6 +27,9 @@ from typing_extensions import override from werkzeug.wrappers.request import Request from werkzeug.wrappers.response import Response +if TYPE_CHECKING: + from fixtures.httpserver import ListenAddress + def test_sharding_smoke( neon_env_builder: NeonEnvBuilder, @@ -759,7 +762,7 @@ def test_sharding_split_smoke( def test_sharding_split_stripe_size( neon_env_builder: NeonEnvBuilder, httpserver: HTTPServer, - httpserver_listen_address, + httpserver_listen_address: ListenAddress, initial_stripe_size: int, ): """ diff --git a/test_runner/regress/test_storage_controller.py b/test_runner/regress/test_storage_controller.py index 02da389809..5f3a7b39d3 100644 --- a/test_runner/regress/test_storage_controller.py +++ b/test_runner/regress/test_storage_controller.py @@ -58,6 +58,8 @@ from werkzeug.wrappers.response import Response if TYPE_CHECKING: from typing import Any + from fixtures.httpserver import ListenAddress + def get_node_shard_counts(env: NeonEnv, tenant_ids): counts: defaultdict[int, int] = defaultdict(int) @@ -563,7 +565,7 @@ def test_storage_controller_onboard_detached(neon_env_builder: NeonEnvBuilder): def test_storage_controller_compute_hook( httpserver: HTTPServer, neon_env_builder: NeonEnvBuilder, - httpserver_listen_address, + httpserver_listen_address: ListenAddress, ): """ Test that the sharding service calls out to the configured HTTP endpoint on attachment changes @@ -681,7 +683,7 @@ NOTIFY_FAILURE_LOGS = [ def test_storage_controller_stuck_compute_hook( httpserver: HTTPServer, neon_env_builder: NeonEnvBuilder, - httpserver_listen_address, + httpserver_listen_address: ListenAddress, ): """ Test the migration process's behavior when the compute hook does not enable it to proceed @@ -818,7 +820,7 @@ def test_storage_controller_stuck_compute_hook( def test_storage_controller_compute_hook_revert( httpserver: HTTPServer, neon_env_builder: NeonEnvBuilder, - httpserver_listen_address, + httpserver_listen_address: ListenAddress, ): """ 'revert' in the sense of a migration which gets reversed shortly after, as may happen during diff --git a/test_runner/regress/test_threshold_based_eviction.py b/test_runner/regress/test_threshold_based_eviction.py index bedbd84aee..c87b520366 100644 --- a/test_runner/regress/test_threshold_based_eviction.py +++ b/test_runner/regress/test_threshold_based_eviction.py @@ -2,6 +2,7 @@ from __future__ import annotations import time from dataclasses import dataclass +from typing import TYPE_CHECKING from fixtures.log_helper import log from fixtures.neon_fixtures import ( @@ -13,12 +14,15 @@ from fixtures.pageserver.http import LayerMapInfo from fixtures.remote_storage import RemoteStorageKind from pytest_httpserver import HTTPServer +if TYPE_CHECKING: + from fixtures.httpserver import ListenAddress + # NB: basic config change tests are in test_tenant_conf.py def test_threshold_based_eviction( httpserver: HTTPServer, - httpserver_listen_address, + httpserver_listen_address: ListenAddress, pg_bin: PgBin, neon_env_builder: NeonEnvBuilder, ):