def test_request_removal(daq_robot):
"""Make some requests, and then verify that old ones get removed."""
context = None
with mock.patch('time.time') as mock_time:
service = DataAcquisitionPluginService(daq_robot, single_capability,
success_plugin_impl)
mock_time.return_value = 0
response1 = service.AcquirePluginData(make_single_request('action 1'),
context)
assert response1.request_id > 0
while True:
feedback_response = service.GetStatus(
data_acquisition_pb2.GetStatusRequest(
request_id=response1.request_id), context)
if feedback_response.status == feedback_response.STATUS_COMPLETE:
break
mock_time.return_value = 100
# The current cleanup triggering happens when a request arrives.
response2 = service.AcquirePluginData(make_single_request('action 1'),
context)
assert response2.request_id != response1.request_id
feedback_response = service.GetStatus(
data_acquisition_pb2.GetStatusRequest(
request_id=response1.request_id), context)
print(feedback_response)
assert feedback_response.status == feedback_response.STATUS_REQUEST_ID_DOES_NOT_EXIST
def get_status_async(self, request_id, **kwargs):
"""Async version of the get_status() RPC."""
request = data_acquisition.GetStatusRequest(request_id=request_id)
return self.call_async(self._stub.GetStatus,
request,
error_from_response=_get_status_error,
**kwargs)
def test_cancel_plugin(daq_robot):
"""Test that cancelling a request actually stops it."""
def run_forever(request, store_helper: DataAcquisitionStoreHelper):
while True:
store_helper.cancel_check()
time.sleep(0.01)
service = DataAcquisitionPluginService(daq_robot, single_capability,
run_forever)
context = None
response = service.AcquirePluginData(make_single_request('action 1'),
context)
assert response.request_id > 0
feedback_request = data_acquisition_pb2.GetStatusRequest(
request_id=response.request_id)
feedback = service.GetStatus(feedback_request, context)
assert feedback.status == feedback.STATUS_ACQUIRING
cancel_response = service.CancelAcquisition(
data_acquisition_pb2.CancelAcquisitionRequest(
request_id=response.request_id), context)
assert cancel_response.status == cancel_response.STATUS_OK
feedback = service.GetStatus(feedback_request, context)
assert feedback.status in (feedback.STATUS_CANCEL_IN_PROGRESS,
feedback.STATUS_ACQUISITION_CANCELLED)
service.executor.shutdown()
feedback = service.GetStatus(feedback_request, context)
assert feedback.status == feedback.STATUS_ACQUISITION_CANCELLED
def test_data_error(daq_robot):
"""Test that a data collection can report a data error."""
def error_impl(request, store_helper: DataAcquisitionStoreHelper):
data_id = data_acquisition_pb2.DataIdentifier(
action_id=request.action_id,
channel=single_capability[0].channel_name)
store_helper.state.add_errors(
[make_error(data_id, 'Failed to acquire data.')])
store_helper.wait_for_stores_complete()
service = DataAcquisitionPluginService(daq_robot, single_capability,
error_impl)
context = None
request = make_single_request('action 1')
response = service.AcquirePluginData(request, context)
assert response.request_id > 0
# Wait until the thread has run to completion
service.executor.shutdown()
feedback = service.GetStatus(
data_acquisition_pb2.GetStatusRequest(request_id=response.request_id),
context)
print(feedback)
assert feedback.status == feedback.STATUS_DATA_ERROR
assert feedback.data_errors[0].data_id.action_id == request.action_id
assert feedback.data_errors[0].data_id.channel == single_capability[
0].channel_name
def test_simple_plugin(daq_robot):
"""Test that a basic plugin that completes right away works."""
service = DataAcquisitionPluginService(daq_robot, single_capability,
success_plugin_impl)
context = None
response = service.AcquirePluginData(make_single_request('action 1'),
context)
assert response.request_id > 0
# Wait until the thread has run to completion
service.executor.shutdown()
feedback = service.GetStatus(
data_acquisition_pb2.GetStatusRequest(request_id=response.request_id),
context)
assert feedback.status == feedback.STATUS_COMPLETE
def test_plugin_stages(daq_robot):
"""Test that we can transition through the stages acquiring->saving->complete"""
# Used to coordinate between the service and client rather than relying on sleeps.
go = threading.Event()
done = threading.Event()
def collect_data(request, store_helper: DataAcquisitionStoreHelper):
while not go.wait(timeout=1):
print('Waiting on go 1')
store_helper.cancel_check()
print('Go 1 Finished.')
go.clear()
store_helper.state.set_status(
data_acquisition_pb2.GetStatusResponse.STATUS_SAVING)
done.set()
while not go.wait(timeout=1):
print('Waiting on go 2')
store_helper.cancel_check()
print('Go 2 Finished.')
store_helper.wait_for_stores_complete()
service = DataAcquisitionPluginService(daq_robot, single_capability,
collect_data)
context = None
response = service.AcquirePluginData(make_single_request('action 1'),
context)
assert response.request_id > 0
try:
feedback_request = data_acquisition_pb2.GetStatusRequest(
request_id=response.request_id)
feedback = service.GetStatus(feedback_request, context)
assert feedback.status == feedback.STATUS_ACQUIRING
go.set()
assert done.wait(timeout=1)
done.clear()
feedback = service.GetStatus(feedback_request, context)
assert feedback.status == feedback.STATUS_SAVING
go.set()
service.executor.shutdown(wait=True)
feedback = service.GetStatus(feedback_request, context)
assert feedback.status == feedback.STATUS_COMPLETE
finally:
# If anything fails, make sure that we kill the background thread.
service.CancelAcquisition(
data_acquisition_pb2.CancelAcquisitionRequest(
request_id=response.request_id), context)
service.executor.shutdown(wait=False)
def get_status(self, request_id, **kwargs):
"""Check the status of a data acquisition based on the request id.
Args:
request_id (int): The request id associated with an AcquireData request.
Raises:
RpcError: Problem communicating with the robot.
RequestIdDoesNotExistError: The request id provided is incorrect.
Returns:
If the RPC is successful, then it will return the full status response, which includes the
status as well as other information about any possible errors.
"""
request = data_acquisition.GetStatusRequest(request_id=request_id)
return self.call(self._stub.GetStatus, request, error_from_response=_get_status_error,
**kwargs)
def test_buggy_plugin(daq_robot):
"""Test that a data collection function that throws an exception returns a DATA_ERROR"""
def buggy(*args):
raise Exception("I'm broken, yo.")
service = DataAcquisitionPluginService(daq_robot, single_capability, buggy)
context = None
response = service.AcquirePluginData(make_single_request('action 1'),
context)
assert response.request_id > 0
# Wait until the thread has run to completion
service.executor.shutdown()
feedback = service.GetStatus(
data_acquisition_pb2.GetStatusRequest(request_id=response.request_id),
context)
print(feedback)
assert feedback.status == feedback.STATUS_INTERNAL_ERROR
