updated tests to match new interface

shared/datastore/tests/integration/base_crud_test.py

@@ -7,7 +7,7 @@ from io import BytesIO
 import minio
 import pytest
 
-from shared.datastore import delete, get, put
+from shared.datastore import Datastore
 
 
 def same_data(
@@ -39,77 +39,72 @@ def same_data(
 
 
 def test_should_get_data(
-    minio_client,
+    datastore: Datastore,
     data_in_minio,
 ):
-    data, bucket_name, object_name = data_in_minio
-    received_data = get(
-        client=minio_client,
-        bucket_name=bucket_name,
+    # ARRANGE
+    data, _, object_name = data_in_minio
+    # ACT
+    received_data = datastore._get(
         object_name=object_name,
     )
-
+    # ASSERT
     assert isinstance(data, BytesIO)
     assert same_data(data, received_data)
 
 
 def test_should_delete_data(
-    minio_client,
+    datastore: Datastore,
     data_in_minio,
 ):
-    _, bucket_name, object_name = data_in_minio
-    delete(
-        client=minio_client,
-        bucket_name=bucket_name,
+    # ARRANGE
+    _, _, object_name = data_in_minio
+    # ACT
+    datastore._delete(
         object_name=object_name,
     )
+    # ASSERT
     with pytest.raises(minio.error.S3Error):
-        _ = get(
-            client=minio_client,
-            bucket_name=bucket_name,
+        _ = datastore._get(
             object_name=object_name,
         )
 
 
 def test_should_put_data(
-    minio_client,
+    datastore: Datastore,
     data,
 ):
-    # prepare variables
-    minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
+    # ARRANGE
     minio_object_name = str(os.getenv('MINIO_OBJECT_NAME'))
     buffer = BytesIO(data)
-    put(
-        client=minio_client,
+    # ACT
+    datastore._put(
+        object_name=minio_object_name,
         buffer=buffer,
-        bucket_name=minio_bucket_name,
-        object_name=minio_object_name,
     )
-    received_data = get(
-        client=minio_client,
-        bucket_name=minio_bucket_name,
+    received_data = datastore._get(
         object_name=minio_object_name,
     )
+    # ASSERT
     assert isinstance(received_data, BytesIO)
     assert same_data(received_data, buffer)
 
 
 def test_should_update_data(
-    minio_client,
+    datastore: Datastore,
     data_in_minio,
 ):
-    buffer, bucket_name, object_name = data_in_minio
-    put(
-        client=minio_client,
+    # ARRANGE
+    buffer, _, object_name = data_in_minio
+    # ACT
+    datastore._put(
+        object_name=object_name,
         buffer=buffer,
-        bucket_name=bucket_name,
-        object_name=object_name,
     )
-    received_data = get(
-        client=minio_client,
-        bucket_name=bucket_name,
+    received_data = datastore._get(
         object_name=object_name,
     )
+    # ASSERT
     assert same_data(received_data, buffer)
 
 

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,
+    )

shared/datastore/tests/integration/connect_test.py

@@ -2,9 +2,9 @@
 
 from minio import Minio
 
-from shared.datastore import connect_minio
+from shared.datastore import Datastore
 
 
 def test_should_return_correct_type():
-    client = connect_minio()
-    assert isinstance(client, Minio)
+    with Datastore() as ds:
+        assert isinstance(ds._client, Minio)

shared/datastore/tests/integration/image_crud_test.py

@@ -1,9 +1,10 @@
 """Integration tests related to image CRUD."""
 
 import numpy as np
+import pytest
 from PIL import Image
 
-# from shared.datastore import connect_minio, get_image, put_image
+from shared.datastore import Datastore
 
 
 def same_image(
@@ -25,49 +26,57 @@ def same_image(
     return True
 
 
-# def test_should_get_image(
-#     image_in_minio,
-# ):
-#     image, object_name = image_in_minio
-#     client = connect_minio()
-#     received_image = get_image(
-#         client=client,
-#         object_name=object_name,
-#     )
-#     assert isinstance(image, Image.Image)
-#     assert received_image == image
+def test_should_get_image(
+    datastore: Datastore,
+    image_in_minio: tuple[Image.Image, str],
+):
+    # ARRANGE
+    image, object_name = image_in_minio
+    # ACT
+    received_image = datastore.get_image(
+        object_name=object_name,
+    )
+    # ASSERT
+    assert isinstance(image, Image.Image)
+    assert same_image(image, received_image)
 
 
-# def test_should_put_image(
-#     image,
-# ):
-#     client = connect_minio()
-#     object_name = put_image(
-#         client=client,
-#         image=image,
-#     )
-#     assert isinstance(object_name, str)
-#     assert len(object_name) > 0
-#     received_image = get_image(
-#         client=client,
-#         object_name=object_name,
-#     )
-#     assert same_image(received_image, image)
+def test_should_put_image(
+    datastore: Datastore,
+    image: Image.Image,
+):
+    # ARRANGE
+    object_name = datastore.put_image(
+        image=image,
+    )
+    assert isinstance(object_name, str)
+    assert len(object_name) > 0
+    # ACT
+    received_image = datastore.get_image(
+        object_name=object_name,
+    )
+    # ASSERT
+    assert same_image(image, received_image)
 
 
-# def test_should_update_image(
-#     image_in_minio,
-# ):
-#     image, object_name = image_in_minio
-#     client = connect_minio()
-#     object_name = put_image(
-#         client=client,
-#         image=image,
-#     )
-#     assert isinstance(object_name, str)
-#     assert len(object_name) > 0
-#     received_image = get_image(
-#         client=client,
-#         object_name=object_name,
-#     )
-#     assert received_image == image
+def test_should_update_image(
+    datastore: Datastore,
+    image_in_minio: tuple[Image.Image, str],
+):
+    # ARRANGE
+    image, object_name = image_in_minio
+    object_name = datastore.put_image(
+        image=image,
+    )
+    assert isinstance(object_name, str)
+    assert len(object_name) > 0
+    # ACT
+    received_image = datastore.get_image(
+        object_name=object_name,
+    )
+    # ASSERT
+    assert same_image(image, received_image)
+
+
+if __name__ == '__main__':
+    pytest.main()

shared/datastore/tests/integration/model_crud_test.py

@@ -1,52 +1,90 @@
 """Integration tests related to model CRUD."""
 
-# from torch.nn import Module
+from collections import OrderedDict
 
-# from shared.datastore import connect_minio, get_model, put_model
+from torch.nn import Module
 
-# def test_should_get_model(
-#     model_in_minio,
-# ):
-#     model, object_name = model_in_minio
-#     client = connect_minio()
-#     received_model = get_model(
-#         client=client,
-#         object_name=object_name,
-#     )
-#     assert isinstance(model, Module)
-#     assert received_model == model
+from model.src.models import VisualCommunicationModel
+from shared.datastore import Datastore
 
 
-# def test_should_put_model(
-#     model,
-# ):
-#     client = connect_minio()
-#     object_name = put_model(
-#         client=client,
-#         model=model,
-#     )
-#     assert isinstance(object_name, str)
-#     assert len(object_name) > 0
-#     received_model = get_model(
-#         client=client,
-#         object_name=object_name,
-#     )
-#     assert received_model == model
+def same_model(
+    model_a: Module,
+    model_b: Module,
+) -> bool:
+    """Check if two models are the same class, have the same number of
+    parameters and contain the same weights."""
+    assert isinstance(model_a, Module)
+    assert isinstance(model_b, Module)
+    # compare model classes
+    assert type(model_a) is type(model_b)
+    # compare number of parameters
+    params_a = list(model_a.parameters())
+    params_b = list(model_b.parameters())
+    if len(params_a) != len(params_b):
+        return False
+    # compare model weights
+    for p_a, p_b in zip(params_a, params_b):
+        if p_a.data.ne(p_b.data).sum() > 0:
+            return False
+    return True
 
 
-# def test_should_update_model(
-#     model_in_minio,
-# ):
-#     model, object_name = model_in_minio
-#     client = connect_minio()
-#     object_name = put_model(
-#         client=client,
-#         model=model,
-#     )
-#     assert isinstance(object_name, str)
-#     assert len(object_name) > 0
-#     received_model = get_model(
-#         client=client,
-#         object_name=object_name,
-#     )
-#     assert received_model == model
+def test_should_get_model(
+    datastore: Datastore,
+    model_in_minio: tuple[Module, str],
+):
+    # ARRANGE
+    model, object_name = model_in_minio
+    # ACT
+    model_data = datastore.get_model(
+        object_name=object_name,
+    )
+    assert isinstance(model_data, OrderedDict)
+    received_model = VisualCommunicationModel().to('cpu')
+    received_model.load_state_dict(model_data)
+    # ASSERT
+    assert same_model(model, received_model)
+
+
+def test_should_put_model(
+    datastore: Datastore,
+    model: Module,
+):
+    # ARRANGE
+    object_name = datastore.put_model(
+        model=model,
+    )
+    assert isinstance(object_name, str)
+    assert len(object_name) > 0
+    # ACT
+    model_data = datastore.get_model(
+        object_name=object_name,
+    )
+    assert isinstance(model_data, OrderedDict)
+    received_model = VisualCommunicationModel().to('cpu')
+    received_model.load_state_dict(model_data)
+    # ASSERT
+    assert same_model(model, received_model)
+
+
+def test_should_update_model(
+    datastore: Datastore,
+    model_in_minio: tuple[Module, str],
+):
+    # ARRANGE
+    model, object_name = model_in_minio
+    object_name = datastore.put_model(
+        model=model,
+    )
+    assert isinstance(object_name, str)
+    assert len(object_name) > 0
+    # ACT
+    model_data = datastore.get_model(
+        object_name=object_name,
+    )
+    assert isinstance(model_data, OrderedDict)
+    received_model = VisualCommunicationModel().to('cpu')
+    received_model.load_state_dict(model_data)
+    # ASSERT
+    assert same_model(model, received_model)