def test_wrong_geography_aggregation_unit_raises_error():
    """
    Test that an invalid aggregation unit in a geography query raises an InvalidGeographyError
    """
    with pytest.raises(
        ValidationError,
        match="aggregation_unit.*Must be one of: admin0, admin1, admin2, admin3.",
    ):
        _ = FlowmachineQuerySchema().load(
            {"query_kind": "geography", "aggregation_unit": "DUMMY_AGGREGATION_UNIT"}
        )
    def __call__(self, value) -> Union[None, str]:
        from flowmachine.core.server.query_schemas import FlowmachineQuerySchema

        if (value is not None) and (value is not missing):
            try:
                (FlowmachineQuerySchema().load(
                    QueryInfoLookup(get_redis()).get_query_params(
                        value))._flowmachine_query_obj)
            except UnkownQueryIdError:
                if not cache_table_exists(get_db(), value):
                    raise ValidationError("Must be None or a valid query id.")

        return value
def test_invalid_sampling_params_raises_error(sampling, message):
    query_spec = {
        "query_kind": "spatial_aggregate",
        "locations": {
            "query_kind": "daily_location",
            "date": "2016-01-01",
            "aggregation_unit": "admin3",
            "method": "last",
            "sampling": sampling,
        },
    }
    with pytest.raises(ValidationError, match=message) as exc:
        _ = FlowmachineQuerySchema().load(query_spec)
    print(exc)
async def test_rerun_query_after_cancelled(server_config, real_connections):
    """
    Test that a query can be rerun after it has been cancelled.
    """
    query_obj = (FlowmachineQuerySchema().load(
        dict(
            query_kind="spatial_aggregate",
            locations=dict(
                query_kind="daily_location",
                date="2016-01-01",
                method="last",
                aggregation_unit="admin3",
            ),
        ))._flowmachine_query_obj)
    query_id = query_obj.query_id
    qsm = QueryStateMachine(get_redis(), query_id, get_db().conn_id)
    qsm.enqueue()
    qsm.cancel()
    assert not query_obj.is_stored
    assert qsm.is_cancelled
    query_info_lookup = QueryInfoLookup(get_redis())
    query_info_lookup.register_query(
        query_id,
        dict(
            query_kind="spatial_aggregate",
            locations=dict(
                query_kind="daily_location",
                date="2016-01-01",
                method="last",
                aggregation_unit="admin3",
            ),
        ),
    )

    msg = await action_handler__run_query(
        config=server_config,
        query_kind="spatial_aggregate",
        locations=dict(
            query_kind="daily_location",
            date="2016-01-01",
            method="last",
            aggregation_unit="admin3",
        ),
    )
    assert msg["status"] == ZMQReplyStatus.SUCCESS
    qsm.wait_until_complete()
    assert query_obj.is_stored
    def deserialize(
        self,
        value: typing.Any,
        attr: str = None,
        data: typing.Mapping[str, typing.Any] = None,
        **kwargs,
    ) -> Union[None, Table]:
        from flowmachine.core.server.query_schemas import FlowmachineQuerySchema

        table_name = super().deserialize(value, attr, data, **kwargs)
        if (table_name is missing) or (table_name is None):
            return table_name
        else:
            try:
                return (FlowmachineQuerySchema().load(
                    QueryInfoLookup(get_redis()).get_query_params(
                        value))._flowmachine_query_obj)
            except UnkownQueryIdError:
                return get_query_object_by_id(get_db(), value)
 def get_query_id_for_query_spec(query_spec):
     return FlowmachineQuerySchema().load(query_spec).query_id
def test_construct_query(expected_md5, query_spec):
    """
    Test that expected query objects are constructed by construct_query_object
    """
    obj = FlowmachineQuerySchema().load(query_spec)
    assert expected_md5 == obj.query_id