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