diff --git a/src/py/mat3ra/api_client/endpoints/properties.py b/src/py/mat3ra/api_client/endpoints/properties.py index 6eaade0..30b88c1 100644 --- a/src/py/mat3ra/api_client/endpoints/properties.py +++ b/src/py/mat3ra/api_client/endpoints/properties.py @@ -51,3 +51,41 @@ def delete(self, id_): def update(self, id_, modifier): raise NotImplementedError + + def list_for_job(self, job_id): + """ + List properties for a job grouped by unit. + + Args: + job_id (str): Job ID. + + Returns: + list[dict]: List of {"unit_id": str, "properties": [str, ...]}. + """ + properties = self.list(query={"source.info.jobId": job_id}) + units = {} + for prop in properties: + unit_id = prop["source"]["info"]["unitId"] + if unit_id not in units: + units[unit_id] = [] + units[unit_id].append(prop["data"]["name"]) + return [{"unit_id": unit_id, "properties": names} for unit_id, names in units.items()] + + def get_for_job(self, job_id, property_name=None, unit_id=None): + """ + Get property data for a job, optionally filtered by property name and/or unit. + + Args: + job_id (str): Job ID. + property_name (str, optional): Property name (e.g., "band_gaps", "total_energy"). + unit_id (str, optional): Unit flowchart ID (e.g., "pw-nscf"). + + Returns: + list[dict]: List of property data dicts. + """ + query = {"source.info.jobId": job_id} + if property_name: + query["data.name"] = property_name + if unit_id: + query["source.info.unitId"] = unit_id + return [prop["data"] for prop in self.list(query=query)] diff --git a/tests/py/unit/test_properties.py b/tests/py/unit/test_properties.py index 9cad06c..62f6ccb 100644 --- a/tests/py/unit/test_properties.py +++ b/tests/py/unit/test_properties.py @@ -1,3 +1,4 @@ +import json from unittest import mock from mat3ra.api_client.endpoints.properties import PropertiesEndpoints @@ -5,6 +6,24 @@ ENDPOINT_NAME = "properties" +JOB_ID = "ukmnfWw9Q5ryXHK4X" +PROPERTY_NAME_0 = "total_energy" +PROPERTY_NAME_1 = "band_gaps" +UNIT_ID_0 = "pw-relax" +UNIT_ID_1 = "pw-nscf" + +MOCK_PROPERTY_0 = { + "data": {"name": PROPERTY_NAME_0, "value": -260.698, "units": "eV"}, + "source": {"info": {"jobId": JOB_ID, "unitId": UNIT_ID_0}}, +} +MOCK_PROPERTY_1 = { + "data": {"name": PROPERTY_NAME_1, "values": [{"type": "direct", "value": 0.5, "units": "eV"}]}, + "source": {"info": {"jobId": JOB_ID, "unitId": UNIT_ID_1}}, +} + +MOCK_PROPERTIES_RESPONSE = json.dumps({"status": "success", "data": [MOCK_PROPERTY_0, MOCK_PROPERTY_1]}) +MOCK_SINGLE_PROPERTY_RESPONSE = json.dumps({"status": "success", "data": [MOCK_PROPERTY_1]}) + class EndpointCharacteristicUnitTest(EntityEndpointsUnitTest): """ @@ -23,3 +42,45 @@ def test_list(self, mock_request): @mock.patch("requests.sessions.Session.request") def test_get(self, mock_request): self.get(mock_request) + + @mock.patch("requests.sessions.Session.request") + def test_list_for_job(self, mock_request): + mock_request.return_value = self.mock_response(MOCK_PROPERTIES_RESPONSE) + result = self.endpoints.list_for_job(JOB_ID) + print(result) + self.assertEqual(result, [ + {"unit_id": UNIT_ID_0, "properties": [PROPERTY_NAME_0]}, + {"unit_id": UNIT_ID_1, "properties": [PROPERTY_NAME_1]}, + ]) + sent_query = json.loads(mock_request.call_args[1]["params"]["query"]) + self.assertEqual(sent_query["source.info.jobId"], JOB_ID) + + @mock.patch("requests.sessions.Session.request") + def test_get_for_job(self, mock_request): + mock_request.return_value = self.mock_response(MOCK_PROPERTIES_RESPONSE) + result = self.endpoints.get_for_job(JOB_ID) + self.assertEqual(len(result), 2) + sent_query = json.loads(mock_request.call_args[1]["params"]["query"]) + self.assertEqual(sent_query["source.info.jobId"], JOB_ID) + self.assertNotIn("data.name", sent_query) + self.assertNotIn("source.info.unitId", sent_query) + + @mock.patch("requests.sessions.Session.request") + def test_get_for_job_filtered_by_name(self, mock_request): + mock_request.return_value = self.mock_response(MOCK_SINGLE_PROPERTY_RESPONSE) + result = self.endpoints.get_for_job(JOB_ID, PROPERTY_NAME_1) + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["name"], PROPERTY_NAME_1) + sent_query = json.loads(mock_request.call_args[1]["params"]["query"]) + self.assertEqual(sent_query["data.name"], PROPERTY_NAME_1) + + + @mock.patch("requests.sessions.Session.request") + def test_get_for_job_filtered_by_unit_id_and_name(self, mock_request): + mock_request.return_value = self.mock_response(MOCK_SINGLE_PROPERTY_RESPONSE) + result = self.endpoints.get_for_job(JOB_ID, PROPERTY_NAME_1, unit_id=UNIT_ID_1) + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["name"], PROPERTY_NAME_1) + sent_query = json.loads(mock_request.call_args[1]["params"]["query"]) + self.assertEqual(sent_query["source.info.unitId"], UNIT_ID_1) + self.assertEqual(sent_query["data.name"], PROPERTY_NAME_1)