diff --git a/pyproject.toml b/pyproject.toml index db5470e..c26d208 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ dynamic = ["version"] dependencies = [ "requests", "boto3", - "pydantic>=2.0", + "pydantic>=2.0,<2.12", "pydantic-settings>=2.0", ] diff --git a/src/aind_data_access_api/document_db.py b/src/aind_data_access_api/document_db.py index a1c6fe1..7bd8fb0 100644 --- a/src/aind_data_access_api/document_db.py +++ b/src/aind_data_access_api/document_db.py @@ -625,6 +625,11 @@ def _deregister_asset_url(self) -> str: """Url to deregister (delete) an asset in DocDB and Code Ocean""" return f"https://{self.host}/{self.version}/assets/deregister" + @property + def _add_qc_evaluation_url(self) -> str: + """Url to add QC evaluation(s) or other QC content to a data asset""" + return f"https://{self.host}/{self.version}/add_qc_evaluation" + def generate_data_summary(self, record_id: str) -> Dict[str, Any]: """Get an LLM-generated summary for a data asset.""" url = f"{self._data_summary_url}/{record_id}" @@ -666,6 +671,25 @@ def deregister_asset(self, s3_location: str) -> Dict[str, Any]: response.raise_for_status() return response.json() + def add_qc_evaluation( + self, data_asset_id: str, qc_contents: Dict[str, Any] + ) -> Dict[str, Any]: + """Add one or more QC evaluations (or other QC content) + to a data asset.""" + + qc_contents_with_id = {**qc_contents, "data_asset_id": data_asset_id} + data = json.dumps(qc_contents_with_id) + signed_header = self._signed_request( + method="POST", url=self._add_qc_evaluation_url, data=data + ) + response = self.session.post( + url=self._add_qc_evaluation_url, + headers=dict(signed_header.headers), + data=data, + ) + response.raise_for_status() + return response.json() + class AnalysisDbClient(Client): """Class to manage reading and writing to analysis db""" diff --git a/src/aind_data_access_api/helpers/data_schema.py b/src/aind_data_access_api/helpers/data_schema.py index 184a661..a5fc244 100644 --- a/src/aind_data_access_api/helpers/data_schema.py +++ b/src/aind_data_access_api/helpers/data_schema.py @@ -1,10 +1,10 @@ """Module for convenience functions for the data access API.""" import json -import pandas as pd -from typing import List, Optional, Union from datetime import datetime, timezone +from typing import List, Optional, Union +import pandas as pd from aind_data_schema.core.quality_control import QualityControl from aind_data_access_api.document_db import MetadataDbClient diff --git a/tests/helpers/test_data_schema.py b/tests/helpers/test_data_schema.py index dafc386..90edf83 100644 --- a/tests/helpers/test_data_schema.py +++ b/tests/helpers/test_data_schema.py @@ -3,26 +3,27 @@ import json import os import unittest -import pandas as pd from datetime import datetime from pathlib import Path from unittest.mock import MagicMock, patch + +import pandas as pd from aind_data_schema.core.quality_control import ( - QualityControl, QCEvaluation, QCMetric, QCStatus, - Status, + QualityControl, Stage, + Status, ) from aind_data_schema_models.modalities import Modality from aind_data_access_api.helpers.data_schema import ( get_quality_control_by_id, get_quality_control_by_name, - get_quality_control_value_df, - get_quality_control_status_df, get_quality_control_by_names, + get_quality_control_status_df, + get_quality_control_value_df, ) TEST_DIR = Path(os.path.dirname(os.path.realpath(__file__))).parent diff --git a/tests/test_document_db.py b/tests/test_document_db.py index 475f7d3..a98cc95 100644 --- a/tests/test_document_db.py +++ b/tests/test_document_db.py @@ -1016,6 +1016,77 @@ def test_deregister_asset( ) self.assertEqual(response_message, response) + @patch("boto3.session.Session") + @patch("botocore.auth.SigV4Auth.add_auth") + @patch("requests.Session.post") + def test_add_qc_evaluation_success( + self, + mock_post: MagicMock, + mock_auth: MagicMock, + mock_session: MagicMock, + ): + """Tests add_qc_evaluation method success case""" + mock_creds = MagicMock() + mock_creds.access_key = "abc" + mock_creds.secret_key = "efg" + mock_session.return_value.region_name = "us-west-2" + mock_session.get_credentials.return_value = mock_creds + mock_response = Response() + mock_response.status_code = 200 + response_message = {"acknowledged": True, "matchedCount": 1} + mock_response._content = json.dumps(response_message).encode("utf-8") + mock_post.return_value = mock_response + + client = MetadataDbClient(**self.example_client_args) + qc_contents = { + "qc_evaluation": { + "modality": {"name": "ecephys", "abbreviation": "ecephys"}, + "stage": "Raw data", + "name": "Test QC", + "metrics": [], + } + } + response = client.add_qc_evaluation("fake-uuid", qc_contents) + mock_auth.assert_called_once() + mock_post.assert_called_once_with( + url="https://example.com/v1/add_qc_evaluation", + headers={"Content-Type": "application/json"}, + data=json.dumps({**qc_contents, "data_asset_id": "fake-uuid"}), + ) + self.assertEqual(response, response_message) + + @patch("boto3.session.Session") + @patch("botocore.auth.SigV4Auth.add_auth") + @patch("requests.Session.post") + def test_add_qc_evaluation_failure( + self, + mock_post: MagicMock, + mock_auth: MagicMock, + mock_session: MagicMock, + ): + """Tests add_qc_evaluation failure case""" + mock_creds = MagicMock() + mock_creds.access_key = "abc" + mock_creds.secret_key = "efg" + mock_session.return_value.region_name = "us-west-2" + mock_session.return_value.get_credentials.return_value = mock_creds + mock_response = Response() + mock_response.status_code = 400 + mock_response._content = json.dumps( + { + "detail": "Invalid QC evaluation submission. " + "'qc_evaluation' field missing required keys.", + "error": "BadRequest", + } + ).encode("utf-8") + mock_post.return_value = mock_response + + client = MetadataDbClient(**self.example_client_args) + qc_contents = {"foo": "bar"} + with self.assertRaises(requests.exceptions.HTTPError) as e: + client.add_qc_evaluation("fake-uuid", qc_contents) + self.assertIn("400 Client Error", str(e.exception)) + class TestAnalysisDbClient(unittest.TestCase): """Test methods in AnalysisDbClient class."""