Update message for text files

Author: pvk-developer

sdv/datasets/demo.py

@@ -197,6 +197,22 @@ def is_direct_json_under_prefix(key):
     )
 
 
+def _download_text_file_error_message(
+    modality,
+    dataset_name,
+    output_filepath=None,
+    bucket=PUBLIC_BUCKET,
+    filename=None,
+    **kwargs,
+):
+    return (
+        f"Could not retrieve '{filename}' for dataset '{dataset_name}' "
+        f"from bucket '{bucket}'. "
+        'Make sure the bucket name is correct. If the bucket is private '
+        'make sure to provide your credentials.'
+    )
+
+
 def _download_error_message(
     modality,
     dataset_name,
@@ -640,6 +656,7 @@ def _save_document(text, output_filepath, filename, dataset_name):
         LOGGER.info(f'Error saving {filename} for dataset {dataset_name}.')
 
 
+@handle_aws_client_errors(_download_text_file_error_message)
 def _get_text_file_content(
     modality, dataset_name, filename, output_filepath=None, bucket=PUBLIC_BUCKET, credentials=None
 ):
@@ -690,7 +707,7 @@ def _get_text_file_content(
     return text
 
 
-@handle_aws_client_errors(_download_error_message)
+@handle_aws_client_errors(_download_text_file_error_message)
 def get_source(
     modality, dataset_name, output_filepath=None, s3_bucket_name=PUBLIC_BUCKET, credentials=None
 ):
@@ -727,7 +744,6 @@ def get_source(
     )
 
 
-@handle_aws_client_errors(_download_error_message)
 def get_readme(
     modality, dataset_name, output_filepath=None, s3_bucket_name=PUBLIC_BUCKET, credentials=None
 ):

tests/unit/datasets/test_demo.py

@@ -1465,7 +1465,8 @@ def test_get_source_with_client_error(mock__create_s3_client):
     mock__create_s3_client.return_value = client
 
     error_msg = (
-        "Could not download dataset 'fake_hotels' from bucket 'private_bucket'. "
+        "Could not retrieve 'SOURCE.txt' for dataset 'fake_hotels' "
+        "from bucket 'private_bucket'. "
         'Make sure the bucket name is correct. If the bucket is private '
         'make sure to provide your credentials.'
     )
@@ -1499,7 +1500,8 @@ def test_get_readme_with_client_error(mock__create_s3_client):
     mock__create_s3_client.return_value = client
 
     error_msg = (
-        "Could not download dataset 'fake_hotels' from bucket 'private_bucket'. "
+        "Could not retrieve 'README.txt' for dataset 'fake_hotels' "
+        "from bucket 'private_bucket'. "
         'Make sure the bucket name is correct. If the bucket is private '
         'make sure to provide your credentials.'
     )