def add_plans_to_queue():
"""
Clear the queue and add 3 fixed plans to the queue.
Raises an exception if clearing the queue or adding plans fails.
"""
resp1, _ = zmq_single_request("queue_clear")
assert resp1["success"] is True, str(resp1)
user_group = "admin"
user = "******"
plan1 = {
"name": "count",
"args": [["det1", "det2"]],
"kwargs": {
"num": 10,
"delay": 1
}
}
plan2 = {"name": "count", "args": [["det1", "det2"]]}
for plan in (plan1, plan2, plan2):
resp2, _ = zmq_single_request("queue_item_add", {
"plan": plan,
"user": user,
"user_group": user_group
})
assert resp2["success"] is True, str(resp2)
def stop_manager(self, timeout=10):
"""
Attempt to exit the subprocess that is running manager in orderly way and kill it
after timeout.
Parameters
----------
p: subprocess.Popen
Subprocess in which the manager is running
timeout: float
Timeout in seconds.
"""
if self._p:
try:
# If the process is already terminated, then don't attempt to communicate with it.
if self._p.poll() is None:
# Try to stop the manager in a nice way first by sending the command
resp, _ = zmq_single_request(method="manager_stop", params=None)
assert resp["success"] is True, f"Request to stop the manager failed: {resp['msg']}."
self._p.wait(timeout)
clear_redis_pool()
except Exception as ex:
# The manager is not responsive, so just kill the process.
self._p.kill()
clear_redis_pool()
assert False, f"RE Manager failed to stop: {str(ex)}"
self._p = None
def get_queue():
"""
Returns current queue.
"""
method, params = "queue_get", None
msg, _ = zmq_single_request(method, params)
if msg is None:
raise TimeoutError("Timeout occurred while loading queue from RE Manager.")
return msg
def test_manager_redis_addr_parameter(re_manager_cmd, redis_addr, success): # noqa: F811
if success:
re_manager_cmd(["--redis-addr", redis_addr])
# Try to communicate with the server to make sure Redis is configure correctly.
# RE Manager has to access Redis in order to prepare 'status'.
resp1, _ = zmq_single_request("status")
assert resp1["items_in_queue"] == 0
assert resp1["items_in_history"] == 0
else:
with pytest.raises(TimeoutError, match="RE Manager failed to start"):
re_manager_cmd(["--redis-addr", redis_addr])
def test_manager_acq_with_0MQ_proxy(re_manager_cmd, zmq_proxy,
zmq_dispatcher): # noqa: F811
re_manager_cmd(["--zmq-data-proxy-addr", "localhost:5567"])
# Open the environment (make sure that the environment loads)
resp1, _ = zmq_single_request("environment_open")
assert resp1["success"] is True
assert wait_for_condition(time=10, condition=condition_environment_created)
# Add the plan to the queue (will fail if incorrect environment is loaded)
params = {"item": _plan1, "user": _user, "user_group": _user_group}
resp2, _ = zmq_single_request("queue_item_add", params)
assert resp2["success"] is True, f"resp={resp2}"
# Start the queue
resp3, _ = zmq_single_request("queue_start")
assert resp3["success"] is True
assert wait_for_condition(time=10,
condition=condition_queue_processing_finished)
# Make sure that the plan was executed
resp4, _ = zmq_single_request("status")
assert resp4["items_in_queue"] == 0
assert resp4["items_in_history"] == 1
# Close the environment
resp5, _ = zmq_single_request("environment_close")
assert resp5["success"] is True, f"resp={resp5}"
assert wait_for_condition(time=5, condition=condition_environment_closed)
# Test if the data was delivered to the consumer.
# Simple test: check if 'start' and 'stop' documents were delivered.
queue = zmq_dispatcher
remote_accumulator = []
while not queue.empty(
): # Since queue is used by one process at a time, queue.empty() should work reliably
remote_accumulator.append(queue.get(timeout=2))
assert len(remote_accumulator) >= 2
assert remote_accumulator[0][0] == "start" # Start document
assert remote_accumulator[-1][0] == "stop" # Stop document
def test_fixture_re_manager_cmd_2(re_manager_cmd, db_catalog): # noqa F811
"""
Test for the fixture ``re_manager_cmd``: start RE Manager with command line parameters.
Subscribe RE to databroker (created by ``db_catalog``, execute the plan and make sure
that the run is recorded by the databroker by comparing Run UIDs from history and
start document.
"""
db_name = db_catalog["catalog_name"]
re_manager_cmd(["--databroker-config", db_name])
cat = db_catalog["catalog"]
plan = {"name": "scan", "args": [["det1", "det2"], "motor", -1, 1, 10], "item_type": "plan"}
# Plan
params1 = {"item": plan, "user": _user, "user_group": _user_group}
resp1, _ = zmq_single_request("queue_item_add", params1)
assert resp1["success"] is True, f"resp={resp1}"
resp2, _ = zmq_single_request("status")
assert resp2["items_in_queue"] == 1
assert resp2["items_in_history"] == 0
# Open the environment.
resp3, _ = zmq_single_request("environment_open")
assert resp3["success"] is True
assert wait_for_condition(time=10, condition=condition_environment_created)
resp4, _ = zmq_single_request("queue_start")
assert resp4["success"] is True
assert wait_for_condition(time=5, condition=condition_manager_idle)
resp5, _ = zmq_single_request("status")
assert resp5["items_in_queue"] == 0
assert resp5["items_in_history"] == 1
resp6, _ = zmq_single_request("history_get")
history = resp6["items"]
assert len(history) == 1
uid = history[-1]["result"]["run_uids"][0]
start_doc = cat[uid].metadata["start"]
assert start_doc["uid"] == uid
# Close the environment.
resp7, _ = zmq_single_request("environment_close")
assert resp7["success"] is True, f"resp={resp7}"
assert wait_for_condition(time=5, condition=condition_environment_closed)
def zmq_secure_request(method, params=None, *, zmq_server_address=None, server_public_key=None):
"""
Wrapper for 'zmq_single_request'. Verifies if environment variable holding server public key is set
and passes the key to 'zmq_single_request' . Simplifies writing tests that use RE Manager in secure mode.
Use functions `set_qserver_zmq_public_key()` and `clear_qserver_zmq_public_key()` to set and
clear the environment variable.
The function also verifies if the environment variable holding ZMQ server address is set, and
passes the address to ``zmq_single_request``. If ``zmq_server_address`` is passed as a parameter, then
the environment variable is ignored (at least in current implementation).
Parameters
----------
method: str
Name of the method called in RE Manager
params: dict or None
Dictionary of parameters (payload of the message). If ``None`` then
the message is sent with empty payload: ``params = {}``.
zmq_server_address: str or None
Address of the ZMQ control socket of RE Manager. An address from the environment variable or
the default address is used if the value is ``None``.
server_public_key: str or None
Server public key (z85-encoded 40 character string). The Valid public key from the server
public/private key pair must be passed if encryption is enabled at the 0MQ server side.
Communication requests will time out if the key is invalid. Exception will be raised if
the key is improperly formatted. The key from the environment or is used if the environment
variable is set, otherwise the encryption is disabled.
"""
# Use the key from env. variable if 'server_public_key' is None
server_public_key = server_public_key or os.environ.get(_name_ev_public_key, None)
# Use the address from env. variable if 'zmq_server_address' is None
zmq_server_address = zmq_server_address or os.environ.get(_name_ev_zmq_address, None)
return zmq_single_request(
method=method, params=params, zmq_server_address=zmq_server_address, server_public_key=server_public_key
)
def zmq_secure_request(method, params=None, *, zmq_server_address=None):
"""
Wrapper for 'zmq_single_request'. Verifies if environment variable holding server public key is set
and passes the key to 'zmq_single_request' . Simplifies writing tests that use RE Manager in secure mode.
Use functions `set_qserver_zmq_public_key()` and `clear_qserver_zmq_public_key()` to set and
clear the environment variable.
The function also verifies if the environment variable holding ZMQ server address is set, and
passes the address to ``zmq_single_request``. If ``zmq_server_address`` is passed as a parameter, then
the environment variable is ignored (at least in current implementation).
"""
server_public_key = None
if _name_ev_public_key in os.environ:
server_public_key = os.environ[_name_ev_public_key]
# Use the address passed in environment variable only if the parameter 'zmq_server_address' is None
if (_name_ev_zmq_address in os.environ) and (zmq_server_address is None):
zmq_server_address = os.environ[_name_ev_zmq_address]
return zmq_single_request(method=method,
params=params,
zmq_server_address=zmq_server_address,
server_public_key=server_public_key)
def test_manager_options_startup_profile(re_manager_cmd, tmp_path, monkeypatch,
option): # noqa: F811
pc_path = copy_default_profile_collection(tmp_path)
# Add extra plan. The original set of startup files will not contain this plan.
append_code_to_last_startup_file(pc_path, additional_code=_sample_plan1)
# Generate the new list of allowed plans and devices and reload them
gen_list_of_plans_and_devices(startup_dir=pc_path,
file_dir=pc_path,
overwrite=True)
# Start manager
if option == "startup_dir":
re_manager_cmd(["--startup-dir", pc_path])
elif option == "profile":
# This option is more complicated: we want to recreate the structure of IPython startup
# directory: <some root dir>/profile_<profile_name>/startup.
root_dir = os.path.split(pc_path)[0]
monkeypatch.setenv("IPYTHONDIR", root_dir)
profile_name = "testing"
startup_path = os.path.join(root_dir, f"profile_{profile_name}",
"startup")
os.makedirs(startup_path)
file_pattern = os.path.join(pc_path, "*")
for fl_path in glob.glob(file_pattern):
shutil.move(fl_path, startup_path)
os.rmdir(pc_path)
# We pass only profile name as a parameter.
re_manager_cmd(["--startup-profile", profile_name])
elif option == "multiple":
# Expected to fail if multiple options are selected.
with pytest.raises(TimeoutError, match="RE Manager failed to start"):
re_manager_cmd(
["--startup-dir", pc_path, "--startup-profile", "some_name"])
return
else:
assert False, f"Unknown option '{option}'"
# Open the environment (make sure that the environment loads)
resp1, _ = zmq_single_request("environment_open")
assert resp1["success"] is True
assert wait_for_condition(time=10, condition=condition_environment_created)
# Add the plan to the queue (will fail if incorrect environment is loaded)
plan = {"name": "simple_sample_plan", "item_type": "plan"}
params = {"item": plan, "user": _user, "user_group": _user_group}
resp2, _ = zmq_single_request("queue_item_add", params)
assert resp2["success"] is True, f"resp={resp2}"
# Start the queue
resp3, _ = zmq_single_request("queue_start")
assert resp3["success"] is True
assert wait_for_condition(time=10,
condition=condition_queue_processing_finished)
# Make sure that the plan was executed
resp4, _ = zmq_single_request("status")
assert resp4["items_in_queue"] == 0
assert resp4["items_in_history"] == 1
# Close the environment
resp5, _ = zmq_single_request("environment_close")
assert resp5["success"] is True, f"resp={resp5}"
assert wait_for_condition(time=5, condition=condition_environment_closed)
monkeypatch.setenv("IPYTHONDIR", "abc")
def get_queue_state():
method, params = "status", None
msg, _ = zmq_single_request(method, params)
if msg is None:
raise TimeoutError("Timeout occurred while reading RE Manager status.")
return msg
