addressing PR comments

duneanalytics · Sep 29, 2023 · 6d2ca2b · 6d2ca2b
1 parent 6fd2a22
commit 6d2ca2b
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 78 deletions.
diff --git a/dune_client/viz/graphs.py b/dune_client/viz/graphs.py
@@ -2,16 +2,17 @@
 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
 
 # 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
@@ -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
@@ -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"}