Raising warning for visualize

sdv-dev · Aug 14, 2023 · a6298a8 · a6298a8
1 parent 0433374
commit a6298a8
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git 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
@@ -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')