diff --git a/tests/blocks/test_system.py b/tests/blocks/test_system.py index 57e332c803c0..ee95b4b43937 100644 --- a/tests/blocks/test_system.py +++ b/tests/blocks/test_system.py @@ -1,16 +1,21 @@ -import pendulum +from typing import Any + import pytest from pydantic import Secret as PydanticSecret from pydantic import SecretStr from prefect.blocks.system import DateTime, Secret -from prefect.types import DateTime as PydanticDateTime +from prefect.types._datetime import DateTime as PydanticDateTime +from prefect.types._datetime import Timezone -def test_datetime(ignore_prefect_deprecation_warnings): - DateTime(value=PydanticDateTime(2022, 1, 1)).save(name="test") +@pytest.mark.usefixtures("ignore_prefect_deprecation_warnings") +def test_datetime(): + DateTime(value=PydanticDateTime(2022, 1, 1, tzinfo=Timezone("UTC"))).save( + name="test" + ) api_block = DateTime.load("test") - assert api_block.value == pendulum.datetime(2022, 1, 1) + assert api_block.value == PydanticDateTime(2022, 1, 1, tzinfo=Timezone("UTC")) @pytest.mark.parametrize( @@ -18,7 +23,7 @@ def test_datetime(ignore_prefect_deprecation_warnings): ["test", {"key": "value"}, ["test"]], ids=["string", "dict", "list"], ) -def test_secret_block(value): +def test_secret_block(value: Any): Secret(value=value).save(name="test") api_block = Secret.load("test") assert isinstance(api_block.value, PydanticSecret) @@ -35,7 +40,7 @@ def test_secret_block(value): ], ids=["secret_string", "secret_dict", "secret_list"], ) -def test_secret_block_with_pydantic_secret(value): +def test_secret_block_with_pydantic_secret(value: Any): Secret(value=value).save(name="test") api_block = Secret.load("test") assert isinstance(api_block.value, PydanticSecret) diff --git a/tests/cli/deployment/test_deployment_run.py b/tests/cli/deployment/test_deployment_run.py index c40c950cde11..5ac1c7d8e8e9 100644 --- a/tests/cli/deployment/test_deployment_run.py +++ b/tests/cli/deployment/test_deployment_run.py @@ -5,16 +5,14 @@ from typing import Any, Generator from unittest.mock import ANY, AsyncMock -import pendulum import pytest -from pendulum.duration import Duration import prefect from prefect.client.schemas.objects import Deployment, FlowRun from prefect.exceptions import FlowRunWaitTimeout from prefect.states import Completed, Failed from prefect.testing.cli import invoke_and_assert -from prefect.types._datetime import DateTime, local_timezone +from prefect.types._datetime import DateTime, Duration, local_timezone, parse_datetime from prefect.utilities.asyncutils import run_sync_in_worker_thread @@ -134,11 +132,11 @@ async def test_start_at_option_displays_scheduled_start_time( @pytest.mark.parametrize( "start_at,expected_start_time", [ - ("5-17-43 5:30pm UTC", pendulum.parse("2043-05-17T17:30:00")), - ("5-20-2020 5:30pm EDT", pendulum.parse("2020-05-20T17:30:00", tz="EST5EDT")), - ("01/31/43 5:30 CST", pendulum.parse("2043-01-31T05:30:00", tz="CST6CDT")), - ("5-20-43 5:30pm PDT", pendulum.parse("2043-05-20T17:30:00", tz="PST8PDT")), - ("01/31/43 5:30 PST", pendulum.parse("2043-01-31T05:30:00", tz="PST8PDT")), + ("5-17-43 5:30pm UTC", parse_datetime("2043-05-17T17:30:00")), + ("5-20-2020 5:30pm EDT", parse_datetime("2020-05-20T17:30:00", tz="EST5EDT")), + ("01/31/43 5:30 CST", parse_datetime("2043-01-31T05:30:00", tz="CST6CDT")), + ("5-20-43 5:30pm PDT", parse_datetime("2043-05-20T17:30:00", tz="PST8PDT")), + ("01/31/43 5:30 PST", parse_datetime("2043-01-31T05:30:00", tz="PST8PDT")), ], ) async def test_start_at_option_with_tz_displays_scheduled_start_time( @@ -214,11 +212,11 @@ async def test_start_at_option_schedules_flow_run( @pytest.mark.parametrize( "start_at,expected_start_time", [ - ("5-17-43 5:30pm UTC", pendulum.parse("2043-05-17T17:30:00")), - ("5-20-2020 5:30pm EDT", pendulum.parse("2020-05-20T17:30:00", tz="EST5EDT")), - ("01/31/43 5:30 CST", pendulum.parse("2043-01-31T05:30:00", tz="CST6CDT")), - ("5-20-43 5:30pm PDT", pendulum.parse("2043-05-20T17:30:00", tz="PST8PDT")), - ("01/31/43 5:30 PST", pendulum.parse("2043-01-31T05:30:00", tz="PST8PDT")), + ("5-17-43 5:30pm UTC", parse_datetime("2043-05-17T17:30:00")), + ("5-20-2020 5:30pm EDT", parse_datetime("2020-05-20T17:30:00", tz="EST5EDT")), + ("01/31/43 5:30 CST", parse_datetime("2043-01-31T05:30:00", tz="CST6CDT")), + ("5-20-43 5:30pm PDT", parse_datetime("2043-05-20T17:30:00", tz="PST8PDT")), + ("01/31/43 5:30 PST", parse_datetime("2043-01-31T05:30:00", tz="PST8PDT")), ], ) async def test_start_at_option_with_tz_schedules_flow_run( @@ -227,7 +225,7 @@ async def test_start_at_option_with_tz_schedules_flow_run( expected_start_time: DateTime, prefect_client: prefect.client.orchestration.PrefectClient, ): - expected_start_time_local = expected_start_time.in_tz(pendulum.tz.local_timezone()) + expected_start_time_local = expected_start_time.in_tz(local_timezone()) expected_display = ( expected_start_time_local.to_datetime_string() + " " @@ -315,13 +313,13 @@ async def test_start_in_option_displays_scheduled_start_time( @pytest.mark.parametrize( "start_in,expected_duration", [ - ("10 minutes", pendulum.duration(minutes=10)), - ("5 days", pendulum.duration(days=5)), - ("3 seconds", pendulum.duration(seconds=3)), - (None, pendulum.duration(seconds=0)), - ("1 year and 3 months", pendulum.duration(years=1, months=3)), - ("2 weeks & 1 day", pendulum.duration(weeks=2, days=1)), - ("27 hours + 4 mins", pendulum.duration(days=1, hours=3, minutes=4)), + ("10 minutes", Duration(minutes=10)), + ("5 days", Duration(days=5)), + ("3 seconds", Duration(seconds=3)), + (None, Duration(seconds=0)), + ("1 year and 3 months", Duration(years=1, months=3)), + ("2 weeks & 1 day", Duration(weeks=2, days=1)), + ("27 hours + 4 mins", Duration(days=1, hours=3, minutes=4)), ], ) async def test_start_in_option_schedules_flow_run( @@ -332,9 +330,7 @@ async def test_start_in_option_schedules_flow_run( expected_duration: Duration, ): expected_start_time = frozen_now + expected_duration - expected_display = expected_start_time.in_tz( - pendulum.tz.local_timezone() - ).to_datetime_string() + expected_display = expected_start_time.in_tz(local_timezone()).to_datetime_string() await run_sync_in_worker_thread( invoke_and_assert, diff --git a/tests/cli/test_artifact.py b/tests/cli/test_artifact.py index 548b3801c8c8..555a11a54b8c 100644 --- a/tests/cli/test_artifact.py +++ b/tests/cli/test_artifact.py @@ -1,16 +1,20 @@ import sys +from typing import TYPE_CHECKING from uuid import uuid4 -import pendulum import pytest from typer import Exit from prefect.server import models, schemas from prefect.testing.cli import invoke_and_assert +from prefect.types._datetime import create_datetime_instance, human_friendly_diff + +if TYPE_CHECKING: + from sqlalchemy.ext.asyncio import AsyncSession @pytest.fixture(autouse=True) -def interactive_console(monkeypatch): +def interactive_console(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr("prefect.cli.artifact.is_interactive", lambda: True) # `readchar` does not like the fake stdin provided by typer isolation so we provide @@ -29,7 +33,7 @@ def readchar(): @pytest.fixture -async def artifact(session): +async def artifact(session: "AsyncSession"): artifact_schema = schemas.core.Artifact( key="voltaic", data={"a": 1}, type="table", description="opens many doors" ) @@ -45,7 +49,7 @@ async def artifact(session): @pytest.fixture -async def artifact_null_field(session): +async def artifact_null_field(session: "AsyncSession"): artifact_schema = schemas.core.Artifact( key="voltaic", data=1, metadata_={"description": "opens many doors"} ) @@ -61,7 +65,7 @@ async def artifact_null_field(session): @pytest.fixture -async def artifacts(session): +async def artifacts(session: "AsyncSession"): model1 = await models.artifacts.create_artifact( session=session, artifact=schemas.core.Artifact( @@ -98,37 +102,47 @@ def test_listing_artifacts_when_none_exist(): ) -def test_listing_artifacts_after_creating_artifacts(artifact): +def test_listing_artifacts_after_creating_artifacts( + artifact: models.artifacts.Artifact, +): + assert artifact.id is not None, "artifact id should not be None" + assert artifact.key is not None, "artifact key should not be None" + assert artifact.updated is not None, "artifact updated should not be None" + invoke_and_assert( ["artifact", "ls"], expected_output_contains=f""" ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓ ┃ ID ┃ Key ┃ Type ┃ Updated ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩ - │ {artifact.id} │ {artifact.key} │ {artifact.type} │ {pendulum.instance(artifact.updated).diff_for_humans()} │ + │ {artifact.id} │ {artifact.key} │ {artifact.type} │ {human_friendly_diff(create_datetime_instance(artifact.updated))} │ └──────────────────────────────────────┴─────────┴───────┴───────────────────┘ """, ) def test_listing_artifacts_after_creating_artifacts_with_null_fields( - artifact_null_field, + artifact_null_field: models.artifacts.Artifact, ): artifact = artifact_null_field + assert artifact.id is not None, "artifact id should not be None" + assert artifact.key is not None, "artifact key should not be None" + assert artifact.updated is not None, "artifact updated should not be None" + invoke_and_assert( ["artifact", "ls"], expected_output_contains=f""" ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━┳━━━━━━━━━━━━━━━━━━━┓ ┃ ID ┃ Key ┃ Type ┃ Updated ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━╇━━━━━━━━━━━━━━━━━━━┩ - │ {artifact.id} │ {artifact.key} │ │ {pendulum.instance(artifact.updated).diff_for_humans()} │ + │ {artifact.id} │ {artifact.key} │ │ {human_friendly_diff(create_datetime_instance(artifact.updated))} │ └──────────────────────────────────────┴─────────┴──────┴───────────────────┘ """, ) def test_listing_artifacts_with_limit( - artifacts, + artifacts: list[models.artifacts.Artifact], ): expected_output = artifacts[2].key invoke_and_assert( @@ -138,7 +152,9 @@ def test_listing_artifacts_with_limit( ) -def test_listing_artifacts_lists_only_latest_versions(artifacts): +def test_listing_artifacts_lists_only_latest_versions( + artifacts: list[models.artifacts.Artifact], +): expected_output = ( f"{artifacts[2].id}", f"{artifacts[1].id}", @@ -152,7 +168,9 @@ def test_listing_artifacts_lists_only_latest_versions(artifacts): ) -def test_listing_artifacts_with_all_set_to_true(artifacts): +def test_listing_artifacts_with_all_set_to_true( + artifacts: list[models.artifacts.Artifact], +): expected_output = ( f"{artifacts[0].id}", f"{artifacts[1].id}", @@ -166,7 +184,9 @@ def test_listing_artifacts_with_all_set_to_true(artifacts): ) -def test_listing_artifacts_with_all_set_to_false(artifacts): +def test_listing_artifacts_with_all_set_to_false( + artifacts: list[models.artifacts.Artifact], +): expected_output = ( f"{artifacts[2].id}", f"{artifacts[1].id}", @@ -180,7 +200,9 @@ def test_listing_artifacts_with_all_set_to_false(artifacts): ) -def test_inspecting_artifact_succeeds(artifacts): +def test_inspecting_artifact_succeeds( + artifacts: list[models.artifacts.Artifact], +): """ We expect to see all versions of the artifact. """ @@ -213,7 +235,9 @@ def test_inspecting_artifact_nonexistent_key_raises(): ) -def test_inspecting_artifact_with_limit(artifacts): +def test_inspecting_artifact_with_limit( + artifacts: list[models.artifacts.Artifact], +): expected_output = ( f"{artifacts[1].key}", f"{artifacts[1].data}", @@ -226,7 +250,9 @@ def test_inspecting_artifact_with_limit(artifacts): ) -def test_deleting_artifact_by_key_succeeds(artifacts): +def test_deleting_artifact_by_key_succeeds( + artifacts: list[models.artifacts.Artifact], +): invoke_and_assert( ["artifact", "delete", str(artifacts[0].key)], prompts_and_responses=[ @@ -247,7 +273,9 @@ def test_deleting_artifact_nonexistent_key_raises(): ) -def test_deleting_artifact_by_key_without_confimation_aborts(artifacts): +def test_deleting_artifact_by_key_without_confimation_aborts( + artifacts: list[models.artifacts.Artifact], +): invoke_and_assert( ["artifact", "delete", str(artifacts[0].key)], user_input="n", @@ -256,7 +284,9 @@ def test_deleting_artifact_by_key_without_confimation_aborts(artifacts): ) -def test_deleting_artifact_by_id_succeeds(artifacts): +def test_deleting_artifact_by_id_succeeds( + artifacts: list[models.artifacts.Artifact], +): invoke_and_assert( ["artifact", "delete", "--id", str(artifacts[0].id)], user_input="y", @@ -275,7 +305,9 @@ def test_deleting_artifact_nonexistent_id_raises(): ) -def test_deleting_artifact_by_id_without_confimation_aborts(artifacts): +def test_deleting_artifact_by_id_without_confimation_aborts( + artifacts: list[models.artifacts.Artifact], +): invoke_and_assert( ["artifact", "delete", "--id", str(artifacts[0].id)], user_input="n", @@ -284,9 +316,13 @@ def test_deleting_artifact_by_id_without_confimation_aborts(artifacts): ) -def test_deleting_artifact_with_key_and_id_raises(artifacts): +def test_deleting_artifact_with_key_and_id_raises( + artifacts: list[models.artifacts.Artifact], +): + assert artifacts[1].key is not None, "artifact key should not be None" + assert artifacts[1].id is not None, "artifact id should not be None" invoke_and_assert( - ["artifact", "delete", artifacts[1].key, "--id", artifacts[1].id], + ["artifact", "delete", artifacts[1].key, "--id", str(artifacts[1].id)], expected_output_contains=( "Please provide either a key or an artifact_id but not both." ), @@ -294,7 +330,9 @@ def test_deleting_artifact_with_key_and_id_raises(artifacts): ) -def test_deleting_artifact_without_key_or_id_raises(artifacts): +def test_deleting_artifact_without_key_or_id_raises( + artifacts: list[models.artifacts.Artifact], +): invoke_and_assert( ["artifact", "delete"], expected_output_contains="Please provide a key or an artifact_id.", diff --git a/tests/cli/test_deploy.py b/tests/cli/test_deploy.py index 988820ca5e6e..4051d41461d7 100644 --- a/tests/cli/test_deploy.py +++ b/tests/cli/test_deploy.py @@ -9,11 +9,10 @@ import tempfile from datetime import timedelta from pathlib import Path -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from unittest import mock from uuid import UUID, uuid4 -import pendulum import pytest import readchar import yaml @@ -58,14 +57,18 @@ ) from prefect.testing.cli import invoke_and_assert from prefect.testing.utilities import AsyncMock +from prefect.types._datetime import parse_datetime from prefect.utilities.asyncutils import run_sync_in_worker_thread from prefect.utilities.filesystem import tmpchdir +if TYPE_CHECKING: + from prefect.server.schemas.core import BlockDocument + TEST_PROJECTS_DIR = prefect.__development_base_path__ / "tests" / "test-projects" @pytest.fixture -def interactive_console(monkeypatch): +def interactive_console(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr("prefect.cli.deploy.is_interactive", lambda: True) # `readchar` does not like the fake stdin provided by typer isolation so we provide @@ -84,8 +87,8 @@ def readchar(): @pytest.fixture -def project_dir(tmp_path): - with tmpchdir(tmp_path): +def project_dir(tmp_path: Path): + with tmpchdir(str(tmp_path)): shutil.copytree(TEST_PROJECTS_DIR, tmp_path, dirs_exist_ok=True) prefect_home = tmp_path / ".prefect" prefect_home.mkdir(exist_ok=True, mode=0o0700) @@ -94,8 +97,8 @@ def project_dir(tmp_path): @pytest.fixture -def project_dir_with_single_deployment_format(tmp_path): - with tmpchdir(tmp_path): +def project_dir_with_single_deployment_format(tmp_path: Path): + with tmpchdir(str(tmp_path)): shutil.copytree(TEST_PROJECTS_DIR, tmp_path, dirs_exist_ok=True) prefect_home = tmp_path / ".prefect" prefect_home.mkdir(exist_ok=True, mode=0o0700) @@ -113,13 +116,13 @@ def project_dir_with_single_deployment_format(tmp_path): @pytest.fixture -def uninitialized_project_dir(project_dir): +def uninitialized_project_dir(project_dir: Path): Path(project_dir, "prefect.yaml").unlink() return project_dir @pytest.fixture -def uninitialized_project_dir_with_git_no_remote(uninitialized_project_dir): +def uninitialized_project_dir_with_git_no_remote(uninitialized_project_dir: Path): subprocess.run(["git", "init"], cwd=uninitialized_project_dir) assert Path(uninitialized_project_dir, ".git").exists() return uninitialized_project_dir @@ -127,7 +130,7 @@ def uninitialized_project_dir_with_git_no_remote(uninitialized_project_dir): @pytest.fixture def uninitialized_project_dir_with_git_with_remote( - uninitialized_project_dir_with_git_no_remote, + uninitialized_project_dir_with_git_no_remote: Path, ): subprocess.run( ["git", "remote", "add", "origin", "https://example.com/org/repo.git"], @@ -137,7 +140,7 @@ def uninitialized_project_dir_with_git_with_remote( @pytest.fixture -async def default_agent_pool(prefect_client): +async def default_agent_pool(prefect_client: PrefectClient) -> WorkPool: try: return await prefect_client.create_work_pool( WorkPoolCreate(name="default-agent-pool", type="prefect-agent") @@ -169,9 +172,9 @@ async def docker_work_pool(prefect_client: PrefectClient) -> WorkPool: @pytest.fixture -async def mock_prompt(monkeypatch): +async def mock_prompt(monkeypatch: pytest.MonkeyPatch): # Mock prompts() where password=True to prevent hanging - def new_prompt(message, password=False, **kwargs): + def new_prompt(message: str, password: bool = False, **kwargs: Any) -> str: if password: return "456" else: @@ -182,8 +185,8 @@ def new_prompt(message, password=False, **kwargs): @pytest.fixture -def mock_provide_password(monkeypatch): - def new_prompt(message, password=False, **kwargs): +def mock_provide_password(monkeypatch: pytest.MonkeyPatch): + def new_prompt(message: str, password: bool = False, **kwargs: Any) -> str: if password: return "my-token" else: @@ -194,7 +197,7 @@ def new_prompt(message, password=False, **kwargs): @pytest.fixture -def mock_build_docker_image(monkeypatch): +def mock_build_docker_image(monkeypatch: pytest.MonkeyPatch): mock_build = mock.MagicMock() mock_build.return_value = {"build-image": {"image": "{{ build-image.image }}"}} @@ -211,7 +214,7 @@ def mock_build_docker_image(monkeypatch): @pytest.fixture -async def aws_credentials(prefect_client): +async def aws_credentials(prefect_client: PrefectClient) -> "BlockDocument": aws_credentials_type = await prefect_client.create_block_type( block_type=BlockTypeCreate( name="AWS Credentials", @@ -256,7 +259,7 @@ def uninitialized_project_dir_with_git_no_remote(self, uninitialized_project_dir @pytest.fixture def uninitialized_project_dir_with_git_with_remote( - self, uninitialized_project_dir_with_git_no_remote + self, uninitialized_project_dir_with_git_no_remote: Path ): subprocess.run( ["git", "remote", "add", "origin", "https://example.com/org/repo.git"], @@ -264,7 +267,9 @@ def uninitialized_project_dir_with_git_with_remote( ) return uninitialized_project_dir_with_git_no_remote - async def test_project_deploy(self, project_dir, prefect_client: PrefectClient): + async def test_project_deploy( + self, project_dir: Path, prefect_client: PrefectClient + ): await prefect_client.create_work_pool( WorkPoolCreate(name="test-pool", type="test") ) @@ -310,7 +315,11 @@ async def test_deploy_with_no_enforce_parameter_schema( assert not deployment.enforce_parameter_schema async def test_deploy_with_active_workers( - self, project_dir, work_pool, prefect_client, monkeypatch + self, + project_dir: Path, + work_pool: WorkPool, + prefect_client: PrefectClient, + monkeypatch: pytest.MonkeyPatch, ): mock_read_workers_for_work_pool = AsyncMock( return_value=[ @@ -337,7 +346,7 @@ async def test_deploy_with_active_workers( ) async def test_deploy_with_wrapped_flow_decorator( - self, project_dir, work_pool, prefect_client + self, project_dir: Path, work_pool: WorkPool, prefect_client: PrefectClient ): await run_sync_in_worker_thread( invoke_and_assert, @@ -359,7 +368,7 @@ async def test_deploy_with_wrapped_flow_decorator( assert deployment.work_pool_name == work_pool.name async def test_deploy_with_missing_imports( - self, project_dir, work_pool, prefect_client + self, project_dir: Path, work_pool: WorkPool, prefect_client: PrefectClient ): await run_sync_in_worker_thread( invoke_and_assert, @@ -381,7 +390,7 @@ async def test_deploy_with_missing_imports( assert deployment.work_pool_name == work_pool.name async def test_project_deploy_with_default_work_pool( - self, project_dir, prefect_client + self, project_dir: Path, prefect_client: PrefectClient ): await prefect_client.create_work_pool( WorkPoolCreate(name="test-pool", type="test") @@ -411,7 +420,7 @@ async def test_project_deploy_with_default_work_pool( assert deployment.enforce_parameter_schema async def test_project_deploy_with_no_deployment_file( - self, project_dir, prefect_client: PrefectClient + self, project_dir: Path, prefect_client: PrefectClient ): await prefect_client.create_work_pool( WorkPoolCreate(name="test-pool", type="test") @@ -436,7 +445,9 @@ async def test_project_deploy_with_no_deployment_file( assert deployment.job_variables == {"env": "prod"} assert deployment.enforce_parameter_schema is True - async def test_project_deploy_with_no_prefect_yaml(self, project_dir, work_pool): + async def test_project_deploy_with_no_prefect_yaml( + self, project_dir: Path, work_pool: WorkPool + ): Path(project_dir, "prefect.yaml").unlink() await run_sync_in_worker_thread( @@ -453,8 +464,11 @@ async def test_project_deploy_with_no_prefect_yaml(self, project_dir, work_pool) ], ) + @pytest.mark.usefixtures("interactive_console") async def test_deploy_does_not_prompt_storage_when_pull_step_exists( - self, project_dir, work_pool, interactive_console + self, + project_dir: Path, + work_pool: WorkPool, ): # write a pull step to the prefect.yaml with open("prefect.yaml", "r") as f: @@ -506,11 +520,11 @@ async def test_deploy_does_not_prompt_storage_when_pull_step_exists( @pytest.mark.usefixtures("interactive_console", "uninitialized_project_dir") async def test_deploy_with_concurrency_limit_and_options( self, - project_dir, + project_dir: Path, prefect_client: PrefectClient, - cli_options, - expected_limit, - expected_strategy, + cli_options: str, + expected_limit: Optional[int], + expected_strategy: Optional[str], ): await prefect_client.create_work_pool( WorkPoolCreate(name="test-pool", type="test") @@ -577,7 +591,10 @@ async def test_deploy_with_concurrency_limit_and_options( class TestGeneratedPullAction: async def test_project_deploy_generates_pull_action( - self, work_pool, prefect_client, uninitialized_project_dir + self, + work_pool: WorkPool, + prefect_client: PrefectClient, + uninitialized_project_dir: Path, ): await run_sync_in_worker_thread( invoke_and_assert, @@ -601,9 +618,9 @@ async def test_project_deploy_generates_pull_action( async def test_project_deploy_with_no_prefect_yaml_git_repo_no_remote( self, - work_pool, - prefect_client, - uninitialized_project_dir_with_git_no_remote, + work_pool: WorkPool, + prefect_client: PrefectClient, + uninitialized_project_dir_with_git_no_remote: Path, ): await run_sync_in_worker_thread( invoke_and_assert, @@ -629,9 +646,9 @@ async def test_project_deploy_with_no_prefect_yaml_git_repo_no_remote( @pytest.mark.usefixtures("interactive_console") async def test_project_deploy_with_no_prefect_yaml_git_repo_user_rejects( self, - work_pool, - prefect_client, - uninitialized_project_dir_with_git_with_remote, + work_pool: WorkPool, + prefect_client: PrefectClient, + uninitialized_project_dir_with_git_with_remote: Path, ): await run_sync_in_worker_thread( invoke_and_assert, @@ -661,7 +678,9 @@ async def test_project_deploy_with_no_prefect_yaml_git_repo_user_rejects( "interactive_console", "uninitialized_project_dir_with_git_with_remote" ) async def test_project_deploy_with_no_prefect_yaml_git_repo( - self, work_pool, prefect_client + self, + work_pool: WorkPool, + prefect_client: PrefectClient, ): await run_sync_in_worker_thread( invoke_and_assert, @@ -723,7 +742,9 @@ async def test_project_deploy_with_no_prefect_yaml_git_repo( "interactive_console", "uninitialized_project_dir_with_git_with_remote" ) async def test_project_deploy_with_no_prefect_yaml_git_repo_user_overrides( - self, work_pool, prefect_client + self, + work_pool: WorkPool, + prefect_client: PrefectClient, ): await run_sync_in_worker_thread( invoke_and_assert, @@ -788,8 +809,8 @@ async def test_project_deploy_with_no_prefect_yaml_git_repo_user_overrides( ) async def test_project_deploy_with_no_prefect_yaml_git_repo_with_token( self, - work_pool, - prefect_client, + work_pool: WorkPool, + prefect_client: PrefectClient, ): await run_sync_in_worker_thread( invoke_and_assert, @@ -851,10 +872,10 @@ async def test_project_deploy_with_no_prefect_yaml_git_repo_with_token( @pytest.mark.usefixtures("interactive_console", "uninitialized_project_dir") async def test_deploy_with_blob_storage_select_existing_credentials( self, - work_pool, - prefect_client, - aws_credentials, - monkeypatch, + work_pool: WorkPool, + prefect_client: PrefectClient, + aws_credentials: "BlockDocument", + monkeypatch: pytest.MonkeyPatch, ): mock_step = mock.MagicMock() monkeypatch.setattr( @@ -2505,11 +2526,15 @@ async def test_passing_interval_schedules_to_deploy( assert len(deployment.schedules) == 1 schedule = deployment.schedules[0].schedule assert schedule.interval == timedelta(seconds=42) - assert schedule.anchor_date == pendulum.parse("2040-02-02") + assert schedule.anchor_date == parse_datetime("2040-02-02") assert schedule.timezone == "America/New_York" @pytest.mark.usefixtures("project_dir") - async def test_interval_schedule_deployment_yaml(self, prefect_client, work_pool): + async def test_interval_schedule_deployment_yaml( + self, + prefect_client: PrefectClient, + work_pool: WorkPool, + ): prefect_yaml = Path("prefect.yaml") with prefect_yaml.open(mode="r") as f: deploy_config = yaml.safe_load(f) @@ -2540,7 +2565,7 @@ async def test_interval_schedule_deployment_yaml(self, prefect_client, work_pool assert len(deployment.schedules) == 1 schedule = deployment.schedules[0].schedule assert schedule.interval == timedelta(seconds=42) - assert schedule.anchor_date == pendulum.parse("2040-02-02") + assert schedule.anchor_date == parse_datetime("2040-02-02") assert schedule.timezone == "America/Chicago" assert deployment.schedules[0].parameters == {"number": 42} assert deployment.schedules[0].slug == "test-slug" @@ -2731,7 +2756,7 @@ async def test_can_provide_multiple_schedules_of_the_same_type_via_command( "An important name/test-name" ) - schedules = set() + schedules: set[str] = set() for deployment_schedule in deployment.schedules: schedule = deployment_schedule.schedule assert isinstance(schedule, CronSchedule) diff --git a/tests/cli/test_version.py b/tests/cli/test_version.py index a572b34137bd..dfeb8dcac89f 100644 --- a/tests/cli/test_version.py +++ b/tests/cli/test_version.py @@ -4,7 +4,6 @@ from textwrap import dedent from unittest.mock import Mock -import pendulum import pydantic import pytest @@ -17,9 +16,11 @@ temporary_settings, ) from prefect.testing.cli import invoke_and_assert +from prefect.types._datetime import parse_datetime -def test_version_ephemeral_server_type(disable_hosted_api_server): +@pytest.mark.usefixtures("disable_hosted_api_server") +def test_version_ephemeral_server_type(): with temporary_settings( { PREFECT_SERVER_ALLOW_EPHEMERAL_MODE: True, @@ -30,7 +31,8 @@ def test_version_ephemeral_server_type(disable_hosted_api_server): ) -def test_version_unconfigured_server_type(disable_hosted_api_server): +@pytest.mark.usefixtures("disable_hosted_api_server") +def test_version_unconfigured_server_type(): invoke_and_assert( ["version"], expected_output_contains="Server type: unconfigured" ) @@ -56,9 +58,11 @@ def test_version_cloud_server_type(): ) -def test_correct_output_ephemeral_sqlite(monkeypatch, disable_hosted_api_server): +@pytest.mark.usefixtures("disable_hosted_api_server") +def test_correct_output_ephemeral_sqlite(monkeypatch: pytest.MonkeyPatch): version_info = prefect.__version_info__ - built = pendulum.parse(prefect.__version_info__["date"]) + assert version_info["date"] is not None, "date is not set" + built = parse_datetime(version_info["date"]) profile = prefect.context.get_settings_context().profile dialect = Mock() @@ -91,9 +95,11 @@ def test_correct_output_ephemeral_sqlite(monkeypatch, disable_hosted_api_server) ) -def test_correct_output_ephemeral_postgres(monkeypatch, disable_hosted_api_server): +@pytest.mark.usefixtures("disable_hosted_api_server") +def test_correct_output_ephemeral_postgres(monkeypatch: pytest.MonkeyPatch): version_info = prefect.__version_info__ - built = pendulum.parse(prefect.__version_info__["date"]) + assert version_info["date"] is not None, "date is not set" + built = parse_datetime(version_info["date"]) profile = prefect.context.get_settings_context().profile dialect = Mock() @@ -128,7 +134,8 @@ def test_correct_output_ephemeral_postgres(monkeypatch, disable_hosted_api_serve @pytest.mark.usefixtures("use_hosted_api_server") def test_correct_output_non_ephemeral_server_type(): version_info = prefect.__version_info__ - built = pendulum.parse(prefect.__version_info__["date"]) + assert version_info["date"] is not None, "date is not set" + built = parse_datetime(version_info["date"]) profile = prefect.context.get_settings_context().profile invoke_and_assert( diff --git a/tests/conftest.py b/tests/conftest.py index bfa87af2a60b..4af02ba8f1f6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -66,6 +66,7 @@ PREFECT_SERVER_LOGGING_LEVEL, PREFECT_UNIT_TEST_LOOP_DEBUG, ) +from prefect.types._datetime import DateTime, now from prefect.utilities.dispatch import get_registry_for_type # isort: split @@ -144,7 +145,9 @@ def pytest_addoption(parser): ] -def pytest_collection_modifyitems(session, config, items): +def pytest_collection_modifyitems( + session: pytest.Session, config: pytest.Config, items: list[pytest.Item] +): """ Update tests to skip in accordance with service requests """ @@ -268,7 +271,7 @@ def tests_dir() -> pathlib.Path: @pytest.hookimpl(hookwrapper=True) -def pytest_runtest_call(item): +def pytest_runtest_call(item: pytest.Item): """ This hook will be called within the test run. Allowing us to raise errors or add assertions to every test. On error, the test will be marked as failed. If we used @@ -303,7 +306,7 @@ def assert_lifespan_is_not_left_open(): TEST_PROFILE_CTX = None -def pytest_sessionstart(session): +def pytest_sessionstart(session: pytest.Session): """ Creates a profile for the scope of the test session that modifies setting defaults. @@ -378,7 +381,7 @@ def pytest_sessionstart(session): # def pytest_sessionfinish(session, exitstatus): @pytest.fixture(scope="session", autouse=True) -def cleanup(drain_log_workers, drain_events_workers): +def cleanup(drain_log_workers: None, drain_events_workers: None): # this fixture depends on other fixtures with important cleanup steps like # draining workers to ensure that the home directory is not deleted before # these steps are completed @@ -519,7 +522,9 @@ def reset_registered_blocks(): @pytest.fixture -def caplog(caplog): +def caplog( + caplog: pytest.LogCaptureFixture, +) -> Generator[pytest.LogCaptureFixture, None, None]: """ Overrides caplog to apply to all of our loggers that do not propagate and consequently would not be captured by caplog. diff --git a/tests/test_flows.py b/tests/test_flows.py index cca83f0bb138..1ae71c253ee9 100644 --- a/tests/test_flows.py +++ b/tests/test_flows.py @@ -13,12 +13,11 @@ from itertools import combinations from pathlib import Path from textwrap import dedent -from typing import List, Optional +from typing import Any, List, Optional from unittest import mock from unittest.mock import ANY, MagicMock, call, create_autospec import anyio -import pendulum import pydantic import pytest import regex as re @@ -87,6 +86,7 @@ get_most_recent_flow_run, ) from prefect.transactions import get_transaction, transaction +from prefect.types._datetime import DateTime, Timezone from prefect.types.entrypoint import EntrypointType from prefect.utilities.annotations import allow_failure, quote from prefect.utilities.callables import parameter_schema @@ -103,7 +103,7 @@ def mock_sigterm_handler(): pytest.skip("Can't test signal handlers from a thread") mock = MagicMock() - def handler(*args, **kwargs): + def handler(*args: Any, **kwargs: Any): mock(*args, **kwargs) prev_handler = signal.signal(signal.SIGTERM, handler) @@ -5555,14 +5555,14 @@ def flow_function(name: str) -> str: def test_annotations_and_defaults_rely_on_imports(self, tmp_path: Path): source_code = dedent( """ - import pendulum import datetime from prefect import flow + from prefect.types import DateTime @flow def f( x: datetime.datetime, - y: pendulum.DateTime = pendulum.datetime(2025, 1, 1), + y: DateTime = DateTime(2025, 1, 1), z: datetime.timedelta = datetime.timedelta(seconds=5), ): return x, y, z @@ -5573,7 +5573,7 @@ def f( assert result is not None assert result(datetime.datetime(2025, 1, 1)) == ( datetime.datetime(2025, 1, 1), - pendulum.datetime(2025, 1, 1), + DateTime(2025, 1, 1, tzinfo=Timezone("UTC")), datetime.timedelta(seconds=5), ) diff --git a/tests/test_logging.py b/tests/test_logging.py index 0d34331977ec..63d13d35b6af 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -10,7 +10,6 @@ from unittest import mock from unittest.mock import ANY, MagicMock -import pendulum import pytest from rich.color import Color, ColorType from rich.console import Console @@ -68,6 +67,7 @@ ) from prefect.testing.cli import temporary_console_width from prefect.testing.utilities import AsyncMock +from prefect.types._datetime import from_timestamp, now from prefect.utilities.names import obfuscate from prefect.workers.base import BaseJobConfiguration, BaseWorker @@ -464,8 +464,7 @@ def test_sets_timestamp_from_record_created_time( log_dict = mock_log_worker.instance().send.call_args[0][0] assert ( - log_dict["timestamp"] - == pendulum.from_timestamp(record.created).to_iso8601_string() + log_dict["timestamp"] == from_timestamp(record.created).to_iso8601_string() ) def test_sets_timestamp_from_time_if_missing_from_recrod( @@ -487,7 +486,7 @@ def drop_created_and_emit(emit, record): log_dict = mock_log_worker.instance().send.call_args[0][0] - assert log_dict["timestamp"] == pendulum.from_timestamp(now).to_iso8601_string() + assert log_dict["timestamp"] == from_timestamp(now).to_iso8601_string() def test_does_not_send_logs_that_opt_out(self, logger, mock_log_worker, task_run): with TaskRunContext.model_construct(task_run=task_run): @@ -821,7 +820,7 @@ def log_dict(self): task_run_id=uuid.uuid4(), name="test.logger", level=10, - timestamp=pendulum.now("utc"), + timestamp=now("utc"), message="hello", ).model_dump(mode="json")