diff --git a/.gitignore b/.gitignore index bf82b7dbd9e..95051f47f25 100644 --- a/.gitignore +++ b/.gitignore @@ -205,9 +205,10 @@ examples/computer_vision/test_runs #generated folder by zenml zenml_tutorial/ -# Claude settings +# AI code tools **/.claude/settings.local.json .claude/ +.repoprompt/ .local/ # PLEASE KEEP THIS LINE AT THE EOF: never include here src/zenml/zen_server/dashboard, since it is affecting release flow \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md index db10e31b180..4bbbcb3c31a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -67,6 +67,12 @@ Note: The MCP server indexes the latest released docs, not the develop branch. F - For full coverage, use CI (see CI section below) - Some tests use: `bash scripts/test-coverage-xml.sh` (but this won't run all tests) +#### CLI Testing Best Practices +- When writing new CLI tests, use the `cli_runner()` helper from `tests/integration/functional/cli/utils.py` instead of directly instantiating `CliRunner()` +- The helper handles Click version compatibility (Click 8.2+ removed the `mix_stderr` parameter) +- Always check `result.exit_code` after invoking CLI commands to catch failures early +- Existing tests with `CliRunner()` (no arguments) are fine - only new tests need the helper + ## Development Workflow ### Prerequisites diff --git a/src/zenml/cli/base.py b/src/zenml/cli/base.py index d80f428b92d..5256e49d0da 100644 --- a/src/zenml/cli/base.py +++ b/src/zenml/cli/base.py @@ -381,6 +381,8 @@ def clean(yes: bool = False, local: bool = False) -> None: global_zen_config = Path(get_global_config_directory()) if fileio.exists(str(global_zen_config)): gc = GlobalConfiguration() + # Ensure any open DB handles (e.g., SQLite) are released before deletion + gc.close_store() for dir_name in fileio.listdir(str(global_zen_config)): if fileio.isdir(str(global_zen_config / str(dir_name))): warning( diff --git a/src/zenml/cli/cli.py b/src/zenml/cli/cli.py index 01346581782..7eb99cee703 100644 --- a/src/zenml/cli/cli.py +++ b/src/zenml/cli/cli.py @@ -137,17 +137,22 @@ def format_commands( help_ = cmd.get_short_help_str(limit=formatter.width) rows.append((tag.value, subcommand, help_)) if rows: - colored_section_title = ( + section_title = ( "[dim cyan]Available ZenML Commands (grouped)[/dim cyan]" ) - with formatter.section(colored_section_title): + with formatter.section(section_title): formatter.write_dl(rows) # type: ignore[arg-type] -@click.group(cls=ZenMLCLI) +@click.group(cls=ZenMLCLI, invoke_without_command=True) @click.version_option(__version__, "--version", "-v") -def cli() -> None: - """CLI base command for ZenML.""" +@click.pass_context +def cli(ctx: click.Context) -> None: + """CLI base command for ZenML. + + Args: + ctx: The click context. + """ set_root_verbosity() source_context.set(SourceContextTypes.CLI) repo_root = Client.find_repository() @@ -158,6 +163,17 @@ def cli() -> None: # packages directory source_utils.set_custom_source_root(source_root=os.getcwd()) + # Manually show help when invoked without a subcommand to ensure our custom + # formatter is used. Without invoke_without_command=True, Click defaults to + # no_args_is_help=True, which in Click 8.2+ raises NoArgsIsHelpError before + # this callback runs. That error triggers Click's default help display, + # bypassing ZenMLCLI.get_help() and losing our custom formatting (categories, + # colors, layout). By using invoke_without_command=True and manually checking + # for no subcommand, we guarantee consistent behavior across Click versions. + if ctx.invoked_subcommand is None and not ctx.resilient_parsing: + ctx.command.get_help(ctx) + ctx.exit(0) + if __name__ == "__main__": cli() diff --git a/src/zenml/config/global_config.py b/src/zenml/config/global_config.py index 4cc8d17fb3b..d764d7e0bf0 100644 --- a/src/zenml/config/global_config.py +++ b/src/zenml/config/global_config.py @@ -170,6 +170,17 @@ def _reset_instance( singleton. If None, the global GlobalConfiguration singleton is reset to an empty value. """ + # Proactively dispose resources held by the current singleton instance + # to avoid file-handle issues on platforms like Windows. + if cls._global_config is not None: + try: + cls._global_config.close_store() + except Exception: + logger.debug( + "Failed to close store during GlobalConfiguration reset.", + exc_info=True, + ) + cls._global_config = config if config: config._write_config() @@ -701,6 +712,29 @@ def set_store( self._configure_store(config, skip_default_registrations, **kwargs) logger.info("Updated the global store configuration.") + def close_store(self) -> None: + """Close and detach the cached store if it has been initialized. + + Important: + - Operates only on the cached store; does not instantiate a new one. + - Logs cleanup errors at debug level and never raises. + """ + store = self._zen_store + if store is None: + return + + try: + store.close() + except Exception as e: + logger.debug( + "Error while closing zen store during cleanup: %s", + e, + exc_info=True, + ) + finally: + # Clear the cached store reference irrespective of cleanup result. + self._zen_store = None + @property def is_initialized(self) -> bool: """Check if the global configuration is initialized. diff --git a/src/zenml/zen_stores/base_zen_store.py b/src/zenml/zen_stores/base_zen_store.py index 1f369a6e817..614993a1115 100644 --- a/src/zenml/zen_stores/base_zen_store.py +++ b/src/zenml/zen_stores/base_zen_store.py @@ -141,6 +141,18 @@ def __init__( else: logger.debug("Skipping database initialization") + def close(self) -> None: + """Release external resources held by this store instance. + + Default no-op implementation. Subclasses should override this to + release any external resources (e.g., database connections, HTTP + sessions, file handles, background workers). Implementations must + remain idempotent (safe to call multiple times) and must not raise; + any cleanup failures should be handled internally and logged at + debug level instead. + """ + return + @staticmethod def get_store_class(store_type: StoreType) -> Type["BaseZenStore"]: """Returns the class of the given store type. diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 24a0e50b3c4..1a276e562f1 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -4344,6 +4344,40 @@ def batch_delete_tag_resource( route=TAG_RESOURCES, ) + def close(self) -> None: + """Release external resources held by this store instance. + + This implementation closes the underlying HTTP session and clears transient + authentication state. It is safe to call multiple times and never raises; + unexpected cleanup issues are logged at debug level instead. + """ + try: + with self._session_lock: + if self._session is not None: + try: + self._session.close() + except Exception as e: + # Cleanup must not raise; surface at debug level only. + logger.debug( + "Error closing REST session during RestZenStore.close(): %s", + e, + exc_info=True, + ) + finally: + # Clear the session reference to avoid reuse and help GC. + self._session = None + + # Clear transient auth state so future use starts cleanly. + self._api_token = None + self._last_authenticated = None + except Exception as e: + # Absolute safety: no exceptions should escape close() + logger.debug( + "Unexpected error during RestZenStore.close(): %s", + e, + exc_info=True, + ) + # ======================= # Internal helper methods # ======================= diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 451007e05eb..3fd2b8795df 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -1248,6 +1248,30 @@ def filter_and_paginate( # Initialization and configuration # -------------------------------- + def close(self) -> None: + """Release DB resources held by this store. + + Safe to call multiple times. Primarily ensures SQLite file handles are + released so the DB file can be removed on platforms like Windows. + """ + try: + engine = self._engine + if engine is not None: + # Disposes pooled connections and releases file handles. + engine.dispose() + except Exception as e: + # Cleanup must not raise; surface at debug level only. + logger.debug( + "Error disposing SQLAlchemy engine during SqlZenStore.close(): %s", + e, + exc_info=True, + ) + finally: + # Clear references to avoid accidental reuse and help GC. + self._engine = None + self._alembic = None + self._migration_utils = None + def _initialize(self) -> None: """Initialize the SQL store.""" logger.debug("Initializing SqlZenStore at %s", self.config.url) diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index 350a74c0387..b79718b5930 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -234,6 +234,19 @@ def _initialize(self) -> None: be used to set up the backend (database, connection etc.). """ + @abstractmethod + def close(self) -> None: + """Release external resources held by this store instance. + + Implementations must: + - Release any external resources such as database connections, HTTP + sessions, open file handles, background workers, etc. + - Be safe to call multiple times (idempotent). Subsequent calls should + be no-ops and must not fail because resources were already released. + - Never raise exceptions. Any cleanup failures should be handled + internally and logged at debug level instead of being propagated. + """ + @abstractmethod def get_store_info(self) -> ServerModel: """Get information about the store. diff --git a/tests/integration/functional/cli/test_base.py b/tests/integration/functional/cli/test_base.py index 7217a227712..84fb90669ee 100644 --- a/tests/integration/functional/cli/test_base.py +++ b/tests/integration/functional/cli/test_base.py @@ -17,8 +17,8 @@ from pathlib import Path import pytest -from click.testing import CliRunner +from tests.integration.functional.cli.utils import cli_runner from zenml.cli.base import ZENML_PROJECT_TEMPLATES, clean, init from zenml.constants import CONFIG_FILE_NAME, REPOSITORY_DIRECTORY_NAME from zenml.utils import yaml_utils @@ -27,8 +27,9 @@ def test_init_creates_zen_folder(tmp_path: Path) -> None: """Check that init command creates a .zen folder.""" - runner = CliRunner() - runner.invoke(init, ["--path", str(tmp_path)]) + runner = cli_runner() + result = runner.invoke(init, ["--path", str(tmp_path), "--test"]) + assert result.exit_code == 0, f"Command failed: {result.output}" assert (tmp_path / REPOSITORY_DIRECTORY_NAME).exists() @@ -44,8 +45,8 @@ def test_init_creates_from_templates( tmp_path: Path, template_name: str ) -> None: """Check that init command checks-out template.""" - runner = CliRunner() - runner.invoke( + runner = cli_runner() + result = runner.invoke( init, [ "--path", @@ -53,8 +54,10 @@ def test_init_creates_from_templates( "--template", template_name, "--template-with-defaults", + "--test", ], ) + assert result.exit_code == 0, f"Command failed: {result.output}" assert (tmp_path / REPOSITORY_DIRECTORY_NAME).exists() files_in_top_level = set([f.lower() for f in os.listdir(str(tmp_path))]) must_have_files = { @@ -75,9 +78,40 @@ def test_clean_user_config(clean_client) -> None: user_id = yaml_contents["user_id"] analytics_opt_in = yaml_contents["analytics_opt_in"] version = yaml_contents["version"] - runner = CliRunner() - runner.invoke(clean, ["--yes", True]) + runner = cli_runner() + result = runner.invoke(clean, ["--yes"]) + assert result.exit_code == 0, f"Command failed: {result.output}" new_yaml_contents = yaml_utils.read_yaml(str(global_zen_config_yaml)) assert user_id == new_yaml_contents["user_id"] assert analytics_opt_in == new_yaml_contents["analytics_opt_in"] assert version == new_yaml_contents["version"] + + +def test_clean_calls_store_close(monkeypatch, clean_client) -> None: + # Import locally to avoid modifying module-level imports + from zenml.config.global_config import GlobalConfiguration + from zenml.zen_stores.sql_zen_store import SqlZenStore + + # Track whether close was called + called = {"closed": False} + + # Keep reference to the original close to preserve behavior + orig_close = SqlZenStore.close + + def tracked_close(self): + called["closed"] = True + # Call through to the original to keep cleanup behavior intact + return orig_close(self) + + # Patch the close method + monkeypatch.setattr(SqlZenStore, "close", tracked_close, raising=True) + + # Ensure the store is initialized in this process so there is something to close + _ = GlobalConfiguration().zen_store + + # Run clean + result = cli_runner().invoke(clean, ["--yes"]) + assert result.exit_code == 0, f"Command failed: {result.output}" + + # Verify that our close hook was executed + assert called["closed"] is True diff --git a/tests/integration/functional/cli/test_cli.py b/tests/integration/functional/cli/test_cli.py index 010f442c928..8da0c6d152d 100644 --- a/tests/integration/functional/cli/test_cli.py +++ b/tests/integration/functional/cli/test_cli.py @@ -16,15 +16,17 @@ import click import pytest -from click.testing import CliRunner +from tests.integration.functional.cli.utils import ( + cli_runner as create_cli_runner, +) from zenml.cli.cli import ZenMLCLI, cli from zenml.cli.formatter import ZenFormatter @pytest.fixture(scope="function") def runner(request): - return CliRunner() + return create_cli_runner() def test_cli_command_defines_a_cli_group() -> None: @@ -57,7 +59,7 @@ def test_cli_sets_custom_source_root_if_outside_of_repository( mock_set_custom_source_root = mocker.patch( "zenml.utils.source_utils.set_custom_source_root" ) - runner = CliRunner() + runner = create_cli_runner() # Invoke a subcommand so the root CLI group gets called runner.invoke(cli, ["version"]) @@ -73,7 +75,7 @@ def test_cli_does_not_set_custom_source_root_if_inside_repository( mock_set_custom_source_root = mocker.patch( "zenml.utils.source_utils.set_custom_source_root" ) - runner = CliRunner() + runner = create_cli_runner() # Invoke a subcommand so the root CLI group gets called runner.invoke(cli, ["version"]) diff --git a/tests/integration/functional/steps/test_step_context.py b/tests/integration/functional/steps/test_step_context.py index f34e53b8e4b..dc258255bbf 100644 --- a/tests/integration/functional/steps/test_step_context.py +++ b/tests/integration/functional/steps/test_step_context.py @@ -1,6 +1,6 @@ import json import os -from typing import Any, Dict, List, Optional, Tuple, Type +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type import pytest from pydantic import BaseModel @@ -8,11 +8,13 @@ from zenml import get_step_context, log_metadata, pipeline, step from zenml.artifacts.artifact_config import ArtifactConfig -from zenml.client import Client from zenml.enums import ArtifactType from zenml.io import fileio from zenml.materializers.base_materializer import BaseMaterializer +if TYPE_CHECKING: + from zenml.client import Client + class ComplexObject(BaseModel): """This is a custom object to be materialized with ComplexObjectMaterializer.""" @@ -44,11 +46,7 @@ def save(self, data: ComplexObject) -> None: f.write(data.model_dump_json()) -@step( - output_materializers=[ - ComplexObjectMaterializer, - ] -) +@step def _output_complex_object_step(): """This step would call `save` of `ComplexObjectMaterializer`. `save` should not fail and have access to the step context""" @@ -68,19 +66,21 @@ def _access_step_context_step(): assert get_step_context().pipeline.name == "bar" -def test_materializer_can_access_step_context(): +def test_materializer_can_access_step_context(clean_client: "Client"): """Call steps using `ComplexObjectMaterializer` to validate that step context is available to Materializers""" @pipeline(name="bar") def _complex_object_materialization_pipeline(): - complex_object = _output_complex_object_step() + complex_object = _output_complex_object_step.with_options( + output_materializers=ComplexObjectMaterializer + )() _input_complex_object_step(complex_object) _complex_object_materialization_pipeline() -def test_step_can_access_step_context(): +def test_step_can_access_step_context(clean_client: "Client"): """Call step using step context directly, before Materializers""" @pipeline(name="bar") @@ -106,7 +106,7 @@ def step_context_metadata_reader_step(my_input: int) -> None: assert my_input_metadata["some_key"] == "some_value" -def test_input_artifacts_property(): +def test_input_artifacts_property(clean_client: "Client"): """Test the `StepContext.inputs` property.""" @pipeline diff --git a/tests/integration/functional/utils_functional/__init__.py b/tests/integration/functional/utils_functional/__init__.py new file mode 100644 index 00000000000..32ba63d3703 --- /dev/null +++ b/tests/integration/functional/utils_functional/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) ZenML GmbH 2025. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. diff --git a/tests/integration/functional/utils/test_metadata_utils.py b/tests/integration/functional/utils_functional/test_metadata_utils.py similarity index 100% rename from tests/integration/functional/utils/test_metadata_utils.py rename to tests/integration/functional/utils_functional/test_metadata_utils.py diff --git a/tests/integration/functional/utils/test_run_utils.py b/tests/integration/functional/utils_functional/test_run_utils.py similarity index 100% rename from tests/integration/functional/utils/test_run_utils.py rename to tests/integration/functional/utils_functional/test_run_utils.py diff --git a/tests/integration/functional/utils/test_tag_utils.py b/tests/integration/functional/utils_functional/test_tag_utils.py similarity index 100% rename from tests/integration/functional/utils/test_tag_utils.py rename to tests/integration/functional/utils_functional/test_tag_utils.py index fe684372b58..be5329a4e69 100644 --- a/tests/integration/functional/utils/test_tag_utils.py +++ b/tests/integration/functional/utils_functional/test_tag_utils.py @@ -17,8 +17,8 @@ from typing import Annotated, Tuple import pytest -from tests.integration.functional.zen_stores.utils import PipelineRunContext +from tests.integration.functional.zen_stores.utils import PipelineRunContext from zenml import ArtifactConfig, Tag, add_tags, pipeline, remove_tags, step from zenml.client import Client from zenml.constants import ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING diff --git a/tests/integration/functional/zen_stores/test_secrets_store.py b/tests/integration/functional/zen_stores/test_secrets_store.py index 877622c8dd9..3f8c0970c82 100644 --- a/tests/integration/functional/zen_stores/test_secrets_store.py +++ b/tests/integration/functional/zen_stores/test_secrets_store.py @@ -557,14 +557,18 @@ def test_reusing_private_secret_name_succeeds(): assert len(private_secrets) == 1 assert other_secret.id == private_secrets[0].id - other_private_secrets = store.list_secrets( - SecretFilter( - name=secret.name, - private=True, - ), - ).items - assert len(other_private_secrets) == 1 - assert secret.id == other_private_secrets[0].id + # After switching back to the original user context, instantiate a fresh client/store + original_client = Client() + original_store = original_client.zen_store + + other_private_secrets = original_store.list_secrets( + SecretFilter( + name=secret.name, + private=True, + ), + ).items + assert len(other_private_secrets) == 1 + assert secret.id == other_private_secrets[0].id def test_update_scope_succeeds():