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