Add tests

sdv/metadata/visualization.py

@@ -75,7 +75,7 @@ def _replace_special_characters(string):
 
 
 def visualize_graph(nodes, edges, filepath=None):
-    """Plot metadata usign graphviz.
+    """Plot metadata using graphviz.
 
     Try to generate a plot using graphviz.
     If a ``filepath`` is provided save the output into a file.

tests/integration/metadata/test_visualization.py

@@ -0,0 +1,31 @@
+import pandas as pd
+
+from sdv.metadata import MultiTableMetadata, SingleTableMetadata
+
+
+def test_visualize_graph_for_single_table():
+    """Test it runs when a column name contains `>`."""
+    # Setup
+    data = pd.DataFrame({'>': ['a', 'b', 'c']})
+    metadata = SingleTableMetadata()
+    metadata.detect_from_dataframe(data)
+
+    # Run
+    metadata.visualize()
+
+
+def test_visualize_graph_for_multi_table():
+    """Test it runs when a column name contains `>`."""
+    # Setup
+    data1 = pd.DataFrame({'>': ['a', 'b', 'c']})
+    data2 = pd.DataFrame({'>': ['a', 'b', 'c']})
+    tables = {'1': data1, '2': data2}
+    metadata = MultiTableMetadata()
+    metadata.detect_from_dataframes(tables)
+    metadata.update_column('1', '>', sdtype='id')
+    metadata.update_column('2', '>', sdtype='id')
+    metadata.set_primary_key('1', '>')
+    metadata.add_relationship('1', '2', '>', '>')
+
+    # Run
+    metadata.visualize()