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