address comments

sdv/datasets/demo.py

@@ -13,6 +13,7 @@
 import pandas as pd
 from botocore import UNSIGNED
 from botocore.client import Config
+from botocore.exceptions import ClientError
 
 from sdv.metadata.multi_table import MultiTableMetadata
 from sdv.metadata.single_table import SingleTableMetadata
@@ -51,7 +52,7 @@ def _download(modality, dataset_name):
     LOGGER.info(f'Downloading dataset {dataset_name} from {dataset_url}')
     try:
         file_content = _get_data_from_bucket(object_key)
-    except Exception:
+    except ClientError:
         raise ValueError(
             f"Invalid dataset name '{dataset_name}'. "
             'Make sure you have the correct modality for the dataset name or '

tests/unit/datasets/test_demo.py

@@ -1,4 +1,3 @@
-import os
 import re
 from unittest.mock import MagicMock, Mock, patch
 
@@ -66,11 +65,6 @@ def test_download_demo_single_table(tmpdir):
 
 
 @patch('boto3.Session')
-@patch.dict(os.environ, {
-    'AWS_ACCESS_KEY_ID': 'access_key',
-    'AWS_SECRET_ACCESS_KEY': 'secret_access_key',
-    'AWS_DEFAULT_REGION': 'region_name',
-})
 @patch('sdv.datasets.demo.BUCKET', 'bucket')
 def test__get_data_from_bucket(session_mock):
     """Test the ``_get_data_from_bucket`` method."""