From f4049bde6227b7529590be1ef26cb661936867f1 Mon Sep 17 00:00:00 2001 From: agaperste Date: Sat, 30 Sep 2023 12:05:08 -0400 Subject: [PATCH] addressing PR comment --- dune_client/viz/graphs.py | 4 +++- tests/unit/test_viz_sankey.py | 6 +++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/dune_client/viz/graphs.py b/dune_client/viz/graphs.py index 837bb2e..d7a8eb6 100644 --- 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 index aeb7fe5..098f80f 100644 --- 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,