SIENTIAPDE-1231

Refactor tests and update model_repository.py for clarity and consistency

- Added a blank line in model_repository.py for improved readability.
- Adjusted formatting in test_mlflow.py to streamline assertions.
- Updated ensure_bucket_exists method tests in test_minio_repository.py to reflect the new return value of None instead of True.
This commit is contained in:
vitor-aignosi
2025-10-16 10:15:53 -03:00
parent d59833a5e5
commit de47820c4a
3 changed files with 5 additions and 6 deletions

View File

@@ -35,6 +35,7 @@ PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl'
INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
def force_memory_release(logger: Logger): def force_memory_release(logger: Logger):
gc.collect() gc.collect()
@@ -629,7 +630,7 @@ class MLFlowRepository:
# Aligns data with treated data indexes to get target variable # Aligns data with treated data indexes to get target variable
aligned_data = data.loc[treated_data.index] aligned_data = data.loc[treated_data.index]
aligned_series = aligned_data[target_name] aligned_series = aligned_data[target_name]
retrain_dataset = pd.merge( # NOSONAR retrain_dataset = pd.merge( # NOSONAR
treated_data, aligned_series, left_index=True, right_index=True treated_data, aligned_series, left_index=True, right_index=True
) )
else: else:

View File

@@ -200,9 +200,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
mock_to_datetime.assert_called_once_with( mock_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
) )
mock_to_datetime.return_value.dt.strftime.assert_called_once_with( mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
DATETIME_FORMAT
)
mock_to_datetime.assert_called_once_with( mock_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ

View File

@@ -67,7 +67,7 @@ def test_close(minio_repository):
def test_ensure_bucket_exists_bucket_exists(minio_repository): def test_ensure_bucket_exists_bucket_exists(minio_repository):
assert minio_repository.ensure_bucket_exists({}) is True assert minio_repository.ensure_bucket_exists({}) is None
minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
@@ -77,7 +77,7 @@ def test_ensure_bucket_exists_bucket_not_exists_create_success(minio_repository)
error_response={'Error': {'Code': '404'}}, operation_name='head_bucket' error_response={'Error': {'Code': '404'}}, operation_name='head_bucket'
) )
assert minio_repository.ensure_bucket_exists({}) is True assert minio_repository.ensure_bucket_exists({}) is None
minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test') minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')