diff --git a/sdv/metadata/multi_table.py b/sdv/metadata/multi_table.py index c7838500c..bcfe41764 100644 --- a/sdv/metadata/multi_table.py +++ b/sdv/metadata/multi_table.py @@ -556,7 +556,7 @@ def visualize(self, show_table_details='full', show_relationship_labels=True, ) show_table_details = None - FutureWarning(future_warning_msg) + warnings.warn(future_warning_msg, FutureWarning) nodes = {} edges = [] diff --git a/tests/unit/metadata/test_multi_table.py b/tests/unit/metadata/test_multi_table.py index c0ece36cb..f3ee1b0a6 100644 --- a/tests/unit/metadata/test_multi_table.py +++ b/tests/unit/metadata/test_multi_table.py @@ -1298,10 +1298,10 @@ def test_visualize_show_relationship_and_details_summarized(self, visualize_grap ] visualize_graph_mock.assert_called_once_with(expected_nodes, expected_edges, None) - @patch('sdv.metadata.multi_table.FutureWarning') + @patch('sdv.metadata.multi_table.warnings') @patch('sdv.metadata.multi_table.visualize_graph') def test_visualize_show_relationship_and_details_warning(self, visualize_graph_mock, - future_warning_mock): + warnings_mock): """Test the ``visualize`` method. If both the ``show_relationship_labels`` and ``show_table_details`` parameters are @@ -1350,18 +1350,18 @@ def test_visualize_show_relationship_and_details_warning(self, visualize_graph_m ('users', 'payments', ' user_id → id') ] visualize_graph_mock.assert_called_once_with(expected_nodes, expected_edges, None) - future_warning_mock.assert_called_once_with( + warnings_mock.warn.assert_called_once_with( 'Using True or False for show_table_details is deprecated. Use ' - "show_table_details='full' to show all table details." + "show_table_details='full' to show all table details.", FutureWarning ) @patch('sdv.metadata.multi_table.visualize_graph') def test_visualize_show_relationship_show_table_details_none(self, visualize_graph_mock): """Test the ``visualize`` method. - If ``show_relationship_labels`` is True but ``show_table_details``is False, + If ``show_relationship_labels`` is True but ``show_table_details``is None, then the edges should have labels and the labels for the nodes should be just - the table name. Also a ``FutureWarning`` should be raised to use ``None`` instead. + the table name. Setup: - Mock the ``visualize_graph`` function. @@ -1395,10 +1395,10 @@ def test_visualize_show_relationship_show_table_details_none(self, visualize_gra ] visualize_graph_mock.assert_called_once_with(expected_nodes, expected_edges, 'output.jpg') - @patch('sdv.metadata.multi_table.FutureWarning') + @patch('sdv.metadata.multi_table.warnings') @patch('sdv.metadata.multi_table.visualize_graph') def test_visualize_show_relationship_only_warning(self, visualize_graph_mock, - future_warning_mock): + warnings_mock): """Test the ``visualize`` method. If ``show_relationship_labels`` is True but ``show_table_details``is False, @@ -1435,9 +1435,9 @@ def test_visualize_show_relationship_only_warning(self, visualize_graph_mock, ('users', 'payments', ' user_id → id') ] visualize_graph_mock.assert_called_once_with(expected_nodes, expected_edges, 'output.jpg') - future_warning_mock.assert_called_once_with( + warnings_mock.warn.assert_called_once_with( "Using True or False for 'show_table_details' is deprecated. " - 'Use show_table_details=None to hide table details.' + 'Use show_table_details=None to hide table details.', FutureWarning ) @patch('sdv.metadata.multi_table.visualize_graph')