diff --git a/brevia/commands.py b/brevia/commands.py index 13d7cf4..a40e7c2 100644 --- a/brevia/commands.py +++ b/brevia/commands.py @@ -197,7 +197,8 @@ def cleanup_jobs(before_date: datetime, dry_run: bool): num = cleanup_async_jobs(before_date=before_date, dry_run=dry_run) if dry_run: - click.echo(f"Dry run completed. {num} {'job' if num == 1 else 'jobs'} would be deleted.") + njob = f"{num} {'job' if num == 1 else 'jobs'}" + click.echo(f"Dry run completed. {njob} would be deleted.") elif num == 0: click.echo("No async jobs to delete.") else: diff --git a/brevia/utilities/output.py b/brevia/utilities/output.py index 626c685..c379d8f 100644 --- a/brevia/utilities/output.py +++ b/brevia/utilities/output.py @@ -45,6 +45,9 @@ def file_url(self, filename: str): """ # Generate the output URL base_url = get_settings().file_output_base_url + if self.job_id: + base_url = f"{base_url}/{self.job_id}" + return f'{base_url}/{filename}' def _s3_upload(self, file_path: str, bucket_name: str, object_name: str): @@ -74,13 +77,13 @@ def write(self, content: str, filename: str): with open(output_path, 'w', encoding='utf-8') as file: file.write(content) - if self.job_id: - filename = f"{self.job_id}/{filename}" base_path = get_settings().file_output_base_path if base_path.startswith('s3://'): # Extract bucket name and object name from S3 path bucket_name = base_path.split('/')[2] object_name = '/'.join(base_path.split('/')[3:]).lstrip('/') + if self.job_id: + object_name += f"/{self.job_id}" object_name += f"/{filename}" self._s3_upload(output_path, bucket_name, object_name.lstrip('/')) # Remove the local temp file diff --git a/tests/utilities/test_output.py b/tests/utilities/test_output.py index 58368ff..f40593c 100644 --- a/tests/utilities/test_output.py +++ b/tests/utilities/test_output.py @@ -42,6 +42,11 @@ def test_file_url(): assert file_url == '/download/test.txt' + output = LinkedFileOutput(job_id='123456789') + file_url = output.file_url('test.txt') + + assert file_url == '/download/123456789/test.txt' + def test_write_local_file(): """Test write method for local file writing."""