From 6d2ca2b3a6e3a0b20f7ce6375e020689f78b4016 Mon Sep 17 00:00:00 2001
From: agaperste <agapehumilduro@gmail.com>
Date: Fri, 29 Sep 2023 19:10:29 -0400
Subject: [PATCH] addressing PR comments

---
 dune_client/viz/graphs.py     |  7 ++--
 requirements/dev.txt          |  2 +
 tests/unit/test_viz_sankey.py | 78 ++---------------------------------
 3 files changed, 9 insertions(+), 78 deletions(-)

diff --git a/dune_client/viz/graphs.py b/dune_client/viz/graphs.py
index 3cba2ba..1670866 100644
--- a/dune_client/viz/graphs.py
+++ b/dune_client/viz/graphs.py
@@ -2,6 +2,7 @@
 Functions you can call to make different graphs
 """
 
+from typing import Union
 import plotly.graph_objects as go
 import colorlover as cl
 import pandas as pd
@@ -9,9 +10,9 @@
 # function to create Sankey diagram
 def create_sankey(
     query_result: pd.DataFrame,
-    predefined_colors: dict,
-    columns: dict,
-    viz_config: dict,
+    predefined_colors: dict[str, str],
+    columns: dict[str, str],
+    viz_config: dict[str, Union[int, float]],
     title: str = "unnamed",
 ):
     """
diff --git a/requirements/dev.txt b/requirements/dev.txt
index a3583c5..9311cc5 100644
--- a/requirements/dev.txt
+++ b/requirements/dev.txt
@@ -6,3 +6,5 @@ pytest>=7.4.1
 python-dotenv>=1.0.0
 mypy>=1.5.1
 aiounittest>=1.4.2
+colorlover>=0.3.0
+plotly>=5.9.0
diff --git a/tests/unit/test_viz_sankey.py b/tests/unit/test_viz_sankey.py
index 9f62041..7d803c3 100644
--- a/tests/unit/test_viz_sankey.py
+++ b/tests/unit/test_viz_sankey.py
@@ -10,87 +10,15 @@ class TestCreateSankey(unittest.TestCase):
     def setUp(self):
         self.df = pd.DataFrame(
             {
-                "source": [
-                    "WBTC",
-                    "USDC",
-                    "USDC",
-                    "USDC",
-                    "USDC",
-                    "COMP",
-                    "DAI",
-                    "DAI",
-                    "USDT",
-                    "WBTC",
-                    "DAI",
-                    "DAI",
-                    "USDC",
-                    "MKR",
-                    "DAI",
-                    "USDT",
-                    "UNI",
-                    "USDT",
-                    "USDT",
-                    "WBTC",
-                    "USDC",
-                ],
-                "target": [
-                    "WETH",
-                    "WBTC",
-                    "COMP",
-                    "MKR",
-                    "DAI",
-                    "WETH",
-                    "COMP",
-                    "WETH",
-                    "MKR",
-                    "USDT",
-                    "MKR",
-                    "USDT",
-                    "USDT",
-                    "WETH",
-                    "WBTC",
-                    "DAI",
-                    "WETH",
-                    "WETH",
-                    "WBTC",
-                    "DAI",
-                    "UNI",
-                ],
-                "value": [
-                    2184,
-                    2076,
-                    447,
-                    158,
-                    4294,
-                    519,
-                    72,
-                    4070,
-                    123,
-                    99,
-                    85,
-                    188,
-                    4675,
-                    352,
-                    281,
-                    230,
-                    59,
-                    4482,
-                    103,
-                    171,
-                    54,
-                ],
+                "source": ["WBTC", "USDC"],
+                "target": ["USDC", "WBTC"],
+                "value": [2184, 2076],
             }
         )
 
         self.predefined_colors = {
             "USDC": "rgb(38, 112, 196)",
-            "USDT": "rgb(0, 143, 142)",
-            "WETH": "rgb(144, 144, 144)",
             "WBTC": "rgb(247, 150, 38)",
-            "COMP": "rgb(32, 217, 152)",
-            "DAI": "rgb(254, 175, 48)",
-            "MKR": "rgb(38, 173, 158)",
-            "UNI": "rgb(255, 21, 126)",
         }
 
         self.columns = {"source": "source", "target": "target", "value": "value"}