def test_mutator_exceptions():
handled = False
def base_plan():
yield Msg('foo')
yield Msg('bar')
def failing_plan():
nonlocal handled
handled = False
yield Msg('pre')
try:
yield Msg('FAIL')
except EchoException:
handled = True
raise
def test_mutator(msg):
if msg.command == 'bar':
return (failing_plan(), single_message_gen(Msg('foo')))
return None, None
# check generator exit behavior
plan = plan_mutator(base_plan(), test_mutator)
next(plan)
plan.close()
# check exception fall through
plan = plan_mutator(base_plan(), test_mutator)
with pytest.raises(EchoException):
EchoRE(plan, debug=True)
assert handled
def test_exception_in_pre_with_tail():
class SnowFlake(Exception):
...
def bad_pre():
yield Msg('pre_bad', None)
raise SnowFlake('this one')
def good_post():
yield Msg('good_post', None)
def test_mutator(msg):
if msg.command == 'TARGET':
return bad_pre(), good_post()
return None, None
def testing_plan():
yield Msg('a', None)
yield Msg('b', None)
try:
yield Msg('TARGET', None)
except SnowFlake:
pass
yield Msg('b', None)
yield Msg('a', None)
plan = plan_mutator(testing_plan(), test_mutator)
msgs = EchoRE(plan, debug=True)
_verify_msg_seq(msgs,
m_len=5,
cmd_sq=['a', 'b', 'pre_bad', 'b', 'a'],
args_sq=[()] * 5,
kwargs_sq=[{}] * 5)
def test_plan_mutator_returns():
def testing_plan():
yield Msg('a', None)
yield Msg('TARGET', None)
yield Msg('b', None)
return 'foobar'
def outer_plan(pln):
ret = (yield from pln)
assert ret == 'foobar'
return ret
def tail_plan():
yield Msg('A', None)
return 'baz'
def test_mutator(msg):
def pre_plan():
yield Msg('pre', None)
yield msg
if msg.command == 'TARGET':
return pre_plan(), tail_plan()
return None, None
plan = plan_mutator(testing_plan(), test_mutator)
msgs = EchoRE(plan)
_verify_msg_seq(msgs,
m_len=5,
cmd_sq=['a', 'pre', 'TARGET', 'A', 'b'],
args_sq=[()] * 5,
kwargs_sq=[{}] * 5)
def test_simple_mutator():
_mut_active = True
pre_count = 3
post_count = 5
pre_cmd = 'pre'
post_cmd = 'post'
def test_mutator(msg):
nonlocal _mut_active
if _mut_active:
_mut_active = False
return (pchain(echo_plan(num=pre_count, command=pre_cmd),
single_message_gen(msg)),
echo_plan(num=post_count, command=post_cmd))
return None, None
num = 5
cmd = 'echo'
plan = plan_mutator(echo_plan(command=cmd, num=num), test_mutator)
msgs = EchoRE(plan)
total = num + pre_count + post_count
cmd_sq = ([pre_cmd] * pre_count + [cmd] + [post_cmd] * post_count + [cmd] *
(num - 1))
_verify_msg_seq(msgs,
m_len=total,
cmd_sq=cmd_sq,
args_sq=[()] * total,
kwargs_sq=[{}] * total)
def test_plan_mutator_exception_propogation():
class ExpectedException(Exception):
pass
num = 5
cmd1 = 'echo1'
cmd2 = 'echo2'
def bad_tail():
yield Msg('one_tail', None)
raise ExpectedException('this is a test')
def sarfing_plan():
try:
yield from echo_plan(command=cmd1, num=num)
except ExpectedException:
print('CAUGHT IT')
_mut_active = True
def test_mutator(msg):
nonlocal _mut_active
if _mut_active:
_mut_active = False
return (pchain(echo_plan(num=2, command=cmd2),
single_message_gen(msg)), bad_tail())
return None, None
plan = plan_mutator(sarfing_plan(), test_mutator)
EchoRE(plan, debug=True)
def test_plan_mutator_exception_propogation():
class ExpectedException(Exception):
pass
num = 5
cmd1 = "echo1"
cmd2 = "echo2"
def bad_tail():
yield Msg("one_tail", None)
raise ExpectedException("this is a test")
def sarfing_plan():
try:
yield from echo_plan(command=cmd1, num=num)
except ExpectedException:
print("CAUGHT IT")
_mut_active = True
def test_mutator(msg):
nonlocal _mut_active
if _mut_active:
_mut_active = False
return (pchain(echo_plan(num=2, command=cmd2), single_message_gen(msg)), bad_tail())
return None, None
plan = plan_mutator(sarfing_plan(), test_mutator)
EchoRE(plan, debug=True)
def test_simple_mutator():
_mut_active = True
pre_count = 3
post_count = 5
pre_cmd = "pre"
post_cmd = "post"
def test_mutator(msg):
nonlocal _mut_active
if _mut_active:
_mut_active = False
return (
pchain(echo_plan(num=pre_count, command=pre_cmd), single_message_gen(msg)),
echo_plan(num=post_count, command=post_cmd),
)
return None, None
num = 5
cmd = "echo"
plan = plan_mutator(echo_plan(command=cmd, num=num), test_mutator)
msgs = EchoRE(plan)
total = num + pre_count + post_count
cmd_sq = [pre_cmd] * pre_count + [cmd] + [post_cmd] * post_count + [cmd] * (num - 1)
_verify_msg_seq(msgs, m_len=total, cmd_sq=cmd_sq, args_sq=[()] * total, kwargs_sq=[{}] * total)
def test_base_excetpion():
class SnowFlake(Exception):
...
def null_mutator(msg):
return None, None
def test_plan():
yield Msg('a', None)
raise SnowFlake('this one')
pln = plan_mutator(test_plan(), null_mutator)
try:
EchoRE(pln)
except SnowFlake as ex:
assert ex.args[0] == 'this one'
def configure_count_time_wrapper(plan, time):
"""
Preprocessor that sets all devices with a `count_time` to the same time.
The original setting is stashed and restored at the end.
Parameters
----------
plan : iterable or iterator
a generator, list, or similar containing `Msg` objects
time : float or None
If None, the plan passes through unchanged.
Yields
------
msg : Msg
messages from plan, with 'set' messages inserted
"""
devices_seen = set()
original_times = {}
def insert_set(msg):
obj = msg.obj
if obj is not None and obj not in devices_seen:
devices_seen.add(obj)
if hasattr(obj, 'count_time'):
# TODO Do this with a 'read' Msg once reads can be
# marked as belonging to a different event stream (or no
# event stream.
original_times[obj] = obj.count_time.get()
# TODO do this with configure
return pchain(mv(obj.count_time, time),
single_gen(msg)), None
return None, None
def reset():
for obj, time in original_times.items():
yield from mv(obj.count_time, time)
if time is None:
# no-op
return (yield from plan)
else:
return (yield from finalize_wrapper(plan_mutator(plan, insert_set),
reset()))
def test_insert_after():
def target():
yield Msg('a', None)
ret = yield Msg('TARGET', None)
yield Msg('b', None)
assert ret is not None
assert ret.command == 'TARGET'
return ret
def insert_after(msg):
if msg.command == 'TARGET':
def post():
yield Msg('post', None)
return None, post()
else:
return None, None
ret = EchoRE(plan_mutator(target(), insert_after))
def periodic_dark(plan):
"""
a plan wrapper that takes a plan and inserts `take_dark`
The `take_dark` plan is inserted on the fly before the beginning of
any new run after a period of time defined by `glbl.dk_window` has passed.
"""
need_dark = True
def insert_take_dark(msg):
now = time.time()
nonlocal need_dark
qualified_dark_uid = _validate_dark(expire_time=glbl.dk_window)
# FIXME: should we do "or" or "and"?
if (not need_dark) and (not qualified_dark_uid):
need_dark = True
if need_dark \
and (not qualified_dark_uid) \
and msg.command == 'open_run' \
and ('dark_frame' not in msg.kwargs):
# We are about to start a new 'run' (e.g., a count or a scan).
# Insert a dark frame run first.
need_dark = False
# Annoying detail: the detector was probably already staged.
# Unstage it (if it wasn't staged, nothing will happen) and
# then take_dark() and then re-stage it.
return bp.pchain(bp.unstage(glbl.area_det),
take_dark(),
bp.stage(glbl.area_det),
bp.single_gen(msg),
bp.abs_set(glbl.shutter, 60, wait=True)), None
elif msg.command == 'open_run' and 'dark_frame' not in msg.kwargs:
return bp.pchain(bp.single_gen(msg),
bp.abs_set(glbl.shutter, 60, wait=True)), None
else:
# do nothing if (not need_dark)
return None, None
return (yield from bp.plan_mutator(plan, insert_take_dark))
def periodic_dark(plan):
"""
a plan wrapper that takes a plan and inserts `take_dark`
The `take_dark` plan is inserted on the fly before the beginning of
any new run after a period of time defined by `glbl.dk_window` has passed.
"""
need_dark = True
def insert_take_dark(msg):
now = time.time()
nonlocal need_dark
qualified_dark_uid = _validate_dark(expire_time=glbl.dk_window)
# FIXME: should we do "or" or "and"?
if (not need_dark) and (not qualified_dark_uid):
need_dark = True
if need_dark \
and (not qualified_dark_uid) \
and msg.command == 'open_run' \
and ('dark_frame' not in msg.kwargs):
# We are about to start a new 'run' (e.g., a count or a scan).
# Insert a dark frame run first.
need_dark = False
# Annoying detail: the detector was probably already staged.
# Unstage it (if it wasn't staged, nothing will happen) and
# then take_dark() and then re-stage it.
return bp.pchain(bp.unstage(glbl.area_det), take_dark(),
bp.stage(glbl.area_det), bp.single_gen(msg),
bp.abs_set(glbl.shutter, 60, wait=True)), None
elif msg.command == 'open_run' and 'dark_frame' not in msg.kwargs:
return bp.pchain(bp.single_gen(msg),
bp.abs_set(glbl.shutter, 60, wait=True)), None
else:
# do nothing if (not need_dark)
return None, None
return (yield from bp.plan_mutator(plan, insert_take_dark))
def test_insert_before():
def target():
yield Msg('a', None)
ret = yield Msg('TARGET', None)
yield Msg('b', None)
assert ret.command == 'TARGET'
return ret
return ret
def insert_before(msg):
if msg.command == 'TARGET':
def pre():
yield Msg('pre', None)
ret = yield msg
assert ret is not None
assert ret.command == 'TARGET'
return ret
return pre(), None
else:
return None, None
ret = EchoRE(plan_mutator(target(), insert_before))
