From 1d40139a3d9cdc7d6d893f49022c7c5063654052 Mon Sep 17 00:00:00 2001
From: Felipe <fealho@gmail.com>
Date: Mon, 13 Nov 2023 08:29:15 -0800
Subject: [PATCH] Add tests

---
 sdv/metadata/visualization.py                 |  2 +-
 .../metadata/test_visualization.py            | 31 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 tests/integration/metadata/test_visualization.py

diff --git a/sdv/metadata/visualization.py b/sdv/metadata/visualization.py
index 88b4ae082..05498637e 100644
--- a/sdv/metadata/visualization.py
+++ b/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.
diff --git a/tests/integration/metadata/test_visualization.py b/tests/integration/metadata/test_visualization.py
new file mode 100644
index 000000000..5a8049bda
--- /dev/null
+++ b/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()