From 11e8bf60f0757f0ff496e8ebc1bb13e54ea88105 Mon Sep 17 00:00:00 2001
From: Alex Morling <atmorling@gmail.com>
Date: Wed, 15 May 2024 13:59:37 +0200
Subject: [PATCH] Adding tests for EcoPlot

---
 ecoscope/plotting/plot.py |  2 +-
 tests/test_ecoplot.py     | 47 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_ecoplot.py

diff --git a/ecoscope/plotting/plot.py b/ecoscope/plotting/plot.py
index 60a58418..03ad4388 100644
--- a/ecoscope/plotting/plot.py
+++ b/ecoscope/plotting/plot.py
@@ -188,7 +188,7 @@ def nsd(relocations):
     relocations = relocations.to_crs(relocations.estimate_utm_crs())
 
     times = relocations["fixtime"]
-    distances = relocations.distance(relocations.geometry[0]) ** 2
+    distances = relocations.distance(relocations.geometry.iat[0]) ** 2
 
     fig = go.FigureWidget()
 
diff --git a/tests/test_ecoplot.py b/tests/test_ecoplot.py
new file mode 100644
index 00000000..16f3e42e
--- /dev/null
+++ b/tests/test_ecoplot.py
@@ -0,0 +1,47 @@
+import numpy as np
+from ecoscope.plotting.plot import EcoPlotData, ecoplot, mcp, nsd, speed
+from ecoscope.base import Trajectory
+
+
+def test_ecoplot(movebank_relocations):
+    traj = Trajectory.from_relocations(movebank_relocations)
+    epd = EcoPlotData(traj.groupby("groupby_col"), "segment_start", "speed_kmhr", line=dict(color="blue"))
+    figure = ecoplot([epd], "EcoPlot")
+
+    habiba = traj.loc[traj["groupby_col"] == "Habiba"]
+    salif = traj.loc[traj["groupby_col"] == "Salif Keita"]
+
+    assert len(figure.data) == 2
+
+    assert figure.data[0].name == "Habiba"
+    assert np.equal(figure.data[0].x, habiba["segment_start"].array).all()
+    assert np.equal(figure.data[0].y, habiba["speed_kmhr"].array).all()
+
+    assert figure.data[1].name == "Salif Keita"
+    assert np.equal(figure.data[1].x, salif["segment_start"].array).all()
+    assert np.equal(figure.data[1].y, salif["speed_kmhr"].array).all()
+
+
+def test_mcp(movebank_relocations):
+    figure = mcp(movebank_relocations)
+
+    assert len(figure.data) == 1
+    assert movebank_relocations["fixtime"].iat[0] == figure.data[0].x[0]
+    assert movebank_relocations["fixtime"].iat[-1] == figure.data[0].x[-1]
+
+
+def test_nsd(movebank_relocations):
+    figure = nsd(movebank_relocations)
+
+    assert len(figure.data) == 1
+    assert len(figure.data[0].x) == len(movebank_relocations)
+    assert len(figure.data[0].y) == len(movebank_relocations)
+
+
+def test_speed(movebank_relocations):
+    traj = Trajectory.from_relocations(movebank_relocations)
+    figure = speed(traj)
+
+    assert len(figure.data) == 1
+    len(figure.data[0].x) == len(traj) * 4
+    len(figure.data[0].y) == len(traj) * 4