Can't set address columns after fitting (#1662)

sdv/data_processing/data_processor.py

@@ -177,6 +177,9 @@ def set_address_transformer(self, column_names, anonymization_level):
         transformer = self._get_address_transformer(anonymization_level)
         transformer._validate_sdtypes(columns_to_sdtypes)
 
+        if self._prepared_for_fitting:
+            self.update_transformers({column_names: transformer})
+
         self.grouped_columns_to_transformers[column_names] = transformer
 
     def get_model_kwargs(self, model_name):

tests/unit/data_processing/test_data_processor.py

@@ -260,6 +260,35 @@ def test__set_address_transformer(self):
         transformer._validate_sdtypes.assert_called_once_with(columns_to_sdtypes)
         assert dp.grouped_columns_to_transformers == {columns: transformer}
 
+    def test__set_address_transformer_prepared_for_fitting_true(self):
+        """Test the ``set_address_transformer`` method when ``_prepared_for_fitting`` is True."""
+        # Setup
+        metadata = SingleTableMetadata().load_from_dict({
+            'columns': {
+                'country_column': {'sdtype': 'country_code'},
+                'city_column': {'sdtype': 'city'},
+            }
+        })
+        dp = DataProcessor(metadata)
+        transformer = Mock()
+        transformer._validate_sdtypes = Mock()
+        columns_to_sdtypes = {'country_column': 'country_code', 'city_column': 'city'}
+        dp._get_address_transformer = Mock(return_value=transformer)
+        columns = ('country_column', 'city_column')
+        dp.update_transformers = Mock()
+        dp._prepared_for_fitting = True
+
+        # Run
+        dp.set_address_transformer(columns, 'full')
+
+        # Assert
+        dp.update_transformers.assert_called_once_with(
+            {('country_column', 'city_column'): transformer}
+        )
+        dp._get_address_transformer.assert_called_once_with('full')
+        transformer._validate_sdtypes.assert_called_once_with(columns_to_sdtypes)
+        assert dp.grouped_columns_to_transformers == {columns: transformer}
+
     def test_filter_valid(self):
         """Test that we are calling the ``filter_valid`` of each constraint over the data."""
         # Setup