From e41f6759c5b4b89116134455558ae4a618f97e0d Mon Sep 17 00:00:00 2001
From: Kyle Dolezal <kyle.d.dolezal@gmail.com>
Date: Mon, 24 Jun 2024 18:42:19 -0500
Subject: [PATCH] Add test for n_worst_best_routes

---
 data_analysis/tests/test_plots.py | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/data_analysis/tests/test_plots.py b/data_analysis/tests/test_plots.py
index 3e8c4a4..aafd3ac 100644
--- a/data_analysis/tests/test_plots.py
+++ b/data_analysis/tests/test_plots.py
@@ -1,7 +1,8 @@
+from typing import Union
 import plotly.graph_objects
-from data_analysis.plots import create_save_path
-from data_analysis.plots import lineplot
+from data_analysis.plots import create_save_path, lineplot, n_worst_best_routes
 import pandas as pd
+import geopandas 
 
 class MockFig:
     def add_trace(_, __):
@@ -28,3 +29,26 @@ def test_lineplot(mocker):
 
     assert mocked_plot.assert_called_once
     assert mocked_plot.call_args[1]["filename"] == "lineplot_of_y_over_x"
+
+def test_n_worst_best_routes(mocker):
+    df = pd.DataFrame(
+        {
+            "Latitude": [-34.58, -15.78, -33.45, 4.60, 10.48],
+            "Longitude": [-58.66, -47.91, -70.66, -74.08, -66.86],
+            "p_percentiles": [1, 2000, 3000, 4, 5],
+            "p": [1, 2000, 3000, 4, 5],
+            "route_id": [1, 2, 300, 4, 5]
+        }
+    )
+    gdf = geopandas.GeoDataFrame(
+        df, geometry=geopandas.points_from_xy(df.Longitude, df.Latitude), crs="EPSG:4326"
+    )
+    best = (n_worst_best_routes(
+        pd.concat([df, gdf]),
+        'p',
+        n=2,
+        percentile=True,
+        worst=False
+    ))
+    
+    assert best.iloc[0]["route_id"] == 1
\ No newline at end of file