diff --git a/predicthq/endpoints/v1/events/schemas.py b/predicthq/endpoints/v1/events/schemas.py index 7c627fe..0d4b5b4 100644 --- a/predicthq/endpoints/v1/events/schemas.py +++ b/predicthq/endpoints/v1/events/schemas.py @@ -60,6 +60,11 @@ class ImpactPattern(BaseModel): impacts: List[ImpactPatternImpacts] +class PHQLabels(BaseModel): + label: str + weight: float + + class Event(BaseModel): cancelled: Optional[datetime] = None category: str @@ -76,6 +81,7 @@ class Event(BaseModel): labels: List[str] location: Optional[Tuple[float, float]] = None parent_event: Optional[ParentEvent] = None + phq_labels: Optional[List[PHQLabels]] = None place_hierarchies: Optional[List[List[str]]] = None postponed: Optional[datetime] = None relevance: Optional[float] = None diff --git a/predicthq/version.py b/predicthq/version.py index f5f41e5..1173108 100644 --- a/predicthq/version.py +++ b/predicthq/version.py @@ -1 +1 @@ -__version__ = "3.1.0" +__version__ = "3.2.0" diff --git a/tests/endpoints/test_schemas.py b/tests/endpoints/test_schemas.py index 5270d1b..9564994 100644 --- a/tests/endpoints/test_schemas.py +++ b/tests/endpoints/test_schemas.py @@ -5,7 +5,7 @@ from predicthq.endpoints import decorators, schemas from predicthq.endpoints.oauth2.schemas import AccessToken -from predicthq.endpoints.v1.events.schemas import Event +from predicthq.endpoints.v1.events.schemas import PHQLabels from predicthq.endpoints.v1.places.schemas import Place from predicthq.endpoints.base import BaseEndpoint @@ -33,6 +33,23 @@ def test_place_schema(): assert Place(id="some_id", type="some_type", name="some_name", location=[32.123, -84.123]).location == (32.123, -84.123) +@pytest.mark.parametrize("phq_labels,raise_validation_error", [({"label": 34, "weight": "holiday"}, True), # wrong type + ({"label": "holiday", "weight": "holiday"}, True), # wrong type + ({"label": 34, "weight": 6}, True), # wrong type + ({"weight": 2.0}, True), # missing label + ({"label": "holiday"}, True), # missing weight + ({"label": "some_string", "weight": 2.0}, False), # correct + ({"label": "another_string", "weight": 7}, False)]) # correct +def test_phq_label_schema(phq_labels, raise_validation_error): + if raise_validation_error: + with pytest.raises(ValidationError) as e: + label = PHQLabels.parse_obj(phq_labels) + else: + label = PHQLabels.parse_obj(phq_labels) + assert label.label == phq_labels["label"] + assert label.weight == phq_labels["weight"] + + def test_resultset(): class ResultExample(BaseModel): value: int