addressing PR comment

duneanalytics · Sep 30, 2023 · f4049bd · f4049bd
1 parent 509874e
commit f4049bd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/dune_client/viz/graphs.py b/dune_client/viz/graphs.py
@@ -21,7 +21,9 @@ def create_sankey(
 ) -> Figure:
     """
     Creates a Sankey diagram based on input query_result,
-    which must contain source, target, value columns
+    which must contain source, target, value columns.
+    Column names don't have to be exact same but there must be
+    these three columns conceptually and value column must be numeric.
     """
     # Check if the dataframe contains required columns
     required_columns = [columns["source"], columns["target"], columns["value"]]

diff --git a/tests/unit/test_viz_sankey.py b/tests/unit/test_viz_sankey.py
@@ -20,7 +20,11 @@ def setUp(self):
             "WBTC": "rgb(247, 150, 38)",
         }
 
-        self.columns = {"source": "source", "target": "target", "value": "value"}
+        self.columns = {
+            "source": "source_col",
+            "target": "target_col",
+            "value": "value_col",
+        }
         self.viz_config: dict = {
             "node_pad": 15,
             "node_thickness": 20,