add support for adding attachments

src/aind_slims_api/core.py

@@ -8,6 +8,7 @@
     methods and integration with SlimsBaseModel subtypes
 """
 
+import base64
 import logging
 from copy import deepcopy
 from functools import lru_cache
@@ -211,9 +212,21 @@ def fetch_attachments(
             ),
         )
 
-    def fetch_attachment_content(self, attachment: SlimsAttachment) -> Response:
-        """Fetch attachment content for a given attachment."""
-        return self.db.slims_api.get(f"repo/{attachment.pk}")
+    def fetch_attachment_content(
+        self,
+        attachment: int | SlimsAttachment,
+    ) -> Response:
+        """Fetch attachment content for a given attachment.
+
+        Parameters
+        -----------
+        attachment: int | SlimsAttachment
+            The primary key of the attachment or an attachment object
+        """
+        if isinstance(attachment, SlimsAttachment):
+            attachment = attachment.pk
+
+        return self.db.slims_api.get(f"repo/{attachment}")
 
     @lru_cache(maxsize=None)
     def fetch_pk(self, table: SLIMS_TABLES, *args, **kwargs) -> int | None:
@@ -302,3 +315,38 @@ def update_model(self, model: SlimsBaseModel, *args, **kwargs):
             ),
         )
         return type(model).model_validate(rtn)
+
+    def add_attachment_content(
+        self,
+        record: SlimsBaseModel,
+        name: str,
+        content: bytes | str,
+    ) -> int:
+        """Add an attachment to a SLIMS record
+
+        Returns
+        -------
+        int: Primary key of the attachment added.
+
+        Notes
+        -----
+        - Returned attachment does not contain the name of the attachment in
+         Slims, this requires a separate fetch.
+        """
+        if record.pk is None:
+            raise ValueError("Cannot add attachment to a record without a pk")
+
+        if isinstance(content, str):
+            content = content.encode("utf-8")
+
+        response = self.db.slims_api.post(
+            url="repo",
+            body={
+                "attm_name": name,
+                "atln_recordPk": record.pk,
+                "atln_recordTable": record._slims_table,
+                "contents": base64.b64encode(content).decode("utf-8"),
+            },
+        )
+        response.raise_for_status()
+        return int(response.text)

src/aind_slims_api/models/attachment.py

@@ -24,5 +24,5 @@ class SlimsAttachment(SlimsBaseModel):
     """
 
     pk: int = Field(..., alias="attm_pk")
-    name: str = Field(..., alias="attm_name")
+    name: str = Field(alias="attm_name")
     _slims_table = "Attachment"

src/aind_slims_api/models/base.py

@@ -6,7 +6,7 @@
 from typing import ClassVar, Optional
 
 from pydantic import BaseModel, ValidationInfo, field_serializer, field_validator
-from slims.internal import Column as SlimsColumn
+from slims.internal import Column as SlimsColumn  # type: ignore
 
 from aind_slims_api.models.utils import _find_unit_spec
 from aind_slims_api.types import SLIMS_TABLES
@@ -42,7 +42,7 @@ class MyModel(SlimsBaseModel):
     def _validate(cls, value, info: ValidationInfo):
         """Validates a field, accounts for Quantities"""
         if isinstance(value, SlimsColumn):
-            if value.datatype == "QUANTITY":
+            if value.datatype == "QUANTITY" and info.field_name is not None:
                 unit_spec = _find_unit_spec(cls.model_fields[info.field_name])
                 if unit_spec is None:
                     msg = (

src/aind_slims_api/models/behavior_session.py

@@ -18,12 +18,43 @@ class SlimsBehaviorSession(SlimsBaseModel):
 
     Examples
     --------
+    Read a session.
+
+    >>> from datetime import datetime
     >>> from aind_slims_api import SlimsClient
     >>> from aind_slims_api.models import SlimsMouseContent
     >>> client = SlimsClient()
     >>> mouse = client.fetch_model(SlimsMouseContent, barcode="00000000")
     >>> behavior_sessions = client.fetch_models(SlimsBehaviorSession,
     ...  mouse_pk=mouse.pk, sort=["date"])
+    >>> curriculum_attachments = client.fetch_attachments(behavior_sessions[0])
+
+    Write a new session.
+    >>> from aind_slims_api.models import SlimsInstrument, SlimsUser
+    >>> trainer = client.fetch_model(SlimsUser, username="LKim")
+    >>> instrument = client.fetch_model(SlimsInstrument, name="323_EPHYS1_OPTO")
+    >>> added = client.add_model(
+    ...  SlimsBehaviorSession(
+    ...      cnvn_fk_content=mouse.pk,
+    ...      cnvn_cf_fk_instrument=instrument.pk,
+    ...      cnvn_cf_fk_trainer=[trainer.pk],
+    ...      cnvn_cf_notes="notes",
+    ...      cnvn_cf_taskStage="stage",
+    ...      cnvn_cf_task="task",
+    ...      cnvn_cf_taskSchemaVersion="0.0.1",
+    ...      cnvn_cf_stageIsOnCurriculum=True,
+    ...      cnvn_cf_scheduledDate=datetime(2021, 1, 2),
+    ...  )
+    ... )
+
+    Add a curriculum attachment to the session (Attachment content isn't
+     available immediately.)
+    >>> import json
+    >>> attachment_pk = client.add_attachment_content(
+    ...  added,
+    ...  "curriculum",
+    ...  json.dumps({"curriculum_key": "curriculum_value"}),
+    ... )
     """
 
     pk: int | None = Field(default=None, alias="cnvn_pk")

tests/test_behavior_session.py

@@ -5,7 +5,7 @@
 import unittest
 from datetime import datetime
 from pathlib import Path
-from unittest.mock import MagicMock, patch, call
+from unittest.mock import MagicMock, call, patch
 
 from slims.internal import Record
 
@@ -128,9 +128,7 @@ def test_fetch_behavior_session_content_events_success_sort_list(
     @patch("aind_slims_api.core.logger")
     @patch("slims.slims.Slims.add")
     def test_write_behavior_session_content_events_success(
-        self,
-        mock_add: MagicMock,
-        mock_log_info: MagicMock
+        self, mock_add: MagicMock, mock_log_info: MagicMock
     ):
         """Test write_behavior_session_content_events success"""
         mock_add.return_value = self.example_write_sessions_response
@@ -143,9 +141,7 @@ def test_write_behavior_session_content_events_success(
         )
         self.assertTrue(all((item.mouse_pk == self.example_mouse.pk for item in added)))
         self.assertTrue(len(added) == len(self.example_behavior_sessions))
-        mock_log_info.assert_has_calls(
-            [call.info('SLIMS Add: ContentEvent/79')]
-        )
+        mock_log_info.assert_has_calls([call.info("SLIMS Add: ContentEvent/79")])
 
 
 if __name__ == "__main__":

tests/test_core.py

@@ -13,6 +13,7 @@
 
 from aind_slims_api.core import SlimsAttachment, SlimsClient
 from aind_slims_api.exceptions import SlimsRecordNotFound
+from aind_slims_api.models.behavior_session import SlimsBehaviorSession
 from aind_slims_api.models.unit import SlimsUnit
 
 RESOURCES_DIR = Path(os.path.dirname(os.path.realpath(__file__))) / "resources"
@@ -26,6 +27,7 @@ class TestSlimsClient(unittest.TestCase):
     example_fetch_mouse_response: list[Record]
     example_fetch_user_response: list[Record]
     example_fetch_attachment_response: list[Record]
+    example_add_attachments_response_text = str
 
     @classmethod
     def setUpClass(cls):
@@ -53,6 +55,9 @@ def get_response(attribute_name: str):
         cls.example_fetch_attachment_response = get_response(
             "example_fetch_attachments_response.json_entity"
         )
+        cls.example_add_attachments_response_text = (
+            RESOURCES_DIR / "example_add_attachment_content_response_text.txt"
+        ).read_text()
 
     def test_rest_link(self):
         """Tests rest_link method with both queries and no queries."""
@@ -247,40 +252,67 @@ def test_fetch_model_no_records(self, mock_slims_fetch: MagicMock):
         with self.assertRaises(SlimsRecordNotFound):
             self.example_client.fetch_model(SlimsUnit)
 
-    def test_fetch_attachments(self):
+    @patch("slims.internal._SlimsApi.get_entities")
+    def test_fetch_attachments(self, mock_get_entities: MagicMock):
         """Tests fetch_attachments method success."""
         # slims_api is dynamically added to slims client
-        assert len(self.example_fetch_attachment_response) == 1
-        with patch.object(
-            self.example_client.db.slims_api,
-            "get_entities",
-            return_value=self.example_fetch_attachment_response,
-        ):
-            unit = SlimsUnit.model_validate(
-                Record(
-                    json_entity=self.example_fetch_unit_response[0].json_entity,
-                    slims_api=self.example_client.db.slims_api,
-                )
+        mock_get_entities.return_value = self.example_fetch_attachment_response
+        unit = SlimsUnit.model_validate(
+            Record(
+                json_entity=self.example_fetch_unit_response[0].json_entity,
+                slims_api=self.example_client.db.slims_api,
             )
-            attachments = self.example_client.fetch_attachments(
-                unit,
-            )
-            assert len(attachments) == 1
+        )
+        attachments = self.example_client.fetch_attachments(
+            unit,
+        )
+        self.assertEqual(1, len(attachments))
 
-    def test_fetch_attachment_content(self):
+    @patch("slims.slims._SlimsApi.get")
+    def test_fetch_attachment_content(self, mock_get: MagicMock):
         """Tests fetch_attachment_content method success."""
-        # slims_api is dynamically added to slims client
-        with patch.object(
-            self.example_client.db.slims_api,
-            "get",
-            return_value=Response(),
-        ):
-            self.example_client.fetch_attachment_content(
-                SlimsAttachment(
-                    attm_name="test",
-                    attm_pk=1,
-                )
+        mock_get.return_value = Response()
+        self.example_client.fetch_attachment_content(
+            SlimsAttachment(
+                attm_name="test",
+                attm_pk=1,
+            )
+        )
+        mock_get.assert_called_once()
+
+    @patch("slims.internal._SlimsApi.post")
+    def test_add_attachment_content(self, mock_post: MagicMock):
+        """Tests add_attachment_content method success."""
+        mock_response = MagicMock()
+        mock_response.text.return_value = self.example_add_attachments_response_text
+        mock_post.return_value = mock_response
+        unit_pk = 1
+        self.example_client.add_attachment_content(
+            SlimsUnit(
+                unit_name="test",
+                unit_pk=unit_pk,
+            ),
+            "test",
+            "some test content",
+        )
+        self.assertEqual(mock_post.call_count, 1)
+        self.assertEqual(
+            mock_post.mock_calls[0].kwargs["body"]["atln_recordPk"], unit_pk
+        )
+
+    @patch("slims.internal._SlimsApi.post")
+    def test_add_attachment_content_no_pk(self, mock_post: MagicMock):
+        """Tests add_attachment_content method failure due to lack of pk."""
+        mock_response = MagicMock()
+        mock_response.text.return_value = self.example_add_attachments_response_text
+        mock_post.return_value = mock_response
+        with self.assertRaises(ValueError):
+            self.example_client.add_attachment_content(
+                SlimsBehaviorSession(),
+                "test",
+                "some test content",
             )
+        mock_post.assert_not_called()
 
     @patch("logging.Logger.error")
     def test__validate_model_invalid_model(self, mock_log: MagicMock):
@@ -303,8 +335,8 @@ def test__validate_model_invalid_model(self, mock_log: MagicMock):
                 ),
             ],
         )
-        assert len(validated) == 1
-        assert mock_log.call_count == 1
+        self.assertEqual(1, len(validated))
+        self.assertEqual(1, mock_log.call_count)
 
     def test_resolve_model_alias_invalid(self):
         """Tests resolve_model_alias method raises expected error with an