From 302119c3f5196c05e6ae94672b03ededaf9b8382 Mon Sep 17 00:00:00 2001
From: Taha Yassine Kraiem <tahayk2@gmail.com>
Date: Tue, 24 Jan 2023 16:35:56 +0100
Subject: [PATCH] feat(chalice): handle wrong startTimestamp endTimestamp
 values

---
 api/schemas.py                              | 18 +++++++++++++++---
 ee/api/chalicelib/core/sessions_insights.py |  6 ++++--
 ee/api/schemas_ee.py                        |  2 +-
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/api/schemas.py b/api/schemas.py
index 740ecb5ba..f1aef97ab 100644
--- a/api/schemas.py
+++ b/api/schemas.py
@@ -92,7 +92,19 @@ class CreateNotificationSchema(BaseModel):
     notifications: List = Field(...)
 
 
-class NotificationsViewSchema(BaseModel):
+class _TimedSchema(BaseModel):
+    startTimestamp: int = Field(default=None)
+    endTimestamp: int = Field(default=None)
+
+    @root_validator
+    def time_validator(cls, values):
+        if values.get("startTimestamp") is not None and values.get("endTimestamp") is not None:
+            assert values.get("startTimestamp") < values.get("endTimestamp"), \
+                "endTimestamp must be greater than startTimestamp"
+        return values
+
+
+class NotificationsViewSchema(_TimedSchema):
     ids: Optional[List] = Field(default=[])
     startTimestamp: Optional[int] = Field(default=None)
     endTimestamp: Optional[int] = Field(default=None)
@@ -816,7 +828,7 @@ class SearchErrorsSchema(FlatSessionsSearchPayloadSchema):
     query: Optional[str] = Field(default=None)
 
 
-class MetricPayloadSchema(BaseModel):
+class MetricPayloadSchema(_TimedSchema):
     startTimestamp: int = Field(TimeUTC.now(delta_days=-1))
     endTimestamp: int = Field(TimeUTC.now())
     density: int = Field(7)
@@ -967,7 +979,7 @@ class MetricOfClickMap(str, Enum):
     click_map_url = "clickMapUrl"
 
 
-class CardSessionsSchema(FlatSessionsSearch, _PaginatedSchema):
+class CardSessionsSchema(FlatSessionsSearch, _PaginatedSchema,_TimedSchema):
     startTimestamp: int = Field(TimeUTC.now(-7))
     endTimestamp: int = Field(TimeUTC.now())
     series: List[CardCreateSeriesSchema] = Field(default=[])
diff --git a/ee/api/chalicelib/core/sessions_insights.py b/ee/api/chalicelib/core/sessions_insights.py
index f2f85ed3a..5de97e3c2 100644
--- a/ee/api/chalicelib/core/sessions_insights.py
+++ b/ee/api/chalicelib/core/sessions_insights.py
@@ -1,5 +1,7 @@
-import schemas, schemas_ee
-from typing import List, Optional
+from typing import Optional
+
+import schemas
+import schemas_ee
 from chalicelib.core import metrics
 from chalicelib.utils import ch_client
 
diff --git a/ee/api/schemas_ee.py b/ee/api/schemas_ee.py
index 60d92817b..b46e0dce5 100644
--- a/ee/api/schemas_ee.py
+++ b/ee/api/schemas_ee.py
@@ -47,7 +47,7 @@ class InsightCategories(str, Enum):
     resources = "resources"
 
 
-class GetInsightsSchema(BaseModel):
+class GetInsightsSchema(schemas._TimedSchema):
     startTimestamp: int = Field(default=TimeUTC.now(-7))
     endTimestamp: int = Field(default=TimeUTC.now())
     metricValue: List[InsightCategories] = Field(default=[])