From 64b4d35be77d10d6b0541f4eb233f393e35f23aa Mon Sep 17 00:00:00 2001
From: Sander Roosendaal <sanderroosendaal@proton.me>
Date: Sun, 21 Apr 2024 11:04:06 +0200
Subject: [PATCH] tasks

---
 rowers/tasks.py                  | 32 ++++++++++++++++----------------
 rowers/tests/mocks.py            |  4 ++--
 rowers/tests/test_async_tasks.py |  4 ++--
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/rowers/tasks.py b/rowers/tasks.py
index 104045e5..8c46ca9f 100644
--- a/rowers/tasks.py
+++ b/rowers/tasks.py
@@ -2553,9 +2553,10 @@ def handle_otwsetpower(self, f1, boattype, boatclass, coastalbrand, weightvalue,
         totaltime = totaltime + rowdata.df.loc[0, ' ElapsedTime (sec)']
     except KeyError:  # pragma: no cover
         pass
-    df = getsmallrowdata_db(
+    df = getsmallrowdata_pd(
         ['power', 'workoutid', 'time'], ids=[workoutid],
         debug=debug)
+
     thesecs = totaltime
     maxt = 1.05 * thesecs
     logarr = datautils.getlogarr(maxt)
@@ -3171,23 +3172,22 @@ graphql_url = "https://rp3rowing-app.com/graphql"
 
 @app.task
 def handle_update_wps(rid, types, ids, mode, debug=False, **kwargs):
-    df = getsmallrowdata_db(['time', 'driveenergy'], ids=ids)
+    df = read_data(['time', 'driveenergy'], ids=ids)
     try:
-        mask = df['driveenergy'] > 100
-    except (KeyError, TypeError):  # pragma: no cover
-        return 0
-    try:
-        wps_median = int(df.loc[mask, 'driveenergy'].median())
+        wps_median = int(df.filter(pl.col("driveenergy")>100)["driveenergy"].median())
+        rower = Rower.objects.get(id=rid)
+        if mode == 'water':
+            rower.median_wps = wps_median
+        else:  # pragma: no cover
+            rower.median_wps_erg = wps_median
+
+        rower.save()
     except ValueError:  # pragma: no cover
-        return 0
-
-    rower = Rower.objects.get(id=rid)
-    if mode == 'water':
-        rower.median_wps = wps_median
-    else:
-        rower.median_wps_erg = wps_median
-
-    rower.save()
+        wps_median = 0
+    except OverflowError:
+        wps_median = 0
+    except ColumnNotFoundError:
+        wps_median = 0
 
     return wps_median
 
diff --git a/rowers/tests/mocks.py b/rowers/tests/mocks.py
index 617c7c38..ed73b421 100644
--- a/rowers/tests/mocks.py
+++ b/rowers/tests/mocks.py
@@ -357,8 +357,8 @@ def mocked_getsmallrowdata_db_water(*args, **kwargs): # pragma: no cover
 
     return df
 
-def mocked_getsmallrowdata_db_wps(*args, **kwargs): # pragma: no cover
-    df = pd.read_csv('rowers/tests/testdata/driveenergies.csv')
+def mocked_read_data_wps(*args, **kwargs): # pragma: no cover
+    df = pl.read_csv('rowers/tests/testdata/driveenergies.csv')
 
     return df
 
diff --git a/rowers/tests/test_async_tasks.py b/rowers/tests/test_async_tasks.py
index 5aab5a37..f7af8923 100644
--- a/rowers/tests/test_async_tasks.py
+++ b/rowers/tests/test_async_tasks.py
@@ -486,8 +486,8 @@ class AsyncTaskTests(TestCase):
 
 
 
-    @patch('rowers.dataprep.getsmallrowdata_db')
-    def test_handle_update_wps(self,mocked_getsmallrowdata_db_wps):
+    @patch('rowers.dataprep.read_data')
+    def test_handle_update_wps(self,mocked_read_data_wps):
       ids = [1,2,3]
 
       result = tasks.handle_update_wps(self.r.id,['water'],ids,mode='water')