SIENTIAPDE-1350: Configure cleanup test script and improve test isolation. This commit configures the cleanup test script to load environment variables from a .env file and use them for Temporal connection. It also adds a fixture to clean up temporary directories created by tests, ensuring better test isolation and preventing potential conflicts. Additionally, it adds the 'uri' property to the MongoDB config.
This commit is contained in:
@@ -15,7 +15,8 @@ import os
|
|||||||
import sys
|
import sys
|
||||||
from datetime import timedelta
|
from datetime import timedelta
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
from pathlib import Path
|
||||||
from temporalio.client import Client
|
from temporalio.client import Client
|
||||||
|
|
||||||
# Ensure project root is on PYTHONPATH when running directly
|
# 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:
|
if ROOT_DIR not in sys.path:
|
||||||
sys.path.insert(0, ROOT_DIR)
|
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:
|
async def main(argv: list[str]) -> None:
|
||||||
@@ -34,11 +42,11 @@ async def main(argv: list[str]) -> None:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
# Config from environment / defaults
|
# Config from environment / defaults
|
||||||
temporal_host = os.getenv('TEMPORAL_HOST', 'localhost:37463')
|
temporal_host = os.getenv('TEMPORAL_HOST')
|
||||||
temporal_namespace = os.getenv('TEMPORAL_NAMESPACE', 'model-manager')
|
temporal_namespace = os.getenv('TEMPORAL_NAMESPACE')
|
||||||
task_queue = 'cleanup-queue'
|
task_queue = os.getenv('CLEANUP_TASK_QUEUE')
|
||||||
|
default_bucket = os.getenv('DEFAULT_CLEANUP_BUCKET')
|
||||||
default_bucket = os.getenv('DEFAULT_CLEANUP_BUCKET', 'model-training')
|
use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true'
|
||||||
|
|
||||||
# Optional CLI: bucket name override
|
# Optional CLI: bucket name override
|
||||||
bucket_name = default_bucket
|
bucket_name = default_bucket
|
||||||
@@ -46,7 +54,11 @@ async def main(argv: list[str]) -> None:
|
|||||||
bucket_name = argv[0]
|
bucket_name = argv[0]
|
||||||
|
|
||||||
print(f"Connecting to Temporal at {temporal_host} (namespace={temporal_namespace})...")
|
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] = {
|
input_data: dict[str, Any] = {
|
||||||
'bucket_name': bucket_name,
|
'bucket_name': bucket_name,
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
"""Unit tests for ModelRepository with 100% coverage."""
|
"""Unit tests for ModelRepository with 100% coverage."""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
|
import shutil
|
||||||
from unittest.mock import MagicMock, patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -8,6 +9,29 @@ import pandas as pd
|
|||||||
import pytest
|
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
|
@pytest.fixture
|
||||||
def mock_logger():
|
def mock_logger():
|
||||||
"""Create a 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
|
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._get_next_run_name = MagicMock(return_value='test_experiment-1')
|
||||||
repo._generate_artifacts = MagicMock(return_value=mock_train_result)
|
repo._generate_artifacts = MagicMock(return_value=mock_train_result)
|
||||||
repo._save_run = MagicMock()
|
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
|
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
|
mock_exists.return_value = True
|
||||||
|
|
||||||
repo.cleanup_run_directory('/tmp/test_run') # noqa: S108
|
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
|
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
|
mock_exists.return_value = False
|
||||||
|
|
||||||
repo.cleanup_run_directory('/tmp/test_run') # noqa: S108
|
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
|
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('')
|
repo.cleanup_run_directory('')
|
||||||
|
|
||||||
mock_logger.info.assert_called_once_with('No run directory specified, skipping cleanup')
|
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_train_result.run_name = None
|
||||||
|
mock_exists.return_value = True # Mock reports directory exists
|
||||||
|
|
||||||
with pytest.raises(ValueError, match='run_name must be set'):
|
# Mock _create_run_directory to avoid creating real directories
|
||||||
repo._generate_artifacts(mock_train_result)
|
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')
|
@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
|
mock_exists.side_effect = exists_side_effect
|
||||||
|
|
||||||
with pytest.raises(FileNotFoundError, match='Header file does not exist'):
|
# Mock _create_run_directory to avoid creating real directories
|
||||||
repo._generate_artifacts(mock_train_result)
|
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')
|
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||||
|
|||||||
@@ -96,6 +96,7 @@ def test_build_mongo_db_config_with_env_vars():
|
|||||||
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
|
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
|
||||||
'database_name': 'test_db',
|
'database_name': 'test_db',
|
||||||
'ttl_index_seconds': 3600,
|
'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',
|
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
|
||||||
'database_name': 'sientia',
|
'database_name': 'sientia',
|
||||||
'ttl_index_seconds': 3600,
|
'ttl_index_seconds': 3600,
|
||||||
|
'uri': 'localhost:27018',
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -209,6 +209,7 @@ async def test_main_successful_startup(
|
|||||||
mock_build_mongodb.return_value = {
|
mock_build_mongodb.return_value = {
|
||||||
'connection_string': 'mongodb://test',
|
'connection_string': 'mongodb://test',
|
||||||
'database_name': 'test_db',
|
'database_name': 'test_db',
|
||||||
|
'uri': 'localhost:27018',
|
||||||
}
|
}
|
||||||
mock_build_postgres.return_value = {}
|
mock_build_postgres.return_value = {}
|
||||||
mock_build_mlflow.return_value = {}
|
mock_build_mlflow.return_value = {}
|
||||||
@@ -290,6 +291,7 @@ async def test_main_handles_exception(
|
|||||||
mock_build_mongodb.return_value = {
|
mock_build_mongodb.return_value = {
|
||||||
'connection_string': 'mongodb://test',
|
'connection_string': 'mongodb://test',
|
||||||
'database_name': 'test_db',
|
'database_name': 'test_db',
|
||||||
|
'uri': 'localhost:27018',
|
||||||
}
|
}
|
||||||
mock_build_postgres.return_value = {}
|
mock_build_postgres.return_value = {}
|
||||||
mock_build_mlflow.return_value = {}
|
mock_build_mlflow.return_value = {}
|
||||||
@@ -372,6 +374,7 @@ async def test_main_temporal_client_configuration(
|
|||||||
mock_build_mongodb.return_value = {
|
mock_build_mongodb.return_value = {
|
||||||
'connection_string': 'mongodb://test',
|
'connection_string': 'mongodb://test',
|
||||||
'database_name': 'test_db',
|
'database_name': 'test_db',
|
||||||
|
'uri': 'localhost:27018',
|
||||||
}
|
}
|
||||||
mock_build_postgres.return_value = {}
|
mock_build_postgres.return_value = {}
|
||||||
mock_build_mlflow.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
|
# Verify Temporal client was configured with correct parameters
|
||||||
mock_client_class.connect.assert_called_once_with(
|
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 = {
|
mock_build_mongodb.return_value = {
|
||||||
'connection_string': 'mongodb://test',
|
'connection_string': 'mongodb://test',
|
||||||
'database_name': 'test_db',
|
'database_name': 'test_db',
|
||||||
|
'uri': 'localhost:27018',
|
||||||
}
|
}
|
||||||
mock_build_postgres.return_value = {}
|
mock_build_postgres.return_value = {}
|
||||||
mock_build_mlflow.return_value = {}
|
mock_build_mlflow.return_value = {}
|
||||||
|
|||||||
Reference in New Issue
Block a user