From 50b16e0a97c60f941e76193d78213d1a6ee870fa Mon Sep 17 00:00:00 2001
From: R-Palazzo <romainpalazzo@gmail.com>
Date: Tue, 21 Nov 2023 16:57:24 -0600
Subject: [PATCH] address comments

---
 sdv/datasets/demo.py             | 3 ++-
 tests/unit/datasets/test_demo.py | 6 ------
 2 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/sdv/datasets/demo.py b/sdv/datasets/demo.py
index 24b53f0df..c202d91c9 100644
--- a/sdv/datasets/demo.py
+++ b/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 '
diff --git a/tests/unit/datasets/test_demo.py b/tests/unit/datasets/test_demo.py
index dc1288f3d..fb0a5446f 100644
--- a/tests/unit/datasets/test_demo.py
+++ b/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."""