test data usage

geo-engine · Jan 9, 2025 · f4a04c5 · f4a04c5
1 parent 204bf03
commit f4a04c5
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 2 deletions.
diff --git a/geoengine/workflow.py b/geoengine/workflow.py
@@ -988,7 +988,7 @@ def data_usage(offset: int = 0, limit: int = 10) -> List[geoengine_openapi_clien
 
 
 def data_usage_summary(granularity: geoengine_openapi_client.UsageSummaryGranularity,
-                       data: Optional[str] = None,
+                       dataset: Optional[str] = None,
                        offset: int = 0, limit: int = 10) -> pd.DataFrame:
     '''
     Get data usage summary
@@ -998,7 +998,7 @@ def data_usage_summary(granularity: geoengine_openapi_client.UsageSummaryGranula
 
     with geoengine_openapi_client.ApiClient(session.configuration) as api_client:
         user_api = geoengine_openapi_client.UserApi(api_client)
-        response = user_api.data_usage_summary_handler(data=data, granularity=granularity,
+        response = user_api.data_usage_summary_handler(dataset=dataset, granularity=granularity,
                                                        offset=offset, limit=limit)
 
         # create dataframe from response

diff --git a/tests/test_data_usage.py b/tests/test_data_usage.py
@@ -0,0 +1,98 @@
+'''Tests for WMS calls'''
+
+import unittest
+import pandas as pd
+from geoengine.datasets import DatasetName, UploadId, StoredDataset
+import geoengine as ge
+from . import UrllibMocker
+
+
+class DataUsageTests(unittest.TestCase):
+    '''Test methods for data usage'''
+
+    def setUp(self) -> None:
+        ge.reset(False)
+
+    def test_data_usage(self):
+
+        with UrllibMocker() as m:
+            m.post('http://mock-instance/login',
+                   expected_request_body={"email": "admin@localhost", "password": "adminadmin"},
+                   json={
+                       "id": "e327d9c3-a4f3-4bd7-a5e1-30b26cae8064",
+                       "user": {
+                           "id": "328ca8d1-15d7-4f59-a989-5d5d72c98744",
+                       },
+                       "created": "2021-06-08T15:22:22.605891994Z",
+                       "validUntil": "2021-06-08T16:22:22.605892183Z",
+                       "project": None,
+                       "view": None
+                   })
+
+            m.get('http://mock-instance/quota/dataUsage?offset=0&limit=10', json=[
+                {
+                    "timestamp": "2025-01-09T16:40:22.933Z",
+                    "userId": "e440bffc-d899-4304-aace-b23fc56828b2",
+                    "computationId": "7b08af4a-8793-4299-83c1-39d0c20560f5",
+                    "data": "land_cover",
+                    "count": 4
+                },
+                {
+                    "timestamp": "2025-01-09T16:40:10.970Z",
+                    "userId": "e440bffc-d899-4304-aace-b23fc56828b2",
+                    "computationId": "57fba95f-d693-432b-a65b-58ac225a384a",
+                    "data": "land_cover",
+                    "count": 4
+                }
+            ])
+
+            ge.initialize("http://mock-instance", ("admin@localhost", "adminadmin"))
+
+            df = ge.data_usage(offset=0, limit=10)
+
+            expected = pd.DataFrame({
+                'computationId': [
+                    '7b08af4a-8793-4299-83c1-39d0c20560f5',
+                    '57fba95f-d693-432b-a65b-58ac225a384a'
+                ],
+                'count': [4, 4],
+                'data': ['land_cover', 'land_cover'],
+                'timestamp': [pd.Timestamp('2025-01-09 16:40:22.933000+0000', tz='UTC'), pd.Timestamp('2025-01-09 16:40:10.970000+0000', tz='UTC')],
+                'userId': ['e440bffc-d899-4304-aace-b23fc56828b2', 'e440bffc-d899-4304-aace-b23fc56828b2']})
+
+            pd.testing.assert_frame_equal(df, expected)
+
+    def test_data_usage_summary(self):
+
+        with UrllibMocker() as m:
+            m.post('http://mock-instance/login',
+                   expected_request_body={"email": "admin@localhost", "password": "adminadmin"},
+                   json={
+                       "id": "e327d9c3-a4f3-4bd7-a5e1-30b26cae8064",
+                       "user": {
+                           "id": "328ca8d1-15d7-4f59-a989-5d5d72c98744",
+                       },
+                       "created": "2021-06-08T15:22:22.605891994Z",
+                       "validUntil": "2021-06-08T16:22:22.605892183Z",
+                       "project": None,
+                       "view": None
+                   })
+
+            m.get('http://mock-instance/quota/dataUsage/summary?granularity=minutes&offset=0&limit=10', json=[
+                {
+                    "timestamp": "2025-01-09T16:40:00.000Z",
+                    "data": "land_cover",
+                    "count": 8
+                }
+            ])
+
+            ge.initialize("http://mock-instance", ("admin@localhost", "adminadmin"))
+
+            df = ge.data_usage_summary(granularity=ge.UsageSummaryGranularity.MINUTES, offset=0, limit=10)
+
+            expected = pd.DataFrame({
+                'count': {0: 8},
+                'data': {0: 'land_cover'},
+                'timestamp': {0: pd.Timestamp('2025-01-09 16:40:00+0000', tz='UTC')}})
+
+            pd.testing.assert_frame_equal(df, expected)