diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index f9e783b..72707c1 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -35,6 +35,7 @@ PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl' INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." + def force_memory_release(logger: Logger): gc.collect() @@ -629,7 +630,7 @@ class MLFlowRepository: # Aligns data with treated data indexes to get target variable aligned_data = data.loc[treated_data.index] 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 ) else: diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 3d3a550..662f0ba 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -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_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ ) - mock_to_datetime.return_value.dt.strftime.assert_called_once_with( - DATETIME_FORMAT - ) + mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) mock_to_datetime.assert_called_once_with( mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ diff --git a/tests/laborious/utils/repository/test_minio_repository.py b/tests/laborious/utils/repository/test_minio_repository.py index 7883c1c..bebf743 100644 --- a/tests/laborious/utils/repository/test_minio_repository.py +++ b/tests/laborious/utils/repository/test_minio_repository.py @@ -67,7 +67,7 @@ def test_close(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') @@ -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' ) - 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.create_bucket.assert_called_once_with(Bucket='test')