From 40d005c5aa64cad1565a332639e92c2f14153ab9 Mon Sep 17 00:00:00 2001
From: R-Palazzo <romainpalazzo@gmail.com>
Date: Thu, 2 Nov 2023 14:46:01 -0600
Subject: [PATCH] visualization to work with multi-foreign keys

---
 .../_properties/relationship_validity.py       |  5 ++++-
 .../_properties/test_relationship_validity.py  | 18 +++++++++++++-----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/sdmetrics/reports/multi_table/_properties/relationship_validity.py b/sdmetrics/reports/multi_table/_properties/relationship_validity.py
index 95145271..80b9ae3a 100644
--- a/sdmetrics/reports/multi_table/_properties/relationship_validity.py
+++ b/sdmetrics/reports/multi_table/_properties/relationship_validity.py
@@ -87,7 +87,10 @@ def _get_table_relationships_plot(self, table_name):
         """
         plot_data = self.get_details(table_name).copy()
         column_name = 'Child → Parent Relationship'
-        plot_data[column_name] = plot_data['Child Table'] + ' → ' + plot_data['Parent Table']
+        plot_data[column_name] = (
+            plot_data['Child Table'] + ' (' + plot_data['Foreign Key'] + ') → ' +
+            plot_data['Parent Table']
+        )
         plot_data = plot_data.drop(['Child Table', 'Parent Table'], axis=1)
 
         average_score = round(plot_data['Score'].mean(), 2)
diff --git a/tests/unit/reports/multi_table/_properties/test_relationship_validity.py b/tests/unit/reports/multi_table/_properties/test_relationship_validity.py
index b984c574..2c8b1b37 100644
--- a/tests/unit/reports/multi_table/_properties/test_relationship_validity.py
+++ b/tests/unit/reports/multi_table/_properties/test_relationship_validity.py
@@ -231,7 +231,9 @@ def test_get_details_with_table_name(self):
         relationship_validity.details = pd.DataFrame({
             'Child Table': ['users_child', 'sessions_child'],
             'Parent Table': ['users_parent', 'sessions_parent'],
-            'Metric': ['RelationshipValidityShapeSimilarity', 'SomeOtherMetric'],
+            'Primary Key': ['user_id', 'user_id'],
+            'Foreign Key': ['user_id', 'user_id'],
+            'Metric': ['ReferentialIntegrity', 'CardinalityBoundaryAdherence'],
             'Score': [1.0, 0.5],
             'Error': [None, 'Some error']
         })
@@ -244,7 +246,9 @@ def test_get_details_with_table_name(self):
         assert details_users_child.equals(pd.DataFrame({
             'Child Table': ['users_child'],
             'Parent Table': ['users_parent'],
-            'Metric': ['RelationshipValidityShapeSimilarity'],
+            'Primary Key': ['user_id'],
+            'Foreign Key': ['user_id'],
+            'Metric': ['ReferentialIntegrity'],
             'Score': [1.0],
             'Error': [None]
         }, index=[0]))
@@ -253,7 +257,9 @@ def test_get_details_with_table_name(self):
         assert details_sessions_parent.equals(pd.DataFrame({
             'Child Table': ['sessions_child'],
             'Parent Table': ['sessions_parent'],
-            'Metric': ['SomeOtherMetric'],
+            'Primary Key': ['user_id'],
+            'Foreign Key': ['user_id'],
+            'Metric': ['CardinalityBoundaryAdherence'],
             'Score': [0.5],
             'Error': ['Some error']
         }, index=[1]))
@@ -268,7 +274,9 @@ def test_get_table_relationships_plot(self):
         instance.details = pd.DataFrame({
             'Child Table': ['users_child', 'sessions_child'],
             'Parent Table': ['users_parent', 'sessions_parent'],
-            'Metric': ['RelationshipValidityShapeSimilarity', 'SomeOtherMetric'],
+            'Primary Key': ['user_id', 'user_id'],
+            'Foreign Key': ['user_id', 'user_id'],
+            'Metric': ['ReferentialIntegrity', 'CardinalityBoundaryAdherence'],
             'Score': [1.0, 0.5],
             'Error': [None, 'Some error']
         })
@@ -279,7 +287,7 @@ def test_get_table_relationships_plot(self):
         # Assert
         assert isinstance(fig, Figure)
 
-        expected_x = ['users_child → users_parent']
+        expected_x = ['users_child (user_id) → users_parent']
         expected_y = [1.0]
         expected_title = 'Data Diagnostic: Relationship Validity (Average Score=1.0)'