diff --git a/mxcubecore/HardwareObjects/MAXIV/ISPyBLims.py b/mxcubecore/HardwareObjects/MAXIV/ISPyBLims.py index d263d8d7ff..e1e88a9e43 100644 --- a/mxcubecore/HardwareObjects/MAXIV/ISPyBLims.py +++ b/mxcubecore/HardwareObjects/MAXIV/ISPyBLims.py @@ -1,241 +1,96 @@ -import json -import logging -from json.decoder import JSONDecodeError -from typing import ( - Dict, - List, - Optional, +# ruff: noqa: TD003, FIX002, ERA001 +from mxcubecore.HardwareObjects.abstract.PyISPyBRestClient import PyISPyBRestClient +from mxcubecore.HardwareObjects.MAXIV.MAXIVPyISPYyBDataAdapter import ( + LAZY_SESSION_PREFIX, + MAXIVPyISPyBDataAdapter, ) -from urllib.parse import urljoin - -import requests -from suds import WebFault - -from mxcubecore.HardwareObjects.abstract.ISPyBDataAdapter import ISPyBDataAdapter from mxcubecore.HardwareObjects.UserTypeISPyBLims import UserTypeISPyBLims -from mxcubecore.model.lims_session import ( - LimsSessionManager, - Session, -) - -log = logging.getLogger("ispyb_client") - -LAZY_SESSION_PREFIX = "lazy" - - -def _get_lazy_session_id(proposal: Dict) -> str: - prop_id = proposal["proposalId"] - return f"{LAZY_SESSION_PREFIX}{prop_id}" - - -def _is_lazy_session_id(session_id: str) -> bool: - return session_id.startswith(LAZY_SESSION_PREFIX) - - -def _check_ispyb_error_message(response): - def _expected_ispyb_err_msg(error_msg): - import re - - match = re.match("^JBAS011843: Failed instantiate.*ldap.*ispyb", error_msg) - return match is not None - - # - # check that we got the 'expected' error message on invalid credentials, - # otherwise log the error message, so we don't swallow new error messages - # - if _expected_ispyb_err_msg(response.text): - # all is fine - return - - log.warning( - "unexpected response from ISPyB\n" - + f"{response.status_code} {response.reason}\n{response.text}" - ) - - -class ISPyBRestClient: - def __init__(self, rest_root: str): - self._rest_root = rest_root - - def authenticate(self, user_name: str, password: str): - """ - authenticate with REST services +from mxcubecore.model.lims_session import Session - Args: - user_name: Username - password: Password - """ - token = None - auth_url = urljoin(self._rest_root, "authenticate?site=MAXIV") - response = requests.post( - auth_url, data={"login": user_name, "password": password} - ) - - try: - # if authentication is successful, we will get - # JSON response containing an auth token - token = response.json().get("token", None) - except JSONDecodeError: - # on invalid credentials, some ISPyB systems will reply with - # an internal error message, as plain text - _check_ispyb_error_message(response) - - if token is None: - # we failed to obtain the auth token, thus we failed to authenticate - raise Exception("invalid credentials") - - -def _create_session_object(proposal, session_id: str, beamline_name: str) -> Session: - return Session( - proposal_id=proposal["proposalId"], - code=proposal["code"], - number=proposal["number"], - session_id=session_id, - beamline_name=beamline_name, - title=proposal["title"], - # - # At MAXIV we don't care if a session is scheduled - # or not, mark all sessions as scheduled. - # - is_scheduled_time=True, - is_scheduled_beamline=True, - ) - - -class CustomISPyBDataAdapter(ISPyBDataAdapter): - """ - Extend the standard ISPyB data adapter with MAXIV specific logic of how to - deal with proposal sessions. - """ - - def _get_proposals(self, username: str): - proposals = json.loads( - self._shipping.service.findProposalsByLoginName(username) - ) - - for proposal in proposals: - if proposal["type"].upper() not in ["MX", "MB"]: - continue - if proposal.get("state", "Open") != "Open": - continue - - yield proposal - - def _get_sessions(self, username: str, beamline_name: str) -> List[Session]: - def list_sessions(): - for proposal in self._get_proposals(username): - sessions = self._collection.service.findSessionsByProposalAndBeamLine( - proposal["code"], proposal["number"], beamline_name - ) - - for sesssion in sessions: - yield _create_session_object( - proposal, sesssion["sessionId"], beamline_name - ) - - # - # A hack to lazily create new sessions. - # - # At MAXIV we don't schedule sessions for proposals ahead of time. Instead, we - # lazily create them as needed. - # - # If a proposal does not contain any active session, create a Session object - # with a special session ID. - # - # If user selects such a session, then we will ask ISPyB to create this session. - # - if len(sessions) == 0: - yield _create_session_object( - proposal, _get_lazy_session_id(proposal), beamline_name - ) - - return sorted(list_sessions(), key=lambda s: f"{s.code}{s.number}") - - def get_sessions_by_username( - self, username: str, beamline_name: str - ) -> LimsSessionManager: - try: - sessions = list(self._get_sessions(username, beamline_name)) - return LimsSessionManager(sessions=sessions) - except WebFault as e: - log.exception(e.message) +class NoSessionException(Exception): + """Exception raised when no expected session found.""" class ISPyBLims(UserTypeISPyBLims): + def __init__(self, name: str): + super().__init__(name) + self._duo_api_url = "" + def init(self): + self._duo_api_url: str = self.get_property("duo_api_url") + pyispyb_rest_root = self.get_property("pyispyb_rest_root") + self._rest_client = PyISPyBRestClient(pyispyb_rest_root) super().init() - self._rest_client = ISPyBRestClient(self.get_property("rest_root")) - - def _create_data_adapter(self) -> ISPyBDataAdapter: - return CustomISPyBDataAdapter( - self.ws_root.strip(), - self.proxy, - self.ws_username, - self.ws_password, + def _create_data_adapter(self) -> MAXIVPyISPyBDataAdapter: + return MAXIVPyISPyBDataAdapter( + self._duo_api_url, + self._rest_client, self.beamline_name, ) def ispyb_login(self, user_name: str, password: str): + """Authenticate with ISPyB REST services. + + In fact password is an access token obtained from Keycloak, but we call it + password to keep the interface consistent with other LIMS implementations. + + Args: + user_name: Username to authenticate with + password: Password (access token) to authenticate with + Returns: + Tuple[bool, Optional[str]]: A tuple containing a boolean indicating + success or failure, and an optional error message. + """ try: - self._rest_client.authenticate(user_name, password) - return True, None + self._rest_client.authenticate(user_name, token=password) except Exception as ex: return False, str(ex) + return True, None def set_active_session_by_id(self, session_id: str) -> Session: - """ - Sets session with session_id to active session + """Sets session with session_id to active session. + + It is possible that user picks the session that does not exist in the database yet, so called lazy session created on the fly. In that case the session POST request is sent to the server in order to create the session in the database and get the proper session id. Args: session_id: session id """ - - def find_session() -> Optional[Session]: - for session in self.session_manager.sessions: - if session.session_id == session_id: - self.session_manager.active_session = session - - return session - - # session not found - return None - - def replace_lazy(sessions: List[Session], new_session: Session): - def gen(): - for session in sessions: - if session.session_id == session_id: - yield new_session - else: - yield session - - return list(gen()) - - session = find_session() - if session is None: - raise Exception(f"no session with ID {session_id} found") - - # - # user selected a session that does not exist yet, - # ask ISPyB to create it - # - if _is_lazy_session_id(session_id): - session = self.adapter.create_session( - session.proposal_id, session.beamline_name - ) - # replace the old lazy-session object, - # with the new proper-session object - self.session_manager.sessions = replace_lazy( - self.session_manager.sessions, session - ) - - return session + session_to_activate = None + for _idx, session in enumerate(self.session_manager.sessions): + if session.session_id == session_id: + session_to_activate = session + break + else: + err = f"No session with ID {session_id} found." + raise NoSessionException(err) + + if session_id.startswith(LAZY_SESSION_PREFIX): + payload = { + "proposalId": session_to_activate.proposal_id, + "startDate": session_to_activate.start_datetime.strftime( + "%Y-%m-%dT%H:%M:%S.%fZ" + ), + "endDate": session_to_activate.end_datetime.strftime( + "%Y-%m-%dT%H:%M:%S.%fZ" + ), + "beamLineName": session_to_activate.beamline_name, + "comments": "Session created by the BCM", + # At MAXIV we consider session created on fly as scheduled + "scheduled": True, + } + new_session = self.adapter.client.post("sessions", json=payload) + session_to_activate.session_id = str(new_session.get("sessionId")) + self.session_manager.sessions[_idx] = session_to_activate + + self.session_manager.active_session = session_to_activate + + return session_to_activate def get_full_user_name(self) -> str: - person = self.adapter.get_person_by_username(self.user_name) - - given_name = person["givenName"] - family_name = person["familyName"] + person = self.adapter.get_current_user_data() + return "%s %s" % (person["givenName"], person["familyName"]) - return f"{given_name} {family_name}" + def xrf_spectrum_results_url(self, spectrum_id: int) -> str: + return self._rest_client.get_xrf_graph_url(spectrum_id) diff --git a/mxcubecore/HardwareObjects/MAXIV/MAXIVPyISPYyBDataAdapter.py b/mxcubecore/HardwareObjects/MAXIV/MAXIVPyISPYyBDataAdapter.py new file mode 100644 index 0000000000..3f7328fa2f --- /dev/null +++ b/mxcubecore/HardwareObjects/MAXIV/MAXIVPyISPYyBDataAdapter.py @@ -0,0 +1,72 @@ +# ruff: noqa: TD003, FIX002, ERA001 +try: + # duo is a part of sdm package that lives at MaxIV private repository + from duo.UO import RestDuo + from sdm.config import DUOPASSWORD, DUOUSER +except ImportError: + RestDuo = None + DUOPASSWORD = None + DUOUSER = None + +from mxcubecore.HardwareObjects.abstract.PyISPyBDataAdapter import PyISPyBDataAdapter +from mxcubecore.model.lims_session import Proposal, Session + +DUO_API_URL = "https://duo-api.maxiv.lu.se" +LAZY_SESSION_PREFIX = "lazy" + + +class MAXIVPyISPyBDataAdapter(PyISPyBDataAdapter): + """Extend the standard ISPyB data adapter with MAXIV specific logic.""" + + def __init__(self, duo_api_url: str, *args, **kwargs): + super().__init__(*args, **kwargs) + self.duo_api_url = duo_api_url + + def get_proposals(self): + """Override method to filter proposals by the type, state and beamline name. + + Include proposals: of type ``MX`` or ``MB`` in ``Open`` state and assigned + to the current beamline. The last is checked via DUO API. + """ + duo = RestDuo(self.duo_api_url) + duo.login(DUOUSER, DUOPASSWORD) + beamline_proposals_ids = set(duo.get_beamline_proposals(self.beamline_name)) + return [ + self._PyISPyBDataAdapter__to_proposal(proposal) + for proposal in self.client.get("proposals") + if proposal.get("proposalCode").upper() in ["MX", "MB"] + and proposal.get("state", "").capitalize() == "Open" + and int(proposal.get("proposalNumber")) in beamline_proposals_ids + ] + + def create_session(self, proposal: Proposal) -> Session: + """Create a new Session object for the given proposal and beamline. + + This is a lazy session creation, done automatically on the fly in case + no appropriate session is found for the user, proposal and current day. + This session is labelled with ``lazy`` prefix and is not posted to + Py-ISPYB service until it is selected. + + Args: + proposal: Proposal object to create session for + + Returns: + Session: Created session object + """ + return Session( + code=proposal.code, + number=proposal.number, + proposal_name=proposal.code + proposal.number, + proposal_id=proposal.proposal_id, + session_id=f"{LAZY_SESSION_PREFIX}{proposal.proposal_id}", + beamline_name=self.beamline_name, + title=proposal.title, + # At MAXIV we don't care if a session is scheduled, set True as default + is_scheduled_time=True, + is_scheduled_beamline=True, + # TODO@dominikatrojanowska: check if we should set start and end time + # for the session created on fly, and if so, what time should be set. + # For now, we just set empty string, and let ISPyB handle it. + # "startDate": start_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + # "endDate": end_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + ) diff --git a/mxcubecore/HardwareObjects/abstract/PyISPyBDataAdapter.py b/mxcubecore/HardwareObjects/abstract/PyISPyBDataAdapter.py new file mode 100644 index 0000000000..84e79934ed --- /dev/null +++ b/mxcubecore/HardwareObjects/abstract/PyISPyBDataAdapter.py @@ -0,0 +1,576 @@ +import logging +from datetime import datetime, timedelta + +from mxcubecore.HardwareObjects.abstract.PyISPyBRestClient import ( + PyISPyBRestClient, + PyISPyBUnsuccessfulResponse, +) +from mxcubecore.model.lims_session import LimsSessionManager, Proposal, Session + + +class PyISPyBDataAdapter: + """Adapter to convert data from PyISPyB REST API.""" + + def __init__( + self, + client: PyISPyBRestClient, + beamline_name: str, + new_session_duration_days: int = 2, + ): + self.client = client + self.beamline_name = beamline_name + self.new_session_duration_days = new_session_duration_days + self.logger = logging.getLogger("pyispyb_adapter") + + def __is_time_between( + self, start_datetime: datetime, end_datetime: datetime + ) -> bool: + """Checks if the current time is between start and end.""" + today = datetime.today() # noqa: DTZ002 + try: + return start_datetime <= today <= end_datetime + except TypeError: + self.logger.exception( + "Invalid date format. start: %s, end: %s, now: %s", + start_datetime, + end_datetime, + today, + ) + return False + + def __to_proposal(self, proposal: dict[str, str]) -> Proposal: + """ + Converts proposal data received from PyISPyB REST API to a Proposal object. + """ + return Proposal( + code=proposal.get("proposalCode").upper(), + number=proposal.get("proposalNumber"), + proposal_id=str(proposal.get("proposalId")), + title=proposal.get("title"), + type=proposal.get("type", ""), + name=proposal.get("proposal"), + ) + + def __to_session(self, session: dict) -> Session: + """Converts session data received from PyISPyB REST API to a Session object.""" + proposal_name = session.get("proposal") + proposal_code = "".join([c for c in proposal_name if not c.isdigit()]) + proposal_number = proposal_name[len(proposal_code) :] + title = session.get("title", "") + start_datetime = datetime.fromisoformat(session.get("startDate")) + end_datetime = datetime.fromisoformat(session.get("endDate")) + return Session( + code=proposal_code, + number=proposal_number, + proposal_name=proposal_name, + proposal_id=str(session.get("proposalId")), + session_id=str(session.get("sessionId")), + beamline_name=session.get("beamLineName", self.beamline_name), + title=title, + comments=session.get("comments"), + start_datetime=start_datetime, + end_datetime=end_datetime, + nb_shifts=str(session.get("nbShifts", "")), + scheduled=str(session.get("scheduled", "False")), + is_scheduled_time=self.__is_time_between(start_datetime, end_datetime), + is_scheduled_beamline=True, # MAX IV does not care about this value + ) + + @staticmethod + def __to_sample(sample_data: dict) -> dict: + """ + Convert sample data get from PyISPyB to internal + sample data structure. + + Args: + sample_data (dict): The raw sample data from PyISPyB. + + Returns: + dict: A dictionary conteining keys required by internal sample structure. + """ + container = sample_data["Container"] + crystal = sample_data["Crystal"] + return { + # experimentType + "sampleId": sample_data["blSampleId"], + "code": sample_data["code"], + "sampleName": sample_data["name"], + "sampleLocation": sample_data["location"], + "containerSampleChangerLocation": container["sampleChangerLocation"], + "containerCode": container["code"], + "proteinAcronym": crystal["Protein"]["acronym"], + "crystalId": crystal["crystalId"], + "crystalSpaceGroup": crystal["space_group"], + "cellA": crystal["cell_a"], + "cellAlpha": crystal["cell_alpha"], + "cellB": crystal["cell_b"], + "cellBeta": crystal["cell_beta"], + "cellC": crystal["cell_c"], + "cellGamma": crystal["cell_gamma"], + "diffractionPlan": sample_data["DiffractionPlan"]["diffractionPlanId"], + } + + def __find_proposal_by_id(self, proposal_id: int) -> Proposal: + return self.__to_proposal( + self.client.get("proposals?proposalId=%s" % (proposal_id))[0] + ) + + def get_current_user_data(self) -> dict: + """Fetches current user details.""" + return self.client.get("user/current") + + def get_proposals(self) -> list[Proposal]: + """Returns proposals to which authenticated user has access.""" + return [ + self.__to_proposal(proposal) for proposal in self.client.get("proposals") + ] + + def find_proposal(self, code: str, number: str) -> Proposal: + """Finds a proposal by its code and number.""" + return self.__to_proposal( + self.client.get("proposals?search=%s%s" % (code, number))[0] + ) + + def get_sessions_by_code_and_number( + self, code: str, number: str, beamline: str + ) -> LimsSessionManager: + """Finds a session by its proposal code and number and beamline name.""" + return LimsSessionManager( + sessions=[ + self.__to_session(session) + for session in self.client.get( + "sessions?proposal=%s%s&beamLineName=%s" % (code, number, beamline) + ) + ] + ) + + def find_sessions_by_proposal_and_beamline_for_today( + self, code: str, number: str, beamline: str + ) -> list[Session]: + """Finds todays sessions by proposal code, number and beamline name.""" + today = datetime.today() # noqa: DTZ002 + day, month, year = today.day, today.month, today.year + return [ + self.__to_session(session) + for session in self.client.get( + "sessions?proposal=%s%s&beamLineName=%s&year=%s&month=%s&day=%s" + % (code, number, beamline, year, month, day) + ) + if self.__is_time_between( + datetime.fromisoformat(session.get("startDate")), + datetime.fromisoformat(session.get("endDate")), + ) + ] + + def get_sessions_by_username( + self, + username: str = "", + beamline_name: str = "", + ) -> LimsSessionManager: + """Get the list of sessions for the authenticated user and current beamline. + + PyISPyB returns only proposals accessible to the authenticated user. + For each proposal, the method fetches sessions for the current month and + picks one overlapping with the current time. If no such session exists, + a new one is created. + Args: + username: Username to fetch sessions for (left for consistency) + beamline_name: Beamline name to fetch sessions for (left for consistency) + + Returns: + LimsSessionManager: A manager containing the list of sessions + """ + sessions: list[Session] = [] + for proposal in self.get_proposals(): + try: + session = self.find_sessions_by_proposal_and_beamline_for_today( + proposal.code, proposal.number, self.beamline_name + )[0] + except IndexError: + self.logger.info( + "No sessions planned for proposal %s%s. Creating new session.", + proposal.code, + proposal.number, + ) + session = self.create_session(proposal) + sessions.append(session) + return LimsSessionManager(sessions=sessions) + + def create_session(self, proposal: Proposal) -> Session: + """Creates new session via PyISPyB REST API for the given proposal.""" + start_time = datetime.today().replace(hour=0, minute=0, second=0, microsecond=0) # noqa: DTZ002 + end_time = start_time + timedelta( + days=self.new_session_duration_days, hours=7, minutes=59, seconds=59 + ) + payload = { + "proposalId": proposal.proposal_id, + "startDate": start_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + "endDate": end_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + "beamLineName": self.beamline_name, + "comments": "Session created by the BCM", + "scheduled": False, + } + return self.__to_session(self.client.post("sessions", json=payload)) + + def _store_data_collection_group(self, group_data: dict) -> dict: + """Stores data collection group in PyISPyB and returns the group id. + + Args: + group_data: A dictionary containing data collection group information. + + Returns: + Response from PyISPyB: { + "dataCollectionGroupId": 0, + "experimentType": "string", + "blSampleId": 0, + "sessionId": 0, + "workflowId": 0 + } + """ + return self.client.post("datacollections/group", json=group_data) + + def _update_data_collection(self, mx_collection: dict) -> tuple[int, int]: + """Updates data collection in PyISPyB. + + Args: + mx_collection: A dictionary containing data collection information. + Expected to contain "collection_id" key. + + Returns: + Tuple containing data collection id and data collection group id. + """ + dc_id, dc_group_id = 0, 0 + if "collection_id" in mx_collection: + group_id = self._store_data_collection_group(mx_collection)[ + "dataCollectionGroupId" + ] + try: + dc_id, dc_group_id, *_ = self.client.post( + "datacollections/datacollection", + json={**mx_collection, "group_id": group_id}, + ) + except PyISPyBUnsuccessfulResponse: + self.logger.exception("Error in _update_data_collection") + else: + self.logger.error( + "Error in _update_data_collection: collection-id missing, " + "the PyISPyB data-collection is not updated." + ) + return dc_id, dc_group_id + + def store_image(self, image_dict: dict) -> int: + """Stores a data collection image in PyISPyB. + + Args: + image_dict: A dictionary containing image information (paths and properties + of data collection). Expected to contain "dataCollectionId" key. + + Returns: + The id of the stored image. 0 if failed or dataCollectionId is missing. + """ + self.logger.info("Storing image in LIMS") + image_id = 0 + if "dataCollectionId" in image_dict: + try: + response = self.client.post("images/image", json=image_dict) + self.logger.debug("PyISPYB store image response: %s", response) + image_id = response["imageId"] + except PyISPyBUnsuccessfulResponse: + self.logger.exception("Exception in store_image") + else: + self.logger.info("Storing image in LIMS completed, id : %s", image_id) + else: + self.logger.error( + "Error in store_image: data_collection_id missing, " + "could not store image in LIMS" + ) + return image_id + + def get_samples(self, proposal_id: int) -> list[dict]: + """Fetches samples for the given proposal id from PyISPyB.""" + try: + proposal = self.__find_proposal_by_id(proposal_id) + samples = self.client.get( + "samples?proposal=%s%s&beamLineName=%s" + % (proposal.code, proposal.number, self.beamline_name), + timeout=10, + ) + except PyISPyBUnsuccessfulResponse: + self.logger.exception("Error in get_samples") + return [] + return [self.__to_sample(sample) for sample in samples] + + def store_robot_action(self, robot_action: dict) -> int: + """Stores robot action. + + Args: + robot_action: A dictionary containing robot action information. + Expected to contain "sampleId" key. + + Returns: + The id of the robot action. 0 if failed or required keys are missing. + """ + self.logger.info("Storing robot actions in LIMS") + robot_action_id = 0 + try: + robot_action["blSampleId"] = robot_action.pop("sampleId") + except KeyError: + self.logger.exception( + "Error in store_robot_action: missing required keys in robot_action" + "dictionary, could not store robot action in LIMS" + ) + else: + robot_action["startTimestamp"] = robot_action.pop("startTime", "") + robot_action["endTimestamp"] = robot_action.pop("endTime", "") + try: + robot_action_id = self.client.post( + "events/robot-action", json=robot_action + ) + except PyISPyBUnsuccessfulResponse: + self.logger.exception("Exception in store_robot_action") + return robot_action_id + + def associate_bl_sample_and_energy_scan(self, entry_dict: dict) -> int: + """Associates bl sample and energy scan in PyISPyB. + + Args: + entry_dict: A dictionary containing energy scan and bl sample information. + + Returns: + The id of the association. -1 if association failed. + """ + payload = { + "energyScanId": entry_dict["energyScanId"], + "blSampleId": entry_dict["blSampleId"], + } + try: + assoc_id = self.client.patch( + "events/energyscan/associate-bl-sample", json=payload + ) + except Exception: + self.logger.exception( + "Failed to associate bl sample and energy scan in PyISPyB" + ) + assoc_id = -1 + return assoc_id + + def get_data_collection(self, data_collection_id: int) -> dict: + """Fetches data collection details from PyISPyB. + + Args: + data_collection_id: The id of the data collection to fetch. + + Returns: + Dictionary containing data collection details or empty if fetching failed. + """ + try: + dc = self.client.get(f"datacollections/{data_collection_id}") + except Exception: + self.logger.exception("Failed to get data collection from PyISPyB") + return {} + if not dc: + return {} + dc["startTime"] = datetime.fromisoformat(dc["startTime"]).strftime( + "%Y-%m-%d %H:%M:%S" + ) + dc["endTime"] = datetime.fromisoformat(dc["endTime"]).strftime( + "%Y-%m-%d %H:%M:%S" + ) + return dc + + def find_detector( + self, + manufacturer, + model, + mode, + type="", # noqa: A002 + ) -> dict | None: + """Finds detector by its type, manufacturer, model and mode. + + Gives the first matching detector or None if no match is found. + Args: + manufacturer: The manufacturer of the detector. + model: The model of the detector. + mode: The mode of the detector. + type: The type of the detector (optional). + + Returns: + A dictionary containing detector details. + """ + try: + return self.client.get( + "detectors?manufacturer=%s&model=%s&mode=%s&type=%s" + % ( + manufacturer, + model, + mode, + type if type else "", + ) + ) + except PyISPyBUnsuccessfulResponse: + self.logger.exception("`Exception in find_detector") + return None + + def update_session(self, session: dict) -> dict: + """Updates Beamline Setup of existing session. + + Returns dictionary containing updated session details from PyISPyB. + If updating failed, returns empty dictionary. + """ + try: + session_id = session["sessionId"] + beamline_setup = session.get("BeamLineSetup") or {} + beamline_setup_id = beamline_setup.get("beamLineSetupId") + if not beamline_setup_id: + msg = "Missing beamLineSetupId" + raise KeyError(msg) + return self.client.patch( + "sessions/%s/associate-beamline-setup?beamLineSetupId=%s" + % ( + session_id, + beamline_setup_id, + ), + ) + except (PyISPyBUnsuccessfulResponse, KeyError): + self.logger.exception("Failed to store or update session") + return {} + + def get_session(self, session_id: int) -> dict: + """Fetches session details from PyISPyB by session id.""" + try: + session = self.client.get("sessions/%s" % session_id) + except PyISPyBUnsuccessfulResponse: + self.logger.exception("Failed to get session from PyISPyB") + session = {} + return session + + def store_beamline_setup(self, session_id: int, bl_config: dict) -> int | None: + """ + Stores beamline setup in PyISPyB and associates it with the session. + + Args: + session_id: The id of the session to associate the beamline setup with. + bl_config: A dictionary containing beamline setup information. + Returns: + The id of the stored beamline setup. None if failed.""" + beamline_setup_id = None + session = self.get_session(session_id) + if session: + try: + beamline_setup_id = self.client.post( + "beamline-setups/beamline-setup", json=bl_config + ) + except PyISPyBUnsuccessfulResponse: + self.logger.exception( + "Failed to store or update beamline setup in PyISPyB" + ) + bl_config["beamlineSetupId"] = beamline_setup_id + session["BeamLineSetup"] = bl_config + self.update_session(session) + return beamline_setup_id + + def store_data_collection(self, mx_collection, bl_config=None): + self.logger.info("Storing datacollection in PyISPyB") + return self._store_data_collection(mx_collection, bl_config) + + def update_data_collection(self, mx_collection): + self.logger.info("Updating datacollection in PyISPyB") + return self._update_data_collection(mx_collection) + + def finalize_data_collection(self, mx_collection): + return self.update_data_collection(mx_collection) + + def _store_data_collection(self, mx_collection: dict, bl_config: dict | None): + """Stores data collection group in PyISPyB and returns the group id. + + Args: + mx_collection: A dictionary containing data collection information. + bl_config: A dictionary containing beamline setup information (optional). + + Returns: + Store data collection id and detector id if found, otherwise 0. + """ + detector_id = 0 + if bl_config: + bl_config["synchrotronMode"] = bl_config.get( + "synchrotronMode", mx_collection.get("synchrotronMode", "") + ) + self.store_beamline_setup(mx_collection["sessionId"], bl_config) + detector = self.find_detector( + bl_config.get("detector_manufacturer", ""), + bl_config.get("detector_model", ""), + bl_config.get("detector_binning_mode", ""), + bl_config.get("detector_type", ""), + ) + if detector: + detector_id = detector.get("detectorId", 0) + + response = self.client.post( + "datacollections/datacollection", + json={**mx_collection, "detectorId": detector_id} + if detector_id + else mx_collection, + ) + collection_id = response["dataCollectionId"] + return collection_id, detector_id + + def store_energy_scan(self, energyscan: dict) -> dict[str, int]: + """Stores energy scan in PyISPyB + + Args: + energyscan: A dictionary containing energy scan information. + + Returns: + Dictionary {"energyScanId": number}, where number is the id of the stored + energy scan. + """ + try: + return { + "energyScanId": self.client.post("events/energyscan", json=energyscan) + } + + except PyISPyBUnsuccessfulResponse: + self.logger.exception("Failed to store energy scan in PyISPyB") + return {"energyScanId": -1} + + def store_xfe_spectrum(self, xfe_spectrum: dict) -> dict[str, int]: + """Stores XFE fluorescence spectrum in PyISPyB + + Args: + xfe_spectrum: A dictionary containing XFE fluorescence spectrum information. + + Returns: + Dictionary {"xfeFluorescenceSpectrumId": number}, where number is the id + of the stored XFE fluorescence spectrum. + """ + try: + return { + "xfeFluorescenceSpectrumId": self.client.post( + "events/xfe-fluorescence-spectrum", json=xfe_spectrum + ) + } + except PyISPyBUnsuccessfulResponse: + self.logger.exception( + "Failed to store XFE fluorescence spectrum in PyISPyB" + ) + return {"xfeFluorescenceSpectrumId": -1} + + def update_bl_sample(self, bl_sample: dict) -> dict: + """Updates beamline sample in PyISPyB. + + Args: + bl_sample: A dictionary containing beamline sample information. + Expected to contain "blSampleId" key. + Returns: + Response from PyISPyB - a dictionary containing updated beamline sample + details or empty if updating failed. + """ + try: + bl_sample_id = bl_sample.get("blSampleId") + if bl_sample_id is None: + self.logger.error("Missing blSampleId") + return {} + return self.client.patch("samples/%s" % bl_sample_id, json=bl_sample) + except PyISPyBUnsuccessfulResponse: + self.logger.exception("Failed to update beamline sample in PyISPyB") + return {} diff --git a/mxcubecore/HardwareObjects/abstract/PyISPyBRestClient.py b/mxcubecore/HardwareObjects/abstract/PyISPyBRestClient.py new file mode 100644 index 0000000000..80a1e30998 --- /dev/null +++ b/mxcubecore/HardwareObjects/abstract/PyISPyBRestClient.py @@ -0,0 +1,154 @@ +import logging +from datetime import datetime, timedelta, timezone +from json.decoder import JSONDecodeError +from urllib.parse import urljoin + +from requests import Session + +log = logging.getLogger("py-ispyb_client") + + +class NoTokenException(Exception): + """Exception raised when no token is returned from authentication.""" + + +class PyISPyBUnsuccessfulResponse(Exception): + """Exception raised when a response from the server is unsuccessful (not 200).""" + + +class AuthenticationExpired(Exception): + """Raised when authentication refresh failed.""" + + +class PyISPyBRestClient: + """REST client for PyISPyB. + + It handles authentication and communication with PyISPyB REST API. + """ + + REFRESH_MARGIN_SECONDS = 60 + + def __init__(self, rest_root: str, timeout: int = 5): + self._rest_root = rest_root + self._session = Session() + self._timeout = timeout + self._username = None + self._access_token = None + self._refresh_token = None + self._token_expiry = None + + def authenticate(self, user_name: str, token: str): + self._username = user_name + response = self.post( + "auth/login", + json={ + "plugin": "keycloak", + "login": user_name, + "token": token, + }, + skip_refresh=True, + ) + self._store_tokens(response) + + def post(self, endpoint, **kwargs): + return self._request(self._session.post, endpoint, **kwargs) + + def get(self, endpoint, **kwargs): + return self._request(self._session.get, endpoint, **kwargs) + + def patch(self, endpoint, **kwargs): + return self._request(self._session.patch, endpoint, **kwargs) + + def update_proxies(self, proxy: dict): + self._session.proxies.update(proxy) + + def _request(self, method, endpoint, skip_refresh=False, **kwargs): # noqa: FBT002 + if not skip_refresh and self._is_token_expired(): + self._refresh_access_token() + timeout = kwargs.pop("timeout", self._timeout) + url = urljoin(self._rest_root, endpoint) + response = method(url, timeout=timeout, **kwargs) + if response.status_code == 401 and not skip_refresh: + log.warning("Received 401. Attempting token refresh.") + self._refresh_access_token() + response = method(url, timeout=timeout, **kwargs) + return self.decode_json_response(response) + + def _refresh_access_token(self): + if not self._refresh_token: + msg = "No refresh token available" + raise AuthenticationExpired(msg) + log.info("Refreshing keycloak access token") + url = urljoin(self._rest_root, "auth/refresh") + response = self._session.post( + url, + timeout=self._timeout, + json={"refreshToken": self._refresh_token}, + ) + if response.status_code != 200: + msg = "Failed to refresh token" + raise AuthenticationExpired(msg) + try: + self._store_tokens(response.json()) + except JSONDecodeError as ex: + msg = "Authentication response malformed." + log.exception(msg) + raise AuthenticationExpired(msg) from ex + + @staticmethod + def decode_json_response(response): + log.debug( + "Received response from %s. Status code: %s, Response text: %s", + response.url, + response.status_code, + response.text, + ) + if response.status_code not in (200, 201): + msg = ( + f"Request to {response.url} failed with code: {response.status_code}. " + f"Response: {response.text}" + ) + raise PyISPyBUnsuccessfulResponse(msg) + try: + response_json = response.json() + except JSONDecodeError: + log.exception( + "Failed to decode JSON response from %s. " + "Status code: %s, Response text: %s", + response.url, + response.status_code, + response.text, + ) + raise + if "results" in response_json: + return response_json["results"] + return response_json + + def _store_tokens(self, response: dict): + if not isinstance(response, dict): + msg = ( + "Authentication response malformed: expected dict, " + f"got {type(response)}" + ) + raise NoTokenException(msg) + access_token = response.get("token") + refresh_token = response.get("refreshToken") + expires_in = response.get("expiresIn", 300) + if not access_token: + msg = "Authentication failed. No access token received." + raise NoTokenException(msg) + self._access_token = access_token + self._refresh_token = refresh_token + self._token_expiry = datetime.now(timezone.utc) + timedelta(seconds=expires_in) + self._set_authorization_header(access_token) + + def _is_token_expired(self) -> bool: + if self._token_expiry is None: + return True + + return datetime.now(timezone.utc) >= ( + self._token_expiry - timedelta(seconds=self.REFRESH_MARGIN_SECONDS) + ) + + def _set_authorization_header(self, token: str): + self._session.headers.update({"Authorization": f"Bearer {token}"}) diff --git a/test/factories/__init__.py b/test/factories/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/factories/pyispyb_factory.py b/test/factories/pyispyb_factory.py new file mode 100644 index 0000000000..818e64276d --- /dev/null +++ b/test/factories/pyispyb_factory.py @@ -0,0 +1,191 @@ +from datetime import datetime, timedelta + +from mxcubecore.model.lims_session import Proposal, Session + + +def current_user_response(): + return { + "givenName": "Admin", + "familyName": "Test_User", + "login": "testusr", + "Permissions": [ + "all_proposals", + "all_sessions", + "own_proposals", + "own_sessions", + ], + "personId": 123456, + "beamLineGroups": [], + "beamLines": [], + } + + +def proposal_response( + proposal_id: 10, + proposal_code: str = "MX", + proposal_number: str = "20090662", +): + return { + "proposalCode": proposal_code, + "proposalNumber": proposal_number, + "proposal": f"{proposal_code.lower()}{proposal_number}", + "title": "Test Proposal", + "state": "open", + "proposalType": "MX", + "proposalId": proposal_id, + } + + +def proposals_response(): + return [ + proposal_response( + proposal_id=10, + proposal_code="mx", + proposal_number="20090662", + ), + proposal_response( + proposal_id=11, + proposal_code="mx", + proposal_number="20210662", + ), + ] + + +def session_response( + session_id=1, + proposal="mx20090662", + proposal_id=10, +): + start_date = datetime.today() + end_date = start_date + timedelta(1) + return { + "proposalId": proposal_id, + "startDate": start_date.isoformat(), + "endDate": end_date.isoformat(), + "beamLineName": "PROXIMA1", + "sessionId": session_id, + "proposal": proposal, + "BeamLineSetup": { + "synchrotronMode": "string", + "beamLineSetupId": 1, + }, + } + + +def default_proposal(): + return Proposal( + proposal_id="10", + person_id="123456", + code="mx", + number="20090662", + title="test proposal", + ) + + +def default_session(): + return Session( + session_id="1", + proposal_id="10", + proposal_name="mx20090662", + code="mx", + number="20090662", + ) + + +def default_image_dict(): + return { + "dataCollectionId": 2, + "imageNumber": 0, + "fileName": "test.cbf", + "fileLocation": "/data/test", + } + + +def sample_response(): + return [ + { + "name": "sample-01", + "code": "smpl-01", + "blSampleId": 1, + "location": 1, + "Container": { + "code": "CONTAINER-01", + "sampleChangerLocation": "1", + }, + "Crystal": { + "Protein": { + "acronym": "TEST", + }, + "crystalId": 12345, + "space_group": None, + "cell_a": 0.0, + "cell_alpha": 0.0, + "cell_b": 0.0, + "cell_beta": 0.0, + "cell_c": 0.0, + "cell_gamma": 0.0, + }, + "DiffractionPlan": { + "diffractionPlanId": 23456, + }, + } + ] + + +def default_data_collection(): + return { + "dataCollectionId": 1, + "dataCollectionGroupId": 99, + "strategySubWedgeOrigId": None, + "detectorId": None, + "blSubSampleId": None, + "startPositionId": 9, + "endPositionId": None, + "dataCollectionNumber": 1, + "startTime": "2015-01-20T16:17:13", + "endTime": "2015-01-20T16:17:13", + "runStatus": "failed", + "sessionId": 123, + } + + +def default_detector(): + return { + "detectorId": 1, + "detectorType": "PixelDetector", + "detectorManufacturer": "ExampleManufacturer", + "detectorModel": "ExampleModel", + "detectorPixelSizeHorizontal": 0.075, + } + + +def default_bl_config(): + return { + "detector_manufacturer": "ExampleManufacturer", + "detector_model": "ExampleModel", + "detector_binning_mode": "1x1", + "detector_type": "PixelDetector", + } + + +def default_energy_scan(): + return { + "fluorescenceDetector": "Pilatus", + "scanFileFullPath": "/data/scan1.dat", + "choochFileFullPath": "/data/chooch1.dat", + "jpegChoochFileFullPath": "/data/chooch1.jpg", + } + + +def default_xfe_spectrum(): + return { + "sessionId": 44692, + "blSampleId": 299, + "fittedDataFileFullPath": "string", + "scanFileFullPath": "string", + "jpegScanFileFullPath": "string", + } + + +def default_bl_sample(): + return {"blSampleId": 1115, "name": "abcf-28", "code": ""} diff --git a/test/test_hwo_maxiv_ispyblims.py b/test/test_hwo_maxiv_ispyblims.py index 6942643cf7..daab1b1931 100644 --- a/test/test_hwo_maxiv_ispyblims.py +++ b/test/test_hwo_maxiv_ispyblims.py @@ -6,23 +6,19 @@ import pytest -ISPYB_AUTH_ERROR_MESSSAGE = ( - "JBAS011843: Failed instantiate InitialContextFactory com.sun.jndi.ldap.LdapCtxFactory " - 'from classloader ModuleClassLoader for Module "deployment.ispyb.ear.ispyb-ws.war:main" ' - "from Service Module Loader" -) +PYISPYB_AUTH_ERROR_MESSSAGE = "Could not verify" REST_ROOT = "http://example.com/rest/" USER = "testusr" -PASS = "testpsd" # noqa: S105 +TOKEN = "testtkn" # noqa: S105 @pytest.fixture @patch("mxcubecore.HardwareObjects.abstract.AbstractLims.HWR") @patch("mxcubecore.HardwareObjects.abstract.ISPyBAbstractLims.HWR") -@patch("mxcubecore.HardwareObjects.abstract.ISPyBDataAdapter.Client") +@patch("mxcubecore.HardwareObjects.abstract.PyISPyBRestClient.Session") @patch("mxcubecore.HardwareObjects.UserTypeISPyBLims.Client") -def ispyb_lims(_suds_mock1, _suds_mock2, _hwr_mock, hwr_mock): +def ispyb_lims(_suds_mock, session_mock, _hwr_mock, hwr_mock): # make sure mocked session HWOBJ have session set hwr_mock.beamline.session.beamline_name = "ID42" @@ -31,12 +27,11 @@ def ispyb_lims(_suds_mock1, _suds_mock2, _hwr_mock, hwr_mock): lims = ISPyBLims(name="dummy") lims._config = lims.HOConfig( # noqa: SLF001 - ws_root="http://example.com/tstWS/", - rest_root=REST_ROOT, + duo_api_url="http://example.com/rest/", pyispyb_rest_root=REST_ROOT ) lims.init() - return lims + return lims, session_mock.return_value.post @pytest.fixture @@ -47,14 +42,21 @@ def json_decode_error(): return JSONDecodeError("mocked-err", doc, 0) -@patch("requests.post") -def test_login_ok(post_mock, ispyb_lims): +# @patch("requests.Session") +def test_login_ok(ispyb_lims): """test the case when user logs in successfully""" + ispyb_lims, post_mock = ispyb_lims - # mocks the ISPyB client to return valid token on POST request - post_mock.return_value.json.return_value = {"token": "dummy-token"} + # mock the ISPyB client to return valid token on POST request + response_mock = MagicMock( + url=f"{REST_ROOT}auth/login", + status_code=200, + text="dummy-text", + ) + post_mock.return_value = response_mock + response_mock.json.return_value = {"token": "dummy-token"} - is_ok, err = ispyb_lims.ispyb_login(USER, PASS) + is_ok, err = ispyb_lims.ispyb_login(USER, TOKEN) # check that we got 'login ok' result assert is_ok @@ -62,27 +64,44 @@ def test_login_ok(post_mock, ispyb_lims): # check that correct POST request was made post_mock.assert_called_once_with( - f"{REST_ROOT}authenticate?site=MAXIV", - data={"login": USER, "password": PASS}, + f"{REST_ROOT}auth/login", + json={ + "plugin": "keycloak", + "login": USER, + "token": TOKEN, + }, + timeout=5, ) -@patch("requests.post") -def test_login_invalid_credentials(post_mock, ispyb_lims, json_decode_error): +def test_login_invalid_credentials(ispyb_lims, json_decode_error): """test the case when user fails to login due to wrong credentials""" + ispyb_lims, post_mock = ispyb_lims - # # mock ISPyB 'login failed' response, - # yes, this is how ISPyB replies at MAXIV - # - response_mock = MagicMock() - response_mock.text = ISPYB_AUTH_ERROR_MESSSAGE - response_mock.json.side_effect = json_decode_error - + response_mock = MagicMock( + url=f"{REST_ROOT}auth/login", + status_code=401, + text=PYISPYB_AUTH_ERROR_MESSSAGE, + ) post_mock.return_value = response_mock + response_mock.json.side_effect = json_decode_error - is_ok, err = ispyb_lims.ispyb_login(USER, PASS) + is_ok, err = ispyb_lims.ispyb_login(USER, TOKEN) # check that we got 'login failed' result assert not is_ok - assert err == "invalid credentials" + assert ( + err == f"Request to {REST_ROOT}auth/login failed with code: 401. " + f"Response: {PYISPYB_AUTH_ERROR_MESSSAGE}" + ) + + post_mock.assert_called_once_with( + f"{REST_ROOT}auth/login", + json={ + "plugin": "keycloak", + "login": USER, + "token": TOKEN, + }, + timeout=5, + ) diff --git a/test/test_pyispyb_data_adapter.py b/test/test_pyispyb_data_adapter.py new file mode 100644 index 0000000000..7c637820bd --- /dev/null +++ b/test/test_pyispyb_data_adapter.py @@ -0,0 +1,816 @@ +from datetime import datetime, timedelta +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from mxcubecore.HardwareObjects.abstract.PyISPyBDataAdapter import PyISPyBDataAdapter +from mxcubecore.HardwareObjects.abstract.PyISPyBRestClient import ( + PyISPyBRestClient, + PyISPyBUnsuccessfulResponse, +) +from mxcubecore.model.lims_session import Session +from test.factories.pyispyb_factory import ( + current_user_response, + default_bl_config, + default_bl_sample, + default_data_collection, + default_detector, + default_energy_scan, + default_image_dict, + default_proposal, + default_session, + default_xfe_spectrum, + proposals_response, + sample_response, + session_response, +) + + +@pytest.fixture +def client(): + return Mock(spec=PyISPyBRestClient) + + +@pytest.fixture +def adapter(client): + adapter = PyISPyBDataAdapter(client=client, beamline_name="PROXIMA1") + adapter.logger = Mock() + return adapter + + +def test_get_current_user_data(adapter, client): + client.get.return_value = current_user_response() + + currrent_user = adapter.get_current_user_data() + + client.get.assert_called_once_with("user/current") + assert currrent_user["login"] == "testusr" + assert currrent_user["givenName"] == "Admin" + assert currrent_user["personId"] == 123456 + + +def test_get_proposals(adapter, client): + client.get.return_value = proposals_response() + + proposals = adapter.get_proposals() + + client.get.assert_called_once_with("proposals") + assert len(proposals) == 2 + assert proposals[0].proposal_id == "10" + assert proposals[1].proposal_id == "11" + + +def test_find_proposal(adapter, client): + client.get.return_value = proposals_response() + + proposal = adapter.find_proposal("mx", "20090662") + + client.get.assert_called_once_with("proposals?search=mx20090662") + assert proposal.proposal_id == "10" + assert proposal.code == "MX" + assert proposal.number == "20090662" + + +def test_get_sessions_by_code_and_number( + adapter, + client, +): + client.get.return_value = [ + session_response( + session_id=1, + proposal="mx20090662", + proposal_id=10, + ), + session_response( + session_id=2, + proposal="mx20090662", + proposal_id=10, + ), + ] + + result = adapter.get_sessions_by_code_and_number( + code="mx", + number="20090662", + beamline="PROXIMA1", + ) + + client.get.assert_called_once_with( + "sessions?proposal=mx20090662&beamLineName=PROXIMA1" + ) + assert len(result.sessions) == 2 + session_1 = result.sessions[0] + assert session_1.session_id == "1" + assert session_1.proposal_id == "10" + assert session_1.proposal_name == "mx20090662" + session_2 = result.sessions[1] + assert session_2.session_id == "2" + assert session_2.proposal_name == "mx20090662" + + +def test_find_sessions_by_proposal_and_beamline_for_today( + adapter, + client, +): + client.get.return_value = [ + session_response(session_id=1), + session_response(session_id=2), + ] + frozen_now = datetime(2026, 6, 2, 14, 30, 45) + + with ( + patch( + "mxcubecore.HardwareObjects.abstract.PyISPyBDataAdapter.datetime" + ) as mock_datetime, + patch.object( + adapter, + "_PyISPyBDataAdapter__is_time_between", + side_effect=[True, True, False], + ), + ): + mock_datetime.today.return_value = frozen_now + + result = adapter.find_sessions_by_proposal_and_beamline_for_today( + code="mx", + number="20090662", + beamline="PROXIMA1", + ) + + client.get.assert_called_once_with( + f"sessions?proposal=mx20090662" + f"&beamLineName=PROXIMA1" + f"&year={frozen_now.year}" + f"&month={frozen_now.month}" + f"&day={frozen_now.day}" + ) + assert len(result) == 1 + session = result[0] + assert session.session_id == "1" + assert session.proposal_id == "10" + assert session.proposal_name == "mx20090662" + assert session.code == "mx" + assert session.number == "20090662" + + +def test_get_sessions_by_username_when_sessions_exist( + adapter, +): + proposal = MagicMock(code="mx", number="20090662") + session = MagicMock(spec=Session) + + with ( + patch.object(adapter, "get_proposals", return_value=[proposal]), + patch.object( + adapter, + "find_sessions_by_proposal_and_beamline_for_today", + return_value=[session], + ) as find_sessions_mock, + patch.object( + adapter, + "create_session", + ) as create_session_mock, + ): + result = adapter.get_sessions_by_username() + + assert len(result.sessions) == 1 + assert result.sessions[0] is session + find_sessions_mock.assert_called_once_with( + "mx", + "20090662", + adapter.beamline_name, + ) + create_session_mock.assert_not_called() + + +def test_get_sessions_by_username_creates_session_when_none_exists( + adapter, +): + proposal = default_proposal() + session = default_session() + + with ( + patch.object(adapter, "get_proposals", return_value=[proposal]), + patch.object( + adapter, "find_sessions_by_proposal_and_beamline_for_today", return_value=[] + ), + patch.object( + adapter, "create_session", return_value=session + ) as create_session_mock, + ): + result = adapter.get_sessions_by_username() + + create_session_mock.assert_called_once_with(proposal) + assert len(result.sessions) == 1 + assert result.sessions[0] is session + + +def test_creat_session( + adapter, + client, +): + proposal = default_proposal() + frozen_now = datetime(2026, 6, 2, 14, 30, 45) + + with patch( + "mxcubecore.HardwareObjects.abstract.PyISPyBDataAdapter.datetime" + ) as mock_datetime: + mock_datetime.today.return_value = frozen_now + client.post.return_value = session_response() + + adapter.create_session(proposal) + + client.post.assert_called_once() + + start_time = frozen_now.replace(hour=0, minute=0, second=0, microsecond=0) + end_time = start_time + timedelta( + days=adapter.new_session_duration_days, hours=7, minutes=59, seconds=59 + ) + (endpoint,) = client.post.call_args.args + payload = client.post.call_args.kwargs["json"] + + assert endpoint == "sessions" + assert payload["proposalId"] == "10" + assert payload["beamLineName"] == adapter.beamline_name + assert payload["startDate"] == start_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ") + assert payload["endDate"] == end_time.strftime("%Y-%m-%dT%H:%M:%S.%fZ") + assert payload["scheduled"] is False + + +def test_store_image_success( + adapter, + client, +): + image_dict = default_image_dict() + client.post.return_value = { + "imageId": 123, + "dataCollectionId": 2, + } + + response = adapter.store_image(image_dict=image_dict) + + client.post.assert_called_once_with( + "images/image", + json=image_dict, + ) + assert response == 123 + + +def test_store_image_returns_zero_when_post_fails( + adapter, + client, +): + image_dict = default_image_dict() + client.post.side_effect = PyISPyBUnsuccessfulResponse("""Exception raised when a + response from the server is unsuccessful (not 200).""") + + image_id = adapter.store_image(image_dict) + + client.post.assert_called_once_with( + "images/image", + json=image_dict, + ) + assert image_id == 0 + + +def test_store_image_returns_zero_when_data_collection_id_missing( + adapter, + client, +): + image_dict = { + "imageNumber": 0, + "fileName": "test.cbf", + } + + image_id = adapter.store_image(image_dict) + + client.post.assert_not_called() + assert image_id == 0 + + +def test_get_samples( + adapter, + client, +): + proposal = default_proposal() + samples = sample_response() + + with ( + patch.object( + adapter, + "_PyISPyBDataAdapter__find_proposal_by_id", + return_value=proposal, + ) as find_proposal_mock, + patch.object( + client, + "get", + return_value=samples, + ) as get_mock, + ): + result_0 = adapter.get_samples(proposal.proposal_id)[0] + + find_proposal_mock.assert_called_once_with(proposal.proposal_id) + get_mock.assert_called_once_with( + "samples?proposal=mx20090662&beamLineName=PROXIMA1", + timeout=10, + ) + sample_0 = samples[0] + assert result_0["sampleId"] == sample_0["blSampleId"] + assert result_0["cellA"] == sample_0["Crystal"]["cell_a"] + assert ( + result_0["diffractionPlan"] == sample_0["DiffractionPlan"]["diffractionPlanId"] + ) + + +def test_get_samples_returns_empty_list_on_error( + adapter, +): + with patch.object( + adapter, + "_PyISPyBDataAdapter__find_proposal_by_id", + side_effect=PyISPyBUnsuccessfulResponse("boom"), + ): + result = adapter.get_samples(10) + + assert result == [] + + +def test_store_robot_action_missing_sample_id(adapter, client): + robot_action = { + "startTime": "2026-03-16T12:59:07.737Z", + "endTime": "2026-04-16T14:00:07.737Z", + } + + robot_action_id = adapter.store_robot_action(robot_action) + + assert robot_action_id == 0 + client.post.assert_not_called() + + +def test_store_robot_action_post_failure(adapter, client): + client.post.side_effect = PyISPyBUnsuccessfulResponse( + "Failed to store robot action" + ) + adapter.logger.exception = Mock() + robot_action = { + "sampleId": 12, + "startTime": "2026-03-16T12:59:07.737Z", + "endTime": "2026-04-16T14:00:07.737Z", + } + robot_action_id = adapter.store_robot_action(robot_action) + + assert robot_action_id == 0 + + client.post.assert_called_once_with( + "events/robot-action", + json={ + "blSampleId": 12, + "startTimestamp": "2026-03-16T12:59:07.737Z", + "endTimestamp": "2026-04-16T14:00:07.737Z", + }, + ) + adapter.logger.exception.assert_called_once_with("Exception in store_robot_action") + + +def test_store_robot_action_sucess( + adapter, + client, +): + client.post.return_value = 77 + robot_action = { + "sampleId": 12, + "startTime": "2025-01-01", + "endTime": "2025-01-02", + } + + robot_action_id = adapter.store_robot_action(robot_action) + + assert robot_action_id == 77 + client.post.assert_called_once_with( + "events/robot-action", + json={ + "blSampleId": 12, + "startTimestamp": "2025-01-01", + "endTimestamp": "2025-01-02", + }, + ) + + +def test_associate_bl_sample_and_energy_scan( + adapter, + client, +): + entry_dict = {"energyScanId": 125, "blSampleId": 1} + client.patch.return_value = entry_dict + + result = adapter.associate_bl_sample_and_energy_scan(entry_dict) + + assert result["energyScanId"] == 125 + client.patch.assert_called_once_with( + "events/energyscan/associate-bl-sample", + json={ + "energyScanId": 125, + "blSampleId": 1, + }, + ) + + +def test_associate_bl_sample_and_energy_scan_failure( + adapter, + client, +): + client.patch.side_effect = Exception("Patch failed") + entry_dict = {"energyScanId": 125, "blSampleId": 1} + + result = adapter.associate_bl_sample_and_energy_scan(entry_dict) + + assert result == -1 + client.patch.assert_called_once_with( + "events/energyscan/associate-bl-sample", + json={ + "energyScanId": 125, + "blSampleId": 1, + }, + ) + adapter.logger.exception.assert_called_once_with( + "Failed to associate bl sample and energy scan in PyISPyB" + ) + + +def test_get_data_collection( + adapter, + client, +): + client.get.return_value = default_data_collection() + + result = adapter.get_data_collection(1) + + client.get.assert_called_once_with("datacollections/1") + assert result == { + "dataCollectionId": 1, + "dataCollectionGroupId": 99, + "strategySubWedgeOrigId": None, + "detectorId": None, + "blSubSampleId": None, + "startPositionId": 9, + "endPositionId": None, + "dataCollectionNumber": 1, + "startTime": "2015-01-20 16:17:13", + "endTime": "2015-01-20 16:17:13", + "runStatus": "failed", + "sessionId": 123, + } + + +def test_get_data_collection_exception( + adapter, + client, +): + client.get.side_effect = Exception("Request failed") + + result = adapter.get_data_collection(1) + + assert result == {} + client.get.assert_called_once_with("datacollections/1") + adapter.logger.exception.assert_called_once_with( + "Failed to get data collection from PyISPyB" + ) + + +def test_get_data_collection_empty_response( + adapter, + client, +): + client.get.return_value = [] + + result = adapter.get_data_collection(999999) + + assert result == {} + client.get.assert_called_once_with("datacollections/999999") + + +def test_get_data_collection_with_string_dates( + adapter, + client, +): + client.get.return_value = { + "startTime": "2015-01-20T16:17:13", + "endTime": "2015-01-20T16:17:13", + } + + result = adapter.get_data_collection(1) + + assert result == { + "startTime": "2015-01-20 16:17:13", + "endTime": "2015-01-20 16:17:13", + } + + +def test_find_detector( + adapter, + client, +): + client.get.return_value = default_detector() + + result = adapter.find_detector( + manufacturer="ExampleManufacturer", + model="ExampleModel", + mode="Standard", + type="PixelDetector", + ) + + client.get.assert_called_once_with( + "detectors?manufacturer=ExampleManufacturer" + "&model=ExampleModel" + "&mode=Standard" + "&type=PixelDetector" + ) + assert result["detectorId"] == 1 + + +def test_find_detector_without_type( + adapter, + client, +): + detector = {"detectorId": 1} + client.get.return_value = detector + + result = adapter.find_detector( + manufacturer="ExampleManufacturer", + model="ExampleModel", + mode="Standard", + ) + + client.get.assert_called_once_with( + "detectors?manufacturer=ExampleManufacturer" + "&model=ExampleModel" + "&mode=Standard" + "&type=" + ) + assert result == detector + + +def test_update_session_success( + adapter, + client, +): + session = { + "sessionId": 46369, + "BeamLineSetup": { + "beamLineSetupId": 2, + }, + } + client.patch.return_value = {"sessionId": 46369, "beamLineSetupId": 2} + + result = adapter.update_session(session) + + assert result["sessionId"] == 46369 + client.patch.assert_called_once_with( + "sessions/46369/associate-beamline-setup?beamLineSetupId=2", + ) + + +def test_update_session_null_beamline_setup( + adapter, + client, +): + session = { + "sessionId": 46369, + "BeamLineSetup": None, + } + + result = adapter.update_session(session) + + assert result == {} + client.post.assert_not_called() + adapter.logger.exception.assert_called_once_with( + "Failed to store or update session" + ) + + +def test_get_sesssion( + adapter, + client, +): + client.get.return_value = session_response() + + result = adapter.get_session(1) + + assert result["sessionId"] == 1 + client.get.assert_called_once_with("sessions/1") + + +def test_store_beamline_setup_success( + adapter, + client, +): + session_id = 1 + session = { + "sessionId": 1, + "BeamLineSetup": { + "beamLineSetupId": 2, + }, + } + client.get.return_value = session + bl_config = default_bl_config() + client.post.return_value = 101 + adapter.update_session = Mock() + + result = adapter.store_beamline_setup(session_id, bl_config) + + assert result == 101 + client.post.assert_called_once_with( + "beamline-setups/beamline-setup", + json=bl_config, + ) + adapter.update_session.assert_called_once_with(session) + + +def test_store_data_collection_without_bl_config(adapter, client): + mx_collection = default_data_collection() + client.post.return_value = {"dataCollectionId": 10, "dataCollectionGroupId": 99} + + result = adapter.store_data_collection(mx_collection) + + assert result == (10, 0) + client.post.assert_called_once_with( + "datacollections/datacollection", + json=mx_collection, + ) + + +def test_store_data_collection_with_detector(adapter, client): + mx_collection = default_data_collection() + bl_config = default_bl_config() + client.post.return_value = {"dataCollectionId": 10, "dataCollectionGroupId": 99} + adapter.store_beamline_setup = Mock() + adapter.find_detector = Mock(return_value={"detectorId": 42}) + + result = adapter.store_data_collection(mx_collection, bl_config) + + assert result == (10, 42) + adapter.store_beamline_setup.assert_called_once_with( + 123, + bl_config, + ) + adapter.find_detector.assert_called_once_with( + "ExampleManufacturer", + "ExampleModel", + "1x1", + "PixelDetector", + ) + client.post.assert_called_once_with( + "datacollections/datacollection", + json={**mx_collection, "detectorId": 42}, + ) + assert ( + mx_collection["detectorId"] is None + ) # ensures input data was not modified within store_data_collection method + + +def test_store_data_collection_without_detector(adapter, client): + mx_collection = default_data_collection() + bl_config = default_bl_config() + client.post.return_value = {"dataCollectionId": 10, "dataCollectionGroupId": 99} + adapter.store_beamline_setup = Mock() + adapter.find_detector = Mock(return_value=None) + + result = adapter.store_data_collection(mx_collection, bl_config) + + assert result == (10, 0) + client.post.assert_called_once_with( + "datacollections/datacollection", json=mx_collection + ) + assert mx_collection["detectorId"] is None + + +def test_update_data_collection_missing_collection_id(adapter, client): + mx_collection = { + "other_key": 1, + } + + result = adapter._update_data_collection(mx_collection) + + assert result == (0, 0) + client.post.assert_not_called() + + +def test_update_data_collection_success( + adapter, + client, +): + mx_collection = { + "collection_id": 1, + } + + client.post.return_value = (1, 99) + adapter._store_data_collection_group = Mock( + return_value={"dataCollectionGroupId": 99} + ) + + result = adapter._update_data_collection(mx_collection) + + assert result == (1, 99) + adapter._store_data_collection_group.assert_called_once_with(mx_collection) + client.post.assert_called_once_with( + "datacollections/datacollection", + json={**mx_collection, "group_id": 99}, + ) + assert ( + mx_collection.get("group_id") is None + ) # ensures input data was not modified within _update_data_collection method + + +def test_store_energy_scan_success( + adapter, + client, +): + energyscan = default_energy_scan() + client.post.return_value = 123 + + result = adapter.store_energy_scan(energyscan) + + assert result == {"energyScanId": 123} + client.post.assert_called_once_with("events/energyscan", json=energyscan) + + +def test_store_energy_scan_failure(adapter, client): + energyscan = default_energy_scan() + client.post.side_effect = PyISPyBUnsuccessfulResponse("error") + + result = adapter.store_energy_scan(energyscan) + + assert result == {"energyScanId": -1} + adapter.logger.exception.assert_called_once_with( + "Failed to store energy scan in PyISPyB" + ) + + +def test_store_xfe_spectrum_success( + adapter, + client, +): + xfe_spectrum = default_xfe_spectrum() + client.post.return_value = 123 + + result = adapter.store_xfe_spectrum(xfe_spectrum) + + assert result == {"xfeFluorescenceSpectrumId": 123} + client.post.assert_called_once_with( + "events/xfe-fluorescence-spectrum", json=xfe_spectrum + ) + + +def test_store_xfe_spectrum_failure( + adapter, + client, +): + xfe_spectrum = default_xfe_spectrum() + client.post.side_effect = PyISPyBUnsuccessfulResponse("error") + + result = adapter.store_xfe_spectrum(xfe_spectrum) + + assert result == {"xfeFluorescenceSpectrumId": -1} + adapter.logger.exception.assert_called_once_with( + "Failed to store XFE fluorescence spectrum in PyISPyB" + ) + + +def test_update_bl_sample_success(adapter, client): + bl_sample = default_bl_sample() + expected_response = default_bl_sample() + client.patch.return_value = expected_response + + result = adapter.update_bl_sample(bl_sample) + + assert result == expected_response + client.patch.assert_called_once_with( + "samples/1115", + json=bl_sample, + ) + + +def test_update_bl_sample_missing_bl_sample_id(adapter, client): + bl_sample = { + "name": "sample1", + } + + result = adapter.update_bl_sample(bl_sample) + + assert result == {} + client.patch.assert_not_called() + adapter.logger.error.assert_called_once_with("Missing blSampleId") + + +def test_update_bl_sample_patch_failure(adapter, client): + bl_sample = default_bl_sample() + client.patch.side_effect = PyISPyBUnsuccessfulResponse("Update failed") + + result = adapter.update_bl_sample(bl_sample) + + assert result == {} + client.patch.assert_called_once_with( + "samples/1115", + json=bl_sample, + ) + adapter.logger.exception.assert_called_once_with( + "Failed to update beamline sample in PyISPyB" + ) diff --git a/test/test_pyispyb_rest_client.py b/test/test_pyispyb_rest_client.py new file mode 100644 index 0000000000..c1949b0133 --- /dev/null +++ b/test/test_pyispyb_rest_client.py @@ -0,0 +1,283 @@ +from datetime import datetime, timedelta, timezone +from json.decoder import JSONDecodeError +from unittest.mock import MagicMock + +import pytest + +from mxcubecore.HardwareObjects.abstract.PyISPyBRestClient import ( + AuthenticationExpired, + NoTokenException, + PyISPyBRestClient, + PyISPyBUnsuccessfulResponse, +) + + +@pytest.fixture +def client(): + return PyISPyBRestClient(rest_root="http://localhost/ispyb/api/v1/") + + +def build_response( + status_code=200, + json_data=None, + text="ok", +): + response = MagicMock() + + response.status_code = status_code + response.text = text + response.url = "http://localhost/ispyb/api/v1/test" + + if isinstance(json_data, Exception): + response.json.side_effect = json_data + else: + response.json.return_value = json_data + + return response + + +# ========================================================= +# AUTHENTICATION +# ========================================================= + + +def test_authenticate_success(client): + login_response = { + "token": "access-token", + "refreshToken": "refresh-token", + "expiresIn": 300, + } + client.post = MagicMock(return_value=login_response) + + client.authenticate(user_name="testusr", token="testtkn") # noqa: S106 + + client.post.assert_called_once_with( + "auth/login", + json={"plugin": "keycloak", "login": "testusr", "token": "testtkn"}, + skip_refresh=True, + ) + assert client._access_token == "access-token" # noqa: S105 + assert client._refresh_token == "refresh-token" # noqa: S105 + assert client._session.headers["Authorization"] == "Bearer access-token" + assert client._token_expiry is not None + + +def test_authenticate_without_token_raises(client): + client.post = MagicMock(return_value={}) + + with pytest.raises(NoTokenException): + client.authenticate(user_name="pyispyb_admin", token="kc-token") # noqa: S106 + + +# ========================================================= +# TOKEN REFRESH +# ========================================================= + + +def test_refresh_access_token_success(client): + client._refresh_token = "refresh-token" # noqa: S105 + response = build_response( + json_data={ + "token": "new-access-token", + "refreshToken": "new-refresh-token", + "expiresIn": 300, + } + ) + client._session.post = MagicMock(return_value=response) + + client._refresh_access_token() + + assert client._access_token == "new-access-token" # noqa: S105 + assert client._refresh_token == "new-refresh-token" # noqa: S105 + assert client._session.headers["Authorization"] == "Bearer new-access-token" + + +def test_refresh_access_token_without_refresh_token_raises(client): + client._refresh_token = None + + with pytest.raises(AuthenticationExpired): + client._refresh_access_token() + + +def test_refresh_access_token_failure_raises(client): + client._refresh_token = "invalid-token" # noqa: S105 + response = build_response( + status_code=401, + text="Unauthorized", + ) + client._session.post = MagicMock(return_value=response) + + with pytest.raises(AuthenticationExpired): + client._refresh_access_token() + + +def test_refresh_access_token_malformed_response(client): + client._refresh_token = "refresh-token" # noqa: S105 + response = build_response( + json_data=JSONDecodeError( + "invalid json", + "doc", + 0, + ) + ) + client._session.post = MagicMock(return_value=response) + + with pytest.raises(AuthenticationExpired): + client._refresh_access_token() + + +# ========================================================= +# GET REQUESTS +# ========================================================= + + +def test_get_success(client): + client._token_expiry = datetime.now(timezone.utc) + timedelta(hours=1) + response = build_response(json_data={"ok": True}) + client._session.get = MagicMock(return_value=response) + + result = client.get("users") + + assert result == {"ok": True} + client._session.get.assert_called_once() + + +def test_get_refreshes_expired_token(client): + client._token_expiry = datetime.now(timezone.utc) - timedelta(seconds=1) + client._refresh_access_token = MagicMock() + response = build_response(json_data={"ok": True}) + client._session.get = MagicMock(return_value=response) + + client.get("users") + + client._refresh_access_token.assert_called_once() + + +def test_get_retries_after_401(client): + client._token_expiry = datetime.now(timezone.utc) + timedelta(hours=1) + unauthorized = build_response( + status_code=401, + text="Unauthorized", + ) + success = build_response(json_data={"ok": True}) + client._session.get = MagicMock(side_effect=[unauthorized, success]) + client._refresh_access_token = MagicMock() + + result = client.get("users") + + assert result == {"ok": True} + client._refresh_access_token.assert_called_once() + assert client._session.get.call_count == 2 + + +# ========================================================= +# POST REQUESTS +# ========================================================= + + +def test_post_success(client): + client._token_expiry = datetime.now(timezone.utc) + timedelta(hours=1) + response = build_response(json_data={"created": True}) + client._session.post = MagicMock(return_value=response) + + result = client.post( + "users", + json={"name": "mohsen"}, + ) + + assert result == {"created": True} + + +# ========================================================= +# RESPONSE DECODING +# ========================================================= + + +def test_decode_json_response_returns_results(client): + response = build_response(json_data={"results": [1, 2, 3]}) + + result = client.decode_json_response(response) + + assert result == [1, 2, 3] + + +def test_decode_json_response_returns_json(client): + response = build_response(json_data={"ok": True}) + + result = client.decode_json_response(response) + + assert result == {"ok": True} + + +def test_decode_json_response_raises_on_http_error(client): + response = build_response( + status_code=500, + text="Server error", + ) + + with pytest.raises(PyISPyBUnsuccessfulResponse): + client.decode_json_response(response) + + +def test_decode_json_response_raises_on_invalid_json(client): + response = build_response( + json_data=JSONDecodeError( + "invalid json", + "doc", + 0, + ) + ) + + with pytest.raises(JSONDecodeError): + client.decode_json_response(response) + + +# ========================================================= +# STORE TOKENS +# ========================================================= + + +def test_store_tokens_success(client): + tokens = { + "token": "new-access-token", + "refreshToken": "new-refresh-token", + "expiresIn": 300, + } + + client._store_tokens(tokens) + + assert client._access_token == "new-access-token" # noqa: S105 + assert client._refresh_token == "new-refresh-token" # noqa: S105 + assert client._session.headers["Authorization"] == "Bearer new-access-token" + + +@pytest.mark.parametrize("tokens", [None, [], 0]) +def test_store_tokens_wrong_response_type(tokens, client): + with pytest.raises(NoTokenException): + client._store_tokens(tokens) + + +def test_store_tokens_missing_access_token(client): + tokens = { + "refreshToken": "new-refresh-token", + "expiresIn": 300, + } + with pytest.raises(NoTokenException): + client._store_tokens(tokens) + + +# ========================================================= +# PROXIES +# ========================================================= + + +def test_update_proxies(client): + proxies = { + "http": "http://proxy", + "https": "https://proxy", + } + + client.update_proxies(proxies) + + assert client._session.proxies["http"] == "http://proxy" + assert client._session.proxies["https"] == "https://proxy"