Add unit tests for write

Author: pvk-developer

sdv/io/local/local.py

@@ -69,6 +69,11 @@ def read(self, folder_name, file_names=None, read_csv_parameters=None):
             file_names (list of str, optional):
                 The names of CSV files to read. If None, all files ending with '.csv'
                 in the folder are read.
+            read_csv_parameters (dict):
+                A dictionary with additional parameters to use when reading the CSVs.
+                The keys are any of the parameter names of the pandas.read_csv function
+                and the values are your inputs. Defaults to
+                `{'parse_dates': False, 'low_memory': False, 'on_bad_lines': 'warn'}`
 
         Returns:
             dict:
@@ -77,6 +82,10 @@ def read(self, folder_name, file_names=None, read_csv_parameters=None):
         Raises:
             FileNotFoundError:
                 If the specified files do not exist in the folder.
+
+            ValueError:
+                If a provided parameter in `read_csv_parameters` is not supported by the
+                `CSVHandler`.
         """
         data = {}
         folder_path = Path(folder_name)
@@ -134,6 +143,11 @@ def write(
                 'x': Write to new files, raising errors if existing files exist with the same name.
                 'w': Write to new files, clearing any existing files that exist.
                 'a': Append the new CSV rows to any existing files.
+
+            to_csv_parameters (dict):
+                A dictionary with additional parameters to use when writing the CSVs.
+                The keys are any of the parameter names of the pandas.to_csv function and
+                the values are your input. Defaults to `{ 'index': False }`.
         """
         folder_path = Path(folder_name)
         to_csv_parameters = to_csv_parameters or {}

tests/integration/io/local/test_local.py

@@ -33,6 +33,32 @@ def test_integration_write_and_read(self, tmpdir):
         pd.testing.assert_frame_equal(data['table1'], synthetic_data['table1'])
         pd.testing.assert_frame_equal(data['table2'], synthetic_data['table2'])
 
+    def test_integration_write_and_read_with_custom_parameters(self, tmpdir):
+        """Test end to end the write and read methods of ``CSVHandler``."""
+        # Prepare synthetic data
+        synthetic_data = {
+            'table1': pd.DataFrame({'col1': [1, 2, 3], 'col2': ['a', 'b', 'c']}),
+            'table2': pd.DataFrame({'col3': [4, 5, 6], 'col4': ['d', 'e', 'f']}),
+        }
+
+        # Write synthetic data to CSV files
+        handler = CSVHandler()
+        write_params = {'sep': ';', 'index': True}
+        handler.write(synthetic_data, tmpdir, to_csv_parameters=write_params)
+
+        # Read data from CSV files
+        read_params = {'nrows': 1, 'sep': ';', 'index_col': 'Unnamed: 0'}
+        data = handler.read(tmpdir, read_csv_parameters=read_params)
+
+        # Check if data was read correctly
+        assert len(data) == 2
+        assert 'table1' in data
+        assert 'table2' in data
+        assert len(data['table1']) == 1
+        assert len(data['table2']) == 1
+        pd.testing.assert_frame_equal(data['table1'], synthetic_data['table1'].head(1))
+        pd.testing.assert_frame_equal(data['table2'], synthetic_data['table2'].head(1))
+
 
 class TestExcelHandler:
     def test_integration_write_and_read(self, tmpdir):

tests/unit/io/local/test_local.py

@@ -151,10 +151,10 @@ def test_read_files_custom_parameters(self, tmpdir):
         """Test the read method of CSVHandler class with custom read parameters."""
         # Setup
         file_path = Path(tmpdir)
-        read_csv_parameters =  {
+        read_csv_parameters = {
             'encoding': 'latin-1',
             'nrows': 1,
-            'escapechar': "\\",
+            'escapechar': '\\',
             'quotechar': '"',
             'sep': ';',
         }
@@ -169,22 +169,18 @@ def test_read_files_custom_parameters(self, tmpdir):
 
         # Run
         data = handler.read(
-            tmpdir,
-            file_names=['parent.csv'],
-            read_csv_parameters=read_csv_parameters
+            tmpdir, file_names=['parent.csv'], read_csv_parameters=read_csv_parameters
         )
 
         # Assert
         assert 'parent' in data
-        pd.testing.assert_frame_equal(
-            data['parent'], pd.DataFrame({'col1': [1], 'col2': ['a']})
-        )
+        pd.testing.assert_frame_equal(data['parent'], pd.DataFrame({'col1': [1], 'col2': ['a']}))
 
     def test_read_files_bad_parameters(self, tmpdir):
         """Test the read method of CSVHandler class with custom read parameters."""
         # Setup
         file_path = Path(tmpdir)
-        read_csv_parameters =  {
+        read_csv_parameters = {
             'filepath_or_buffer': 'myfile',
             'nrows': 1,
             'sep': ';',
@@ -205,11 +201,7 @@ def test_read_files_bad_parameters(self, tmpdir):
             'parameters instead.'
         )
         with pytest.raises(ValueError, match=error_msg):
-            handler.read(
-                tmpdir,
-                file_names=['parent.csv'],
-                read_csv_parameters=read_csv_parameters
-            )
+            handler.read(tmpdir, file_names=['parent.csv'], read_csv_parameters=read_csv_parameters)
 
     def test_write(self, tmpdir):
         """Test the write functionality of a CSVHandler."""
@@ -288,6 +280,29 @@ def test_write_file_exists_mode_is_w(self, tmpdir):
         expected_dataframe = pd.DataFrame({'col1': [1, 2, 3], 'col2': ['a', 'b', 'c']})
         pd.testing.assert_frame_equal(dataframe, expected_dataframe)
 
+    def test_write_file_with_custom_params(self, tmpdir):
+        """Test the write functionality of a CSVHandler when the mode is ``w``."""
+        # Setup
+        table_one_mock = Mock()
+        table_two_mock = Mock()
+
+        synthetic_data = {'table1': table_one_mock, 'table2': table_two_mock}
+
+        os.makedirs(tmpdir / 'synthetic_data')
+        handler = CSVHandler()
+        write_parameters = {'index': True, 'sep': ';'}
+
+        # Run
+        handler.write(synthetic_data, tmpdir / 'synthetic_data', to_csv_parameters=write_parameters)
+
+        # Assert
+        table_one_mock.to_csv.assert_called_once_with(
+            tmpdir / 'synthetic_data' / 'table1.csv', index=True, sep=';', mode='x'
+        )
+        table_two_mock.to_csv.assert_called_once_with(
+            tmpdir / 'synthetic_data' / 'table2.csv', index=True, sep=';', mode='x'
+        )
+
 
 class TestExcelHandler:
     def test___init__(self):