[PBE-1321] align query API

GetStream · Feb 21, 2024 · a50ca3d · a50ca3d
1 parent 1dae3c3
commit a50ca3d
Show file tree

Hide file tree

Showing 11 changed files with 220 additions and 24 deletions.
diff --git a/stream_chat/async_chat/client.py b/stream_chat/async_chat/client.py
@@ -19,6 +19,7 @@
 
 from stream_chat.async_chat.campaign import Campaign
 from stream_chat.async_chat.segment import Segment
+from stream_chat.types.base import SortParam
 from stream_chat.types.campaign import CampaignData, QueryCampaignsOptions
 from stream_chat.types.segment import (
     QuerySegmentsOptions,
@@ -575,9 +576,18 @@ async def get_segment(self, segment_id: str) -> StreamResponse:
         return await self.get(f"segments/{segment_id}")
 
     async def query_segments(
-        self, filter_conditions: Dict, options: QuerySegmentsOptions
+        self,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
+        options: Optional[QuerySegmentsOptions] = None
     ) -> StreamResponse:
-        payload = {"filter": filter_conditions, **options}
+        payload = {}
+        if filter_conditions is not None:
+            payload["filter"] = filter_conditions
+        if sort is not None:
+            payload["sort"] = sort
+        if options is not None:
+            payload.update(cast(dict, options))
         return await self.post("segments/query", data=payload)
 
     async def update_segment(
@@ -603,12 +613,15 @@ async def add_segment_targets(
     async def query_segment_targets(
         self,
         segment_id: str,
-        filter_conditions: Optional[Dict] = None,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
         options: Optional[QuerySegmentTargetsOptions] = None,
     ) -> StreamResponse:
         payload = {}
         if filter_conditions is not None:
             payload["filter"] = filter_conditions
+        if sort is not None:
+            payload["sort"] = sort
         if options is not None:
             payload.update(cast(dict, options))
         return await self.post(f"segments/{segment_id}/targets/query", data=payload)
@@ -621,12 +634,12 @@ async def remove_segment_targets(
         )
 
     def campaign(  # type: ignore
-        self, campaign_id: Optional[str] = None, data: CampaignData = None
+        self, campaign_id: Optional[str] = None, data: Optional[CampaignData] = None
     ) -> Campaign:
         return Campaign(client=self, campaign_id=campaign_id, data=data)
 
     async def create_campaign(
-        self, campaign_id: Optional[str] = None, data: CampaignData = None
+        self, campaign_id: Optional[str] = None, data: Optional[CampaignData] = None
     ) -> StreamResponse:
         payload = {"id": campaign_id}
         if data is not None:
@@ -637,9 +650,18 @@ async def get_campaign(self, campaign_id: str) -> StreamResponse:
         return await self.get(f"campaigns/{campaign_id}")
 
     async def query_campaigns(
-        self, filter_conditions: Dict[str, Any], options: QueryCampaignsOptions = None
+        self,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
+        options: QueryCampaignsOptions = None
     ) -> StreamResponse:
-        payload = {"filter": filter_conditions, **options}
+        payload = {}
+        if filter_conditions is not None:
+            payload["filter"] = filter_conditions
+        if sort is not None:
+            payload["sort"] = sort
+        if options is not None:
+            payload.update(cast(dict, options))
         return await self.post("campaigns/query", data=payload)
 
     async def update_campaign(

diff --git a/stream_chat/async_chat/segment.py b/stream_chat/async_chat/segment.py
@@ -1,6 +1,7 @@
-from typing import Dict, Optional
+from typing import Dict, Optional, List
 
 from stream_chat.base.segment import SegmentInterface
+from stream_chat.types.base import SortParam
 from stream_chat.types.segment import QuerySegmentTargetsOptions, SegmentData
 from stream_chat.types.stream_response import StreamResponse
 
@@ -49,11 +50,13 @@ async def add_targets(self, target_ids: list) -> StreamResponse:
     async def query_targets(
         self,
         filter_conditions: Optional[Dict] = None,
+        sort: Optional[List[SortParam]] = None,
         options: Optional[QuerySegmentTargetsOptions] = None,
     ) -> StreamResponse:
         return await self.client.query_segment_targets(  # type: ignore
             segment_id=self.segment_id,
             filter_conditions=filter_conditions,
+            sort=sort,
             options=options,
         )
 

diff --git a/stream_chat/base/client.py b/stream_chat/base/client.py
@@ -7,6 +7,7 @@
 import sys
 from typing import Any, Awaitable, Dict, Iterable, List, Optional, TypeVar, Union
 
+from stream_chat.types.base import SortParam
 from stream_chat.types.campaign import CampaignData, QueryCampaignsOptions
 from stream_chat.types.segment import (
     QuerySegmentsOptions,
@@ -968,7 +969,10 @@ def get_segment(
 
     @abc.abstractmethod
     def query_segments(
-        self, filter_conditions: Dict, options: QuerySegmentsOptions
+        self,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
+        options: Optional[QuerySegmentsOptions] = None,
     ) -> Union[StreamResponse, Awaitable[StreamResponse]]:
         """
         Query segments
@@ -1015,7 +1019,8 @@ def add_segment_targets(
     def query_segment_targets(
         self,
         segment_id: str,
-        filter_conditions: Optional[Dict] = None,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
         options: Optional[QuerySegmentTargetsOptions] = None,
     ) -> Union[StreamResponse, Awaitable[StreamResponse]]:
         """
@@ -1064,7 +1069,10 @@ def get_campaign(
 
     @abc.abstractmethod
     def query_campaigns(
-        self, filter_conditions: Dict[str, Any], options: QueryCampaignsOptions = None
+        self,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
+        options: Optional[QueryCampaignsOptions] = None
     ) -> Union[StreamResponse, Awaitable[StreamResponse]]:
         """
         Query campaigns

diff --git a/stream_chat/client.py b/stream_chat/client.py
@@ -8,6 +8,7 @@
 
 from stream_chat.campaign import Campaign
 from stream_chat.segment import Segment
+from stream_chat.types.base import SortParam
 from stream_chat.types.campaign import CampaignData, QueryCampaignsOptions
 from stream_chat.types.segment import (
     QuerySegmentsOptions,
@@ -554,9 +555,18 @@ def get_segment(self, segment_id: str) -> StreamResponse:
         return self.get(f"segments/{segment_id}")
 
     def query_segments(
-        self, filter_conditions: Dict, options: QuerySegmentsOptions
+        self,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
+        options: Optional[QuerySegmentsOptions] = None
     ) -> StreamResponse:
-        payload = {"filter": filter_conditions, **options}
+        payload = {}
+        if filter_conditions is not None:
+            payload["filter"] = filter_conditions
+        if sort is not None:
+            payload["sort"] = sort
+        if options is not None:
+            payload.update(cast(dict, options))
         return self.post("segments/query", data=payload)
 
     def update_segment(self, segment_id: str, data: SegmentData) -> StreamResponse:
@@ -578,12 +588,15 @@ def add_segment_targets(
     def query_segment_targets(
         self,
         segment_id: str,
-        filter_conditions: Optional[Dict] = None,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
         options: Optional[QuerySegmentTargetsOptions] = None,
     ) -> StreamResponse:
         payload = {}
         if filter_conditions is not None:
             payload["filter"] = filter_conditions
+        if sort is not None:
+            payload["sort"] = sort
         if options is not None:
             payload.update(cast(dict, options))
         return self.post(f"segments/{segment_id}/targets/query", data=payload)
@@ -612,9 +625,18 @@ def get_campaign(self, campaign_id: str) -> StreamResponse:
         return self.get(f"campaigns/{campaign_id}")
 
     def query_campaigns(
-        self, filter_conditions: Dict[str, Any], options: QueryCampaignsOptions = None
+        self,
+        filter_conditions: Optional[Dict[str, Any]] = None,
+        sort: Optional[List[SortParam]] = None,
+        options: Optional[QueryCampaignsOptions] = None
     ) -> StreamResponse:
-        payload = {"filter": filter_conditions, **options}
+        payload = {}
+        if filter_conditions is not None:
+            payload["filter"] = filter_conditions
+        if sort is not None:
+            payload["sort"] = sort
+        if options is not None:
+            payload.update(cast(dict, options))
         return self.post("campaigns/query", data=payload)
 
     def update_campaign(self, campaign_id: str, data: CampaignData) -> StreamResponse:

diff --git a/stream_chat/segment.py b/stream_chat/segment.py
@@ -1,6 +1,7 @@
-from typing import Dict, Optional
+from typing import Dict, Optional, List
 
 from stream_chat.base.segment import SegmentInterface
+from stream_chat.types.base import SortParam
 from stream_chat.types.segment import QuerySegmentTargetsOptions, SegmentData
 from stream_chat.types.stream_response import StreamResponse
 
@@ -47,10 +48,12 @@ def add_targets(self, target_ids: list) -> StreamResponse:
     def query_targets(
         self,
         filter_conditions: Optional[Dict] = None,
+        sort: Optional[List[SortParam]] = None,
         options: Optional[QuerySegmentTargetsOptions] = None,
     ) -> StreamResponse:
         return self.client.query_segment_targets(  # type: ignore
             segment_id=self.segment_id,
+            sort=sort,
             filter_conditions=filter_conditions,
             options=options,
         )

diff --git a/stream_chat/tests/async_chat/test_campaign.py b/stream_chat/tests/async_chat/test_campaign.py
@@ -107,3 +107,48 @@ async def test_campaign_start_stop(
         assert deleted.is_ok()
 
         await client.delete_segment(segment_id=segment_id)
+
+    async def test_query_campaigns(self, client: StreamChatAsync, random_user: Dict):
+        segment_created = await client.create_segment(segment_type=SegmentType.USER)
+        segment_id = segment_created["segment"]["id"]
+
+        sender_id = random_user["id"]
+
+        target_added = await client.add_segment_targets(
+            segment_id=segment_id, target_ids=[sender_id]
+        )
+        assert target_added.is_ok()
+
+        created = await client.create_campaign(
+            data={
+                "message_template": {
+                    "text": "{Hello}",
+                },
+                "segment_ids": [segment_id],
+                "sender_id": sender_id,
+                "name": "some name",
+            }
+        )
+        assert created.is_ok()
+        assert "campaign" in created
+        assert "id" in created["campaign"]
+        assert "name" in created["campaign"]
+        campaign_id = created["campaign"]["id"]
+
+        query_campaigns = await client.query_campaigns(
+            filter_conditions={
+                "id": {
+                    "$eq": campaign_id,
+                }
+            }
+        )
+        assert query_campaigns.is_ok()
+        assert "campaigns" in query_campaigns
+        assert len(query_campaigns["campaigns"]) == 1
+        assert query_campaigns["campaigns"][0]["id"] == campaign_id
+
+        deleted = await client.delete_campaign(campaign_id=campaign_id)
+        assert deleted.is_ok()
+
+        segment_deleted = await client.delete_segment(segment_id=segment_id)
+        assert segment_deleted.is_ok()
diff --git a/stream_chat/tests/async_chat/test_segment.py b/stream_chat/tests/async_chat/test_segment.py
@@ -73,10 +73,10 @@ async def test_segment_targets(self, client: StreamChatAsync):
 
         query_targets_2 = await segment.query_targets(
             filter_conditions={"target_id": {"$lte": "<user_id>"}},
+            sort=[{"field": "target_id", "direction": SortOrder.DESC}],
             options={
                 "limit": 3,
                 "next": query_targets_1["next"],
-                "sort": [{"field": "target_id", "direction": SortOrder.DESC}],
             },
         )
         assert query_targets_2.is_ok()
@@ -89,3 +89,26 @@ async def test_segment_targets(self, client: StreamChatAsync):
 
         deleted = await segment.delete()
         assert deleted.is_ok()
+
+    async def test_query_segments(self, client: StreamChatAsync):
+        created = await client.create_segment(segment_type=SegmentType.USER)
+        assert created.is_ok()
+        assert "segment" in created
+        assert "id" in created["segment"]
+        assert "name" in created["segment"]
+        segment_id = created["segment"]["id"]
+
+        target_ids = [str(uuid.uuid4()) for _ in range(10)]
+        target_added = await client.add_segment_targets(segment_id=segment_id, target_ids=target_ids)
+        assert target_added.is_ok()
+
+        query_segments = await client.query_segments(filter_conditions={"id": {"$eq": segment_id}})
+        assert query_segments.is_ok()
+        assert "segments" in query_segments
+        assert len(query_segments["segments"]) == 1
+
+        target_deleted = await client.remove_segment_targets(segment_id=segment_id, target_ids=target_ids)
+        assert target_deleted.is_ok()
+
+        deleted = await client.delete_segment(segment_id=segment_id)
+        assert deleted.is_ok()
diff --git a/stream_chat/tests/test_campaign.py b/stream_chat/tests/test_campaign.py
@@ -62,7 +62,8 @@ def test_campaign_crud(self, client: StreamChat, random_user: Dict):
         deleted = campaign.delete()
         assert deleted.is_ok()
 
-        client.delete_segment(segment_id=segment_id)
+        segment_deleted = client.delete_segment(segment_id=segment_id)
+        assert segment_deleted.is_ok()
 
     def test_campaign_start_stop(self, client: StreamChat, random_user: Dict):
         segment = client.create_segment(segment_type=SegmentType.USER)
@@ -110,3 +111,48 @@ def test_campaign_start_stop(self, client: StreamChat, random_user: Dict):
         assert deleted.is_ok()
 
         client.delete_segment(segment_id=segment_id)
+
+    def test_query_campaigns(self, client: StreamChat, random_user: Dict):
+        segment_created = client.create_segment(segment_type=SegmentType.USER)
+        segment_id = segment_created["segment"]["id"]
+
+        sender_id = random_user["id"]
+
+        target_added = client.add_segment_targets(
+            segment_id=segment_id, target_ids=[sender_id]
+        )
+        assert target_added.is_ok()
+
+        created = client.create_campaign(
+            data={
+                "message_template": {
+                    "text": "{Hello}",
+                },
+                "segment_ids": [segment_id],
+                "sender_id": sender_id,
+                "name": "some name",
+            }
+        )
+        assert created.is_ok()
+        assert "campaign" in created
+        assert "id" in created["campaign"]
+        assert "name" in created["campaign"]
+        campaign_id = created["campaign"]["id"]
+
+        query_campaigns = client.query_campaigns(
+            filter_conditions={
+                "id": {
+                    "$eq": campaign_id,
+                }
+            }
+        )
+        assert query_campaigns.is_ok()
+        assert "campaigns" in query_campaigns
+        assert len(query_campaigns["campaigns"]) == 1
+        assert query_campaigns["campaigns"][0]["id"] == campaign_id
+
+        deleted = client.delete_campaign(campaign_id=campaign_id)
+        assert deleted.is_ok()
+
+        segment_deleted = client.delete_segment(segment_id=segment_id)
+        assert segment_deleted.is_ok()