diff --git a/scripts/run_cleanup_test.py b/scripts/run_cleanup_test.py index 6084daa..16e8a87 100644 --- a/scripts/run_cleanup_test.py +++ b/scripts/run_cleanup_test.py @@ -15,7 +15,8 @@ import os import sys from datetime import timedelta from typing import Any - +from dotenv import load_dotenv +from pathlib import Path from temporalio.client import Client # Ensure project root is on PYTHONPATH when running directly @@ -23,7 +24,14 @@ ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if ROOT_DIR not in sys.path: sys.path.insert(0, ROOT_DIR) -from model_manager.workflows.cleanup_files import CleanupFiles # noqa: E402 +from model_manager.workflows.cleanup_files import CleanupFiles + + +# Carrega variáveis de ambiente do arquivo .env na raiz do projeto +PROJECT_ROOT = Path(__file__).resolve().parent.parent +ENV_PATH = PROJECT_ROOT / '.env' +if ENV_PATH.exists(): + load_dotenv(dotenv_path=ENV_PATH) async def main(argv: list[str]) -> None: @@ -34,11 +42,11 @@ async def main(argv: list[str]) -> None: """ # Config from environment / defaults - temporal_host = os.getenv('TEMPORAL_HOST', 'localhost:37463') - temporal_namespace = os.getenv('TEMPORAL_NAMESPACE', 'model-manager') - task_queue = 'cleanup-queue' - - default_bucket = os.getenv('DEFAULT_CLEANUP_BUCKET', 'model-training') + temporal_host = os.getenv('TEMPORAL_HOST') + temporal_namespace = os.getenv('TEMPORAL_NAMESPACE') + task_queue = os.getenv('CLEANUP_TASK_QUEUE') + default_bucket = os.getenv('DEFAULT_CLEANUP_BUCKET') + use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true' # Optional CLI: bucket name override bucket_name = default_bucket @@ -46,7 +54,11 @@ async def main(argv: list[str]) -> None: bucket_name = argv[0] print(f"Connecting to Temporal at {temporal_host} (namespace={temporal_namespace})...") - client = await Client.connect(temporal_host, namespace=temporal_namespace) + client = await Client.connect( + target_host=temporal_host, + namespace=temporal_namespace, + tls=use_tls, + ) input_data: dict[str, Any] = { 'bucket_name': bucket_name, diff --git a/tests/utils/repository/test_model_repository.py b/tests/utils/repository/test_model_repository.py index a31cb42..61d1cea 100644 --- a/tests/utils/repository/test_model_repository.py +++ b/tests/utils/repository/test_model_repository.py @@ -1,6 +1,7 @@ """Unit tests for ModelRepository with 100% coverage.""" import os +import shutil from unittest.mock import MagicMock, patch import numpy as np @@ -8,6 +9,29 @@ import pandas as pd import pytest +@pytest.fixture(autouse=True) +def cleanup_temp_directories(): + """Clean up temporary directories after each test.""" + # Get the temp directory path + current_file_dir = os.path.dirname(os.path.abspath(__file__)) + model_manager_dir = os.path.dirname(os.path.dirname(os.path.dirname(current_file_dir))) + temp_dir = os.path.join(model_manager_dir, 'reports', 'temp') + + # Run the test + yield + + # Clean up after test + if os.path.exists(temp_dir): + for item in os.listdir(temp_dir): + item_path = os.path.join(temp_dir, item) + if os.path.isdir(item_path) and item.startswith('test_run_'): + try: + shutil.rmtree(item_path) + except (OSError, PermissionError): + # Ignore cleanup errors + pass + + @pytest.fixture def mock_logger(): """Create a mock logger.""" @@ -79,6 +103,9 @@ def test_save_model_success(mock_model_serving_class, mock_logger, mock_train_re url='http://mlflow.test', username='user', password='pass', logger=mock_logger ) + # Reset mock after initialization to focus on method-specific calls + mock_logger.reset_mock() + repo._get_next_run_name = MagicMock(return_value='test_experiment-1') repo._generate_artifacts = MagicMock(return_value=mock_train_result) repo._save_run = MagicMock() @@ -105,6 +132,9 @@ def test_cleanup_run_directory_exists( url='http://mlflow.test', username='user', password='pass', logger=mock_logger ) + # Reset mock after initialization to focus on method-specific calls + mock_logger.reset_mock() + mock_exists.return_value = True repo.cleanup_run_directory('/tmp/test_run') # noqa: S108 @@ -124,6 +154,9 @@ def test_cleanup_run_directory_not_exists(mock_exists, mock_model_serving_class, url='http://mlflow.test', username='user', password='pass', logger=mock_logger ) + # Reset mock after initialization to focus on method-specific calls + mock_logger.reset_mock() + mock_exists.return_value = False repo.cleanup_run_directory('/tmp/test_run') # noqa: S108 @@ -141,6 +174,9 @@ def test_cleanup_run_directory_empty_path(mock_model_serving_class, mock_logger) url='http://mlflow.test', username='user', password='pass', logger=mock_logger ) + # Reset mock after initialization to focus on method-specific calls + mock_logger.reset_mock() + repo.cleanup_run_directory('') mock_logger.info.assert_called_once_with('No run directory specified, skipping cleanup') @@ -756,9 +792,14 @@ def test_generate_artifacts_no_run_name( ) mock_train_result.run_name = None + mock_exists.return_value = True # Mock reports directory exists - with pytest.raises(ValueError, match='run_name must be set'): - repo._generate_artifacts(mock_train_result) + # Mock _create_run_directory to avoid creating real directories + with patch.object(repo, '_create_run_directory') as mock_create_dir: + mock_create_dir.return_value = '/mock/run/dir' + + with pytest.raises(ValueError, match='run_name must be set'): + repo._generate_artifacts(mock_train_result) @patch('model_manager.utils.repository.model_repository.ModelServing') @@ -800,8 +841,12 @@ def test_generate_artifacts_header_not_found( mock_exists.side_effect = exists_side_effect - with pytest.raises(FileNotFoundError, match='Header file does not exist'): - repo._generate_artifacts(mock_train_result) + # Mock _create_run_directory to avoid creating real directories + with patch.object(repo, '_create_run_directory') as mock_create_dir: + mock_create_dir.return_value = '/mock/run/dir' + + with pytest.raises(FileNotFoundError, match='Header file does not exist'): + repo._generate_artifacts(mock_train_result) @patch('model_manager.utils.repository.model_repository.ModelServing') diff --git a/tests/utils/test_connectors_config.py b/tests/utils/test_connectors_config.py index 15adb7d..2679020 100644 --- a/tests/utils/test_connectors_config.py +++ b/tests/utils/test_connectors_config.py @@ -96,6 +96,7 @@ def test_build_mongo_db_config_with_env_vars(): 'connection_string': 'mongodb://sientia1:sientia1@localhost:27018', 'database_name': 'test_db', 'ttl_index_seconds': 3600, + 'uri': 'localhost:27018', } @@ -109,6 +110,7 @@ def test_build_mongo_db_config_with_defaults(): 'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018', 'database_name': 'sientia', 'ttl_index_seconds': 3600, + 'uri': 'localhost:27018', } diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py index 0e036c2..1a3d5d0 100644 --- a/tests/worker/test_worker.py +++ b/tests/worker/test_worker.py @@ -209,6 +209,7 @@ async def test_main_successful_startup( mock_build_mongodb.return_value = { 'connection_string': 'mongodb://test', 'database_name': 'test_db', + 'uri': 'localhost:27018', } mock_build_postgres.return_value = {} mock_build_mlflow.return_value = {} @@ -290,6 +291,7 @@ async def test_main_handles_exception( mock_build_mongodb.return_value = { 'connection_string': 'mongodb://test', 'database_name': 'test_db', + 'uri': 'localhost:27018', } mock_build_postgres.return_value = {} mock_build_mlflow.return_value = {} @@ -372,6 +374,7 @@ async def test_main_temporal_client_configuration( mock_build_mongodb.return_value = { 'connection_string': 'mongodb://test', 'database_name': 'test_db', + 'uri': 'localhost:27018', } mock_build_postgres.return_value = {} mock_build_mlflow.return_value = {} @@ -404,7 +407,10 @@ async def test_main_temporal_client_configuration( # Verify Temporal client was configured with correct parameters mock_client_class.connect.assert_called_once_with( - target_host='temporal.example.com:7233', namespace='production', runtime=mock_runtime + target_host='temporal.example.com:7233', + namespace='production', + runtime=mock_runtime, + tls=False, ) @@ -445,6 +451,7 @@ async def test_main_worker_configuration( mock_build_mongodb.return_value = { 'connection_string': 'mongodb://test', 'database_name': 'test_db', + 'uri': 'localhost:27018', } mock_build_postgres.return_value = {} mock_build_mlflow.return_value = {}