tests ensure saved_location_id is sent as expected where relevant.

Beam location now handles either a Location or a location_id

Author: Paul Duncan

predicthq/endpoints/v1/beam/schemas.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel, Field, ConfigDict
 from datetime import datetime
 from predicthq.endpoints.schemas import ArgKwargResultSet
-from typing import Optional, List
+from typing import Optional, List, Union
 
 # Python < 3.11 does not have StrEnum in the enum module
 import sys
@@ -61,6 +61,10 @@ class Location(BaseModel):
     google_place_id: Optional[str] = None
 
 
+class LocationId(BaseModel):
+    saved_location_id: str
+
+
 class RankLevel(BaseModel):
     model_config: ConfigDict = ConfigDict(extra="allow")
 
@@ -192,7 +196,7 @@ class Analysis(BaseModel):
 
     analysis_id: Optional[str] = None
     name: str
-    location: Location
+    location: Union[Location, LocationId]
     rank: Rank
     user_id: Optional[str] = None
     access_type: str

tests/endpoints/v1/test_events.py

@@ -44,6 +44,7 @@ def test_search_params_underscores(self, client):
             updated__gte="2016-03-01",
             updated__lt="2016-04-01",
             updated__tz="Pacific/Auckland",
+            saved_location__location_id="UVel7RGxM7U2ISE1AQOwJQ",
         )
 
         client.request.assert_called_once_with(
@@ -80,6 +81,7 @@ def test_search_params_underscores(self, client):
                 "updated.gte": "2016-03-01",
                 "updated.lt": "2016-04-01",
                 "updated.tz": "Pacific/Auckland",
+                "saved_location.location_id": "UVel7RGxM7U2ISE1AQOwJQ",
             },
             verify=True,
         )
@@ -165,6 +167,7 @@ def test_search_params_dicts(self, client):
             end={"gte": "2016-05-01", "lt": "2016-06-01", "tz": "Pacific/Auckland"},
             active={"gte": "2016-03-01", "lt": "2016-04-01", "tz": "Pacific/Auckland"},
             updated={"gte": "2016-03-01", "lt": "2016-04-01", "tz": "Pacific/Auckland"},
+            saved_location={"location_id": "UVel7RGxM7U2ISE1AQOwJQ"},
         )
 
         client.request.assert_called_once_with(
@@ -200,6 +203,7 @@ def test_search_params_dicts(self, client):
                 "updated.gte": "2016-03-01",
                 "updated.lt": "2016-04-01",
                 "updated.tz": "Pacific/Auckland",
+                "saved_location.location_id": "UVel7RGxM7U2ISE1AQOwJQ",
             },
             verify=True,
         )

tests/endpoints/v1/test_features.py

@@ -865,3 +865,69 @@ def test_spend_request_params_dicts(self, client):
             verify=True,
             params=None,
         )
+
+    @with_mock_client(
+        request_returns=load_fixture(
+            "requests_responses/features_test/test_empty_search"
+        )
+    )
+    def test_saved_location_ids_params_dicts(self, client):
+        client.features.obtain_features(
+            active={"gte": "2017-12-31", "lte": "2018-01-02"},
+            location={
+                "saved_location_id": [
+                    "BN7ZSw8xza9FviPVfyCycd",
+                    "X3uyTFbDOUaX2q_Qh5i31b",
+                    "X3uyTFbDOUhrfq_Qh5i31A",
+                ]
+            },
+        )
+
+        client.request.assert_called_once_with(
+            "post",
+            "/v1/features/",
+            json={
+                "active": {"gte": "2017-12-31", "lte": "2018-01-02"},
+                "location": {
+                    "saved_location_id": [
+                        "BN7ZSw8xza9FviPVfyCycd",
+                        "X3uyTFbDOUaX2q_Qh5i31b",
+                        "X3uyTFbDOUhrfq_Qh5i31A",
+                    ]
+                },
+            },
+            verify=True,
+            params=None,
+        )
+
+    @with_mock_client(
+        request_returns=load_fixture(
+            "requests_responses/features_test/test_empty_search"
+        )
+    )
+    def test_saved_location_ids_params_undersores(self, client):
+        client.features.obtain_features(
+            active={"gte": "2017-12-31", "lte": "2018-01-02"},
+            location__saved_location_id=[
+                "BN7ZSw8xza9FviPVfyCycd",
+                "X3uyTFbDOUaX2q_Qh5i31b",
+                "X3uyTFbDOUhrfq_Qh5i31A",
+            ],
+        )
+
+        client.request.assert_called_once_with(
+            "post",
+            "/v1/features/",
+            json={
+                "active": {"gte": "2017-12-31", "lte": "2018-01-02"},
+                "location": {
+                    "saved_location_id": [
+                        "BN7ZSw8xza9FviPVfyCycd",
+                        "X3uyTFbDOUaX2q_Qh5i31b",
+                        "X3uyTFbDOUhrfq_Qh5i31A",
+                    ]
+                },
+            },
+            verify=True,
+            params=None,
+        )