updated tests to match new interface

shared/datastore/tests/integration/conftest.py

@@ -3,14 +3,18 @@
 import os
 import random
 from collections.abc import Iterator
+from hashlib import md5
 from io import BytesIO
 from pathlib import Path
 
 import pytest
+import torch
 from dotenv import load_dotenv
-from minio import Minio
 from PIL import Image
 
+from model.src.models import VisualCommunicationModel
+from shared.datastore import Datastore
+
 env_var_map = {
     'MINIO_BUCKET_NAME': 'test-bucket',
     'MINIO_OBJECT_NAME': '46KXJMFIAPVLM0TKRFZR5YPPTVJ6PJNX',
@@ -39,35 +43,34 @@ def populate_env(
 
 
 @pytest.fixture(scope='session')
-def minio_client(
+def datastore(
     populate_env,
-) -> Iterator[Minio]:
+) -> Iterator[Datastore]:
     # prepare arguments
-    minio_endpoint = str(os.getenv('MINIO_ENDPOINT'))
-    minio_access_key = str(os.getenv('MINIO_ACCESS_KEY'))
-    minio_secret_key = str(os.getenv('MINIO_SECRET_KEY'))
     minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
     # connect to minio
-    client = Minio(
-        endpoint=minio_endpoint,
-        access_key=minio_access_key,
-        secret_key=minio_secret_key,
-        secure=False,
-    )
+    datastore_client = Datastore()
+    datastore_client.connect()
+    assert datastore_client._client is not None
     # ensure bucket exists
-    if not client.bucket_exists(bucket_name=minio_bucket_name):
-        client.make_bucket(bucket_name=minio_bucket_name)
+    if not datastore_client._client.bucket_exists(bucket_name=minio_bucket_name):
+        datastore_client._client.make_bucket(bucket_name=minio_bucket_name)
     # expose client
-    yield client
+    yield datastore_client
     # remove objects left behind by tests
-    for obj in client.list_objects(bucket_name=minio_bucket_name, recursive=True):
-        client.remove_object(
+    for obj in datastore_client._client.list_objects(
+        bucket_name=minio_bucket_name,
+        recursive=True,
+    ):
+        datastore_client._client.remove_object(
             bucket_name=obj.bucket_name,
             object_name=obj.object_name,
         )
     # remove bucket
-    client.remove_bucket(bucket_name=minio_bucket_name)
-    assert not client.bucket_exists(bucket_name=minio_bucket_name)
+    datastore_client._client.remove_bucket(bucket_name=minio_bucket_name)
+    assert not datastore_client._client.bucket_exists(bucket_name=minio_bucket_name)
+    # disconnect from minio
+    datastore_client.close()
 
 
 @pytest.fixture
@@ -81,9 +84,10 @@ def data() -> Iterator[bytes]:
 
 @pytest.fixture
 def data_in_minio(
-    minio_client,
-    data,
+    datastore: Datastore,
+    data: bytes,
 ) -> Iterator[tuple[BytesIO, str, str]]:
+    assert datastore._client is not None
     # prepare arguments
     minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
     minio_object_name = str(os.getenv('MINIO_OBJECT_NAME'))
@@ -93,7 +97,7 @@ def data_in_minio(
     num_bytes = len(buffer.getvalue())
     buffer.seek(0)
     # send data to bucket
-    minio_client.put_object(
+    datastore._client.put_object(
         bucket_name=minio_bucket_name,
         object_name=minio_object_name,
         length=num_bytes,
@@ -102,7 +106,7 @@ def data_in_minio(
     # expose data
     yield buffer, minio_bucket_name, minio_object_name
     # clean up
-    minio_client.remove_object(
+    datastore._client.remove_object(
         bucket_name=minio_bucket_name,
         object_name=minio_object_name,
     )
@@ -116,13 +120,77 @@ def image() -> Iterator[Image.Image]:
     yield image
 
 
-# @pytest.fixture
-# def image_in_minio(
-#     image: Image.Image,
-# ) -> tuple(Image.Image, str):
-#     #
+@pytest.fixture
+def image_in_minio(
+    datastore: Datastore,
+    image: Image.Image,
+) -> Iterator[tuple[Image.Image, str]]:
+    assert datastore._client is not None
+    # prepare arguments
+    minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
+    # save data to buffer
+    buffer = BytesIO()
+    image.save(buffer, 'png')
+    # get md5 of buffer
+    checksum = md5(buffer.getbuffer()).hexdigest()
+    # build object path
+    object_path = f'images/{checksum}'
+    # prepare for saving
+    num_bytes = len(buffer.getvalue())
+    buffer.seek(0)
+    # send data to bucket
+    datastore._client.put_object(
+        bucket_name=minio_bucket_name,
+        object_name=object_path,
+        length=num_bytes,
+        data=buffer,
+    )
+    # expose image and object name
+    yield image, checksum
+    # cleanup
+    datastore._client.remove_object(
+        bucket_name=minio_bucket_name,
+        object_name=object_path,
+    )
 
 
-#     # expose image and object name
-#     yield image, object_name
-#     # cleanup
+@pytest.fixture
+def model() -> Iterator[torch.nn.Module]:
+    # generate model
+    model = VisualCommunicationModel().to('cpu')
+    # expose model
+    yield model
+
+
+@pytest.fixture
+def model_in_minio(
+    datastore: Datastore,
+    model: torch.nn.Module,
+) -> Iterator[tuple[torch.nn.Module, str]]:
+    assert datastore._client is not None
+    # prepare arguments
+    minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
+    # save data to buffer
+    buffer = BytesIO()
+    torch.save(model.state_dict(), buffer)
+    # get md5 of image
+    checksum = md5(buffer.getbuffer()).hexdigest()
+    # build object path
+    object_path = f'models/{checksum}'
+    # prepare for saving
+    num_bytes = len(buffer.getvalue())
+    buffer.seek(0)
+    # send data to bucket
+    datastore._client.put_object(
+        bucket_name=minio_bucket_name,
+        object_name=object_path,
+        length=num_bytes,
+        data=buffer,
+    )
+    # expose model and object name
+    yield model, checksum
+    # cleanup
+    datastore._client.remove_object(
+        bucket_name=minio_bucket_name,
+        object_name=object_path,
+    )