object_based_docstore_approach #68

Merged
brian merged 64 commits from object_based_docstore_approach into main 2025-04-16 00:03:25 +02:00
45 changed files with 2630 additions and 0 deletions
Showing only changes of commit 090a1cd8bb - Show all commits
+83
View File
@@ -0,0 +1,83 @@
"""Definition of docstore interface."""
from __future__ import annotations
from abc import ABC, abstractmethod
from shared.docstore.src.classes import ModelData, VisualCommunication
class DocstoreInterface(ABC):
"""Docstore interface class."""
@abstractmethod
def connect(
self,
) -> None:
pass
@abstractmethod
def close(
self,
) -> None:
pass
@abstractmethod
def __enter__(
self,
) -> DocstoreInterface:
pass
@abstractmethod
def __exit__(
self,
exc_type,
exc_val,
exc_tb,
) -> None:
pass
@abstractmethod
def count_documents(
self,
only_with_annotation: bool,
) -> int:
pass
@abstractmethod
def list_names(
self,
only_with_annotation: bool,
) -> list[str]:
pass
@abstractmethod
def upsert_visual_comminucations(
self,
visual_communication_list: list[VisualCommunication],
) -> None:
pass
@abstractmethod
def upsert_annotations(
self,
visual_communication_name: str,
annotations: ModelData,
) -> None:
pass
@abstractmethod
def upsert_prediction(
self,
visual_communication_name: str,
predictions: ModelData,
) -> None:
pass
@abstractmethod
def get_visual_communication(
self,
with_annotation: bool,
name: str | None = None,
) -> VisualCommunication:
pass
+241
View File
@@ -0,0 +1,241 @@
"""Definition of docstore mongodb implementation."""
from __future__ import annotations
import logging
import os
from typing import Any
from pymongo import MongoClient
from pymongo.collection import Collection
from pymongo.database import Database
from shared.utils import check_env
from .classes import ModelData, VisualCommunication
from .docstore_interface import DocstoreInterface
from .exceptions import NoDocumentFoundException
class DocstoreMongo(DocstoreInterface):
"""Docstore interface."""
def __init__(self):
# ensure necessary env vars available
var_list = {
'MONGO_ENDPOINT',
'MONGO_DB',
'MONGO_COLLECTION',
}
check_env(var_list)
# prepare internal variables
self.client: MongoClient | None = None
self.db: Database | None = None
self.collection: Collection | None = None
def connect(self) -> None:
"""Connect to Mongo server."""
# prepare arguments
mongo_endpoint = str(os.getenv('MONGO_ENDPOINT'))
mongo_database = str(os.getenv('MONGO_DB'))
mongo_collection = str(os.getenv('MONGO_COLLECTION'))
# connect client
client = MongoClient(mongo_endpoint)
database = client[mongo_database]
collection = database[mongo_collection]
# set unique index on 'name'
collection.create_index(keys='name', unique=True)
# persist state
self._client = client
self._database = database
self._collection = collection
def close(self) -> None:
"""Close connection to Mongo server."""
self._client.close()
self._client = None
self._database = None
self._collection = None
def __enter__(self) -> DocstoreMongo:
self.connect()
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
if any(
(
exc_type is not None,
exc_val is not None,
exc_tb is not None,
),
):
logging.error('error while exiting context')
self.close()
def count_documents(
self,
only_with_annotation: bool = False,
):
"""Get the total number of Visual Communication documents matching the
filter."""
assert isinstance(only_with_annotation, bool)
# build query
query = {}
if only_with_annotation:
query['annotation'] = {'$ne': None}
# execute query
num_docs = self._collection.count_documents(filter=query)
return num_docs
def list_names(
self,
only_with_annotation: bool,
) -> list[str]:
"""List names of Visual Communication documents that match the
filters."""
assert isinstance(only_with_annotation, bool)
assert isinstance(self._collection, Collection)
# build query
query = {}
if only_with_annotation:
query['annotation'] = {'$ne': None}
# execute query
res_list = self._collection.find(
filter=query,
projection={
'_id': False, # don't get id
'name': True, # include document name
},
)
# extract info
name_list = [elem['name'] for elem in res_list]
return name_list
def upsert_visual_comminucations(
self,
visual_communication_list: list[VisualCommunication],
) -> None:
"""Upsert Visual Communication document in the database."""
assert isinstance(visual_communication_list, list)
assert all(
isinstance(vis_com, VisualCommunication)
for vis_com in visual_communication_list
)
assert isinstance(self._collection, Collection)
# convert to dict
doc_list = [vis_com.model_dump() for vis_com in visual_communication_list]
# insert documents
response = self._collection.insert_many(documents=doc_list)
# check response
if not response.acknowledged:
raise OSError('failed inserting documents')
def upsert_annotations(
self,
visual_communication_name: str,
annotations: ModelData,
) -> None:
"""Upsert annotation for the document with matching name."""
assert isinstance(visual_communication_name, str)
assert len(visual_communication_name) > 0
assert isinstance(annotations, ModelData)
assert isinstance(self._collection, Collection)
# convert to dict
doc = annotations.model_dump()
# build query
query = {
'name': visual_communication_name,
}
update = {
'$set': {
'annotation': doc,
},
}
# execute query
res = self._collection.update_one(
filter=query,
update=update,
upsert=True,
)
# check response
if not res.acknowledged:
raise OSError(
f'failed upserting annotations for {visual_communication_name}',
)
def upsert_prediction(
self,
visual_communication_name: str,
predictions: ModelData,
) -> None:
"""Upsert prediction for the document with matching name."""
assert isinstance(visual_communication_name, str)
assert len(visual_communication_name) > 0
assert isinstance(annotations, ModelData)
assert isinstance(self._collection, Collection)
# convert to dict
doc = predictions.model_dump()
# build query
query = {
'name': visual_communication_name,
}
update = {
'$set': {
'prediction': doc,
},
}
# execute query
res = self._collection.update_one(
filter=query,
update=update,
upsert=True,
)
# check response
if not res.acknowledged:
raise OSError(
f'failed upserting predictions for {visual_communication_name}',
)
def get_visual_communication(
self,
with_annotation: bool = False,
name: str | None = None,
) -> VisualCommunication:
"""Get a random Visual Communication document that matches annotation
filter.
If name is not specified, a random document matching filter is
returned.
"""
assert isinstance(with_annotation, bool)
if name is not None:
assert isinstance(name, str)
assert len(name) > 0
# build query
query: dict[str, Any] = {}
if with_annotation:
query['annotation'] = {'$ne': None}
else:
query['annotation'] = {'$eq': None}
if name is not None:
query['name'] = {'$eq': name}
# execute query
res_list = self._collection.aggregate(
pipeline=[
{
'$match': query, # find using filters
},
{
'$sample': {
'size': 1, # get one random
},
},
],
)
doc_list = list(res_list)
# check result
if len(doc_list) == 0:
raise NoDocumentFoundException()
# convert
vis_com = VisualCommunication.model_validate(doc_list[0])
return vis_com
@@ -0,0 +1,45 @@
"""Definition of unittests for Docstore MongoDB implementation."""
import os
from unittest import TestCase
from unittest.mock import MagicMock
from pymongo import MongoClient
from pymongo.collection import Collection
from pymongo.database import Database
from shared.docstore.src.docstore_mongo import DocstoreMongo
class TestDocstoreMongoImplementation(TestCase):
def setUp(self):
# define relevant env vars
self.env_var_map = {
'MONGO_ENDPOINT': '192.168.1.2:27017',
'MONGO_DB': 'visual_critical_discourse_analysis',
'MONGO_COLLECTION': 'test-collection',
}
# set env vars
for key, val in self.env_var_map.items():
os.environ[key] = val
# set other variables
self.mongo_client_mock = MagicMock(MongoClient)
self.mongo_database_mock = MagicMock(Database)
self.mongo_collection_mock = MagicMock(Collection)
def tearDown(self):
# clear env vars
for key in self.env_var_map:
_ = os.environ.pop(key, default=None)
# reset reuseable mocks
self.mongo_client_mock.reset_mock()
self.mongo_database_mock.reset_mock()
self.mongo_collection_mock.reset_mock()
def test_instantiation_should_fail_when_env_not_set(self):
# ensure env not set
self.tearDown()
# run test
with self.assertRaises(OSError):
_ = DocstoreMongo()
+8
View File
@@ -0,0 +1,8 @@
from .src import ImageRepository, ModelRepository, VisualCommunicationRepository
from .src.dto import (
HexadecimalString,
ImageData,
ModelData,
VisualCommunicationData,
VisualCommunicationValues,
)
+3
View File
@@ -0,0 +1,3 @@
from .image_repository import ImageRepository
from .model_repository import ModelRepository
from .visual_communication_repository import VisualCommunicationRepository
+7
View File
@@ -0,0 +1,7 @@
from .hexadecimal_string import HexadecimalString
from .image_data import ImageData
from .model_data import ModelData
from .visual_communication_data import (
VisualCommunicationData,
VisualCommunicationValues,
)
@@ -0,0 +1,37 @@
"""Definition of BytesIO Pydantic Annotation."""
from io import BytesIO
from typing import Any
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema
class BytesIOPydanticAnnotation:
"""Pydantic annotation that defines input validation, as well as general
and json serialization."""
@classmethod
def validate_input(cls, v: Any, handler) -> BytesIO:
"""Pydantic-related function to validate input on instantiation."""
if isinstance(v, BytesIO):
return v
s = handler(v)
return BytesIO(s)
@classmethod
def __get_pydantic_core_schema__(
cls,
source_type,
_handler,
) -> core_schema.CoreSchema:
assert source_type is BytesIO
return core_schema.no_info_wrap_validator_function(
function=cls.validate_input,
schema=core_schema.str_schema(),
serialization=core_schema.to_string_ser_schema(),
)
@classmethod
def __get_pydantic_json_schema__(cls, _core_schema, handler) -> JsonSchemaValue:
return handler(core_schema.str_schema())
@@ -0,0 +1,24 @@
"""Definition of Checksum DTO."""
from __future__ import annotations
import re
class HexadecimalString(str):
"""Hexadecimal-string class."""
def __new__(cls, string):
# ensure proper input format
pattern = r'[0-9-a-fA-F]{32}'
match = re.match(pattern, string)
if match is None:
raise ValueError(f'format does not match a hexadecimal-string: {string}')
return super().__new__(cls, string)
def __repr__(self) -> str:
class_name = self.__class__.__name__
return f"{class_name}('{self}')"
def __reduce__(self):
return self.__class__, (self,)
@@ -0,0 +1,38 @@
"""Definition of HexadecimalString Pydantic Annotation."""
from typing import Any
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema
from .hexadecimal_string import HexadecimalString
class HexadecimalStringPydanticAnnotation:
"""Pydantic annotation that defines input validation, as well as general
and json serialization."""
@classmethod
def validate_input(cls, v: Any, handler) -> HexadecimalString:
"""Pydantic-related function to validate input on instantiation."""
if isinstance(v, HexadecimalString):
return v
s = handler(v)
return HexadecimalString(s)
@classmethod
def __get_pydantic_core_schema__(
cls,
source_type,
_handler,
) -> core_schema.CoreSchema:
assert source_type is HexadecimalString
return core_schema.no_info_wrap_validator_function(
function=cls.validate_input,
schema=core_schema.str_schema(),
serialization=core_schema.to_string_ser_schema(),
)
@classmethod
def __get_pydantic_json_schema__(cls, _core_schema, handler) -> JsonSchemaValue:
return handler(core_schema.str_schema())
+16
View File
@@ -0,0 +1,16 @@
"""Definition of VisualData DTO."""
from typing import Annotated
from PIL import Image
from pydantic import Field
from .image_pydantic_annotation import ImagePydanticAnnotation
from .type_checking_base_model import TypeCheckingBaseModel
class ImageData(TypeCheckingBaseModel):
"""Visual data class."""
image: Annotated[Image.Image, ImagePydanticAnnotation]
name: str = Field(min_length=1)
@@ -0,0 +1,37 @@
"""Definition of BytesIO Pydantic Annotation."""
from typing import Any
from PIL import Image
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema
class ImagePydanticAnnotation:
"""Pydantic annotation that defines input validation, as well as general
and json serialization."""
@classmethod
def validate_input(cls, v: Any, handler) -> Image.Image:
"""Pydantic-related function to validate input on instantiation."""
if isinstance(v, Image.Image):
return v
s = handler(v)
return Image.open(s)
@classmethod
def __get_pydantic_core_schema__(
cls,
source_type,
_handler,
) -> core_schema.CoreSchema:
assert source_type is Image.Image
return core_schema.no_info_wrap_validator_function(
function=cls.validate_input,
schema=core_schema.str_schema(),
serialization=core_schema.to_string_ser_schema(),
)
@classmethod
def __get_pydantic_json_schema__(cls, _core_schema, handler) -> JsonSchemaValue:
return handler(core_schema.str_schema())
+60
View File
@@ -0,0 +1,60 @@
"""Definition of ModelData DTO."""
from __future__ import annotations
from hashlib import md5
from io import BytesIO
from typing import Annotated
import torch
from pydantic import Field
from .bytes_io_pydantic_annotation import BytesIOPydanticAnnotation
from .hexadecimal_string import HexadecimalString
from .hexadecimal_string_pydantic_annotation import HexadecimalStringPydanticAnnotation
from .type_checking_base_model import TypeCheckingBaseModel
class ModelData(TypeCheckingBaseModel):
"""Model Data DTO."""
buffer: Annotated[BytesIO, BytesIOPydanticAnnotation]
buffer_checksum: Annotated[HexadecimalString, HexadecimalStringPydanticAnnotation]
class_name: str = Field(
min_length=1,
description='name of model class to generate data.',
)
@staticmethod
def calculate_checksum(buffer: BytesIO) -> HexadecimalString:
"""Calculate buffer checksum."""
checksum = md5(buffer.getbuffer()).hexdigest()
return HexadecimalString(checksum)
@staticmethod
def model_to_buffer(model: torch.nn.Module) -> BytesIO:
"""Save model to buffer."""
assert isinstance(model, torch.nn.Module)
buffer = BytesIO()
torch.save(model.state_dict(), buffer)
return buffer
@classmethod
def from_model(cls, model: torch.nn.Module) -> ModelData:
"""Instantiate from torch module."""
assert isinstance(model, torch.nn.Module)
# get model name
class_name = type(model).__name__
# save data to buffer
buffer = cls.model_to_buffer(model)
buffer = BytesIO()
torch.save(model.state_dict(), buffer)
# calculate checksum
buffer_checksum = cls.calculate_checksum(buffer)
# instantiate from buffer
data = cls(
buffer=buffer,
buffer_checksum=buffer_checksum,
class_name=class_name,
)
return data
@@ -0,0 +1,12 @@
"""Definition of TypeCheckingBaseModel class."""
from pydantic import BaseModel, ConfigDict
class TypeCheckingBaseModel(BaseModel):
"""BaseModel with added type checking on input types."""
model_config = ConfigDict(
validate_assignment=True, # argument type checking
frozen=True, # ensure data immutability
)
@@ -0,0 +1,2 @@
from .visual_communication_data import VisualCommunicationData
from .visual_communication_values import VisualCommunicationValues
@@ -0,0 +1,11 @@
"""Definition of AngleValues DTO."""
from .values_model import ValuesModel
class AngleValues(ValuesModel):
"""Angle values DTO."""
high: float
eye_level: float
low: float
@@ -0,0 +1,10 @@
"""Definition of ContactValues DTO."""
from .values_model import ValuesModel
class ContactValues(ValuesModel):
"""Contact values DTO."""
offer: float
demand: float
@@ -0,0 +1,11 @@
"""Definition of DistanceValues DTO."""
from .values_model import ValuesModel
class DistanceValues(ValuesModel):
"""Distance values DTO."""
long: float
medium: float
close: float
@@ -0,0 +1,12 @@
"""Definition of FramingValues DTO."""
from .values_model import ValuesModel
class FramingValues(ValuesModel):
"""Framing values DTO."""
frame_lines: float
empty_space: float
colour_contrast: float
form_contrast: float
@@ -0,0 +1,11 @@
"""Definition of InformationValueValues DTO."""
from .values_model import ValuesModel
class InformationValueValues(ValuesModel):
"""Information value values DTO."""
given_new: float
ideal_real: float
central_marginal: float
@@ -0,0 +1,11 @@
"""Definition of ModalityColorValues DTO."""
from .values_model import ValuesModel
class ModalityColorValues(ValuesModel):
"""Modality color values DTO."""
high: float
medium: float
low: float
@@ -0,0 +1,11 @@
"""Definition of ModalityDepthValues DTO."""
from .values_model import ValuesModel
class ModalityDepthValues(ValuesModel):
"""Modality depth values DTO."""
high: float
medium: float
low: float
@@ -0,0 +1,11 @@
"""Definition of ModalityLightingValues DTO."""
from .values_model import ValuesModel
class ModalityLightingValues(ValuesModel):
"""Modality lighting values DTO."""
high: float
medium: float
low: float
@@ -0,0 +1,10 @@
"""Definition of PointOfViewValues DTO."""
from .values_model import ValuesModel
class PointOfViewValues(ValuesModel):
"""Point-of-view values DTO."""
frontal: float
oblique: float
@@ -0,0 +1,13 @@
"""Definition of SalienceValues DTO."""
from .values_model import ValuesModel
class SalienceValues(ValuesModel):
"""Salience values DTO."""
size: float
colour: float
tone: float
form: float
positioning: float
@@ -0,0 +1,60 @@
"""Definition of ValuesModel base class."""
from __future__ import annotations
import random
from pydantic import BaseModel, ConfigDict
from torch import Tensor
class ValuesModel(BaseModel):
"""ValuesModel base class."""
model_config = ConfigDict(
validate_assignment=True, # argument type checking
frozen=True, # ensure data immutability
)
@classmethod
def list_fields(cls) -> list[str]:
"""List options that are stored as attributes."""
return list(cls.model_fields.keys())
@classmethod
def from_random(cls):
"""Instantiate with random numbers."""
kwargs = {field: random.random() for field in cls.list_fields()}
return cls(**kwargs)
@classmethod
def from_choice(cls, option: str) -> ValuesModel:
"""Instantiate from choice."""
assert isinstance(option, str)
assert len(option) > 0
allowed_options_list = cls.list_fields()
if option not in allowed_options_list:
raise ValueError(f'option {option} must be in {allowed_options_list}')
# generate field values
kwargs = {field: 0 for field in allowed_options_list}
# set chosen value to max probability
kwargs[option] = 1
return cls(**kwargs)
@classmethod
def from_tensor(cls, tensor: Tensor):
"""Instantiate from list of values."""
assert tensor.size(dim=0) == 1, f'tensor batch larger than 1: {tensor}'
data_list = [float(t.item()) for t in tensor[0]]
kwargs = dict(zip(cls.list_fields(), data_list))
return cls(**kwargs)
def highest_score_field(self) -> str:
"""Return name of field with highest score."""
model_dict = self.model_dump()
return max(model_dict, key=lambda k: model_dict[k])
def highest_score_value(self) -> float:
"""Return value of field with highest score."""
model_dict = self.model_dump()
return max(model_dict.values())
@@ -0,0 +1,14 @@
"""Definition of VisualCommunicationData DTO."""
from pydantic import Field
from ..type_checking_base_model import TypeCheckingBaseModel
from .visual_communication_values import VisualCommunicationValues
class VisualCommunicationData(TypeCheckingBaseModel):
"""Visual communication data class."""
name: str = Field(min_length=1)
annotation: VisualCommunicationValues | None = None
prediction: VisualCommunicationValues | None = None
@@ -0,0 +1,49 @@
"""Definition of VisualCommunicationValues DTO."""
from __future__ import annotations
from ..type_checking_base_model import TypeCheckingBaseModel
from .angle_values import AngleValues
from .contact_values import ContactValues
from .distance_values import DistanceValues
from .framing_values import FramingValues
from .information_value_values import InformationValueValues
from .modality_color_values import ModalityColorValues
from .modality_depth_values import ModalityDepthValues
from .modality_lighting_values import ModalityLightingValues
from .point_of_view_values import PointOfViewValues
from .salience_values import SalienceValues
from .visual_syntax_values import VisualSyntaxValues
class VisualCommunicationValues(TypeCheckingBaseModel):
"""Visual communication values class."""
visual_syntax: VisualSyntaxValues
contact: ContactValues
angle: AngleValues
point_of_view: PointOfViewValues
distance: DistanceValues
modality_lighting: ModalityLightingValues
modality_color: ModalityColorValues
modality_depth: ModalityDepthValues
information_value: InformationValueValues
framing: FramingValues
salience: SalienceValues
@classmethod
def from_random(cls) -> VisualCommunicationValues:
"""Create a random instance."""
return cls(
visual_syntax=VisualSyntaxValues.from_random(),
contact=ContactValues.from_random(),
angle=AngleValues.from_random(),
point_of_view=PointOfViewValues.from_random(),
distance=DistanceValues.from_random(),
modality_lighting=ModalityLightingValues.from_random(),
modality_color=ModalityColorValues.from_random(),
modality_depth=ModalityDepthValues.from_random(),
information_value=InformationValueValues.from_random(),
framing=FramingValues.from_random(),
salience=SalienceValues.from_random(),
)
@@ -0,0 +1,26 @@
"""Definition of VisualSyntaxValues DTO."""
from .values_model import ValuesModel
class VisualSyntaxValues(ValuesModel):
"""Visual syntax values DTO."""
non_transactional_action: float
non_transactional_reaction: float
unidirectional_transactional_action: float
unidirectional_transactional_reaction: float
bidirectional_transactional_action: float
bidirectional_transactional_reaction: float
conversion: float
speech_process: float
classification_overt_taxonomy: float
analytical_exhaustive: float
analytical_disarranged: float
analytical_temporal: float
analytical_distributed: float
analytical_topological: float
analytical_exploded: float
analytical_inclusive: float
symbolic_suggestive: float
symbolic_attributive: float
@@ -0,0 +1,74 @@
"""Definition of ImageRepository class."""
from __future__ import annotations
from io import BytesIO
from pathlib import Path
from PIL import Image
from .dto import ImageData
from .implementations import MinioImplementation
from .interfaces import ImageInterface
class ImageRepository(ImageInterface, MinioImplementation):
"""Image repository class that handles CRUD functionality for
VisualData."""
def __enter__(self) -> ImageRepository:
self.connect()
return self
@staticmethod
def _build_path(image_name: str) -> Path:
"""Build object path."""
assert isinstance(image_name, str)
path = Path('images') / image_name
return path
def get_data(self, image_name: str) -> ImageData | None:
"""Get Visual data."""
assert isinstance(image_name, str)
assert len(image_name) > 0
# build path
path = self._build_path(image_name)
# get object from bucket
buffer = self._get(path)
# handle if no data found
if not buffer:
return None
# convert data
image = Image.open(buffer)
data = ImageData(image=image, name=image_name)
return data
def put_data(self, data: ImageData) -> None:
"""Put visual data."""
assert isinstance(data, ImageData)
# build path
path = self._build_path(data.name)
# save image to buffer
buffer = BytesIO()
data.image.save(buffer, 'png')
# put object in bucket
self._put(path, buffer)
def remove_data(self, image_name: str) -> None:
"""Remove visual data."""
assert isinstance(image_name, str)
assert len(image_name) > 0
# build path
path = self._build_path(image_name)
# remove object
self._delete(path)
def list_names(self) -> list[str]:
"""List names of all images."""
# build path
path = self._build_path('')
# list object paths
obj_path_list = self._list_objects(path)
# strip prefix
name_list = [obj_path.split('/')[-1] for obj_path in obj_path_list]
return name_list
@@ -0,0 +1,2 @@
from .minio_implementation import MinioImplementation
from .mongo_implementation import MongoImplementation
@@ -0,0 +1,186 @@
"""MinIO implementation."""
from __future__ import annotations
import logging
import os
from io import BytesIO
from pathlib import Path
from traceback import format_exc
from minio import Minio
from shared.utils import check_env
from ..interfaces import DatabaseInterface
class MinioImplementation(DatabaseInterface):
"""MinIO basic CRUD implementation."""
def __init__(self):
# ensure necessary env vars available
var_list = {
'MINIO_ENDPOINT',
'MINIO_ACCESS_KEY',
'MINIO_SECRET_KEY',
'MINIO_BUCKET_NAME',
}
check_env(var_list)
# prepare internal variables
self._client: Minio | None = None
self._bucket_name: str | None = None
def connect(self):
"""Connect to MinIO server."""
# 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 client
client = Minio(
endpoint=minio_endpoint,
access_key=minio_access_key,
secret_key=minio_secret_key,
secure=False,
)
# ensure bucket exists
if not client.bucket_exists(bucket_name=minio_bucket_name):
logging.debug('creating bucket: %s', minio_bucket_name)
client.make_bucket(bucket_name=minio_bucket_name)
# persist state
self._client = client
self._bucket_name = minio_bucket_name
def close(self) -> None:
"""Close connection to MinIO server.
N.B. MinIO connection cannot be closed manually.
"""
self._client = None
self._bucket_name = None
def connected(self):
"""Check connection to Minio."""
res = isinstance(self._client, Minio)
logging.debug(res)
return res
def __enter__(self) -> MinioImplementation:
self.connect()
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
if any(
(
exc_type is not None,
exc_val is not None,
exc_tb is not None,
),
):
logging.error('error while exiting context')
self.close()
def _put(
self,
path: Path,
buffer: BytesIO,
) -> None:
"""Save in-memory buffer as object in MinIO."""
assert isinstance(path, Path)
assert isinstance(buffer, BytesIO)
assert isinstance(self._client, Minio)
assert isinstance(self._bucket_name, str)
# prepare for saving
num_bytes = len(buffer.getvalue())
buffer.seek(0)
# send data to bucket
try:
self._client.put_object(
bucket_name=self._bucket_name,
object_name=path.as_posix(),
length=num_bytes,
data=buffer,
)
logging.debug('saved data to %s', path)
except Exception as exc:
logging.error('failed saving data to MinIO')
raise exc
def _get(
self,
path: Path,
) -> BytesIO | None:
"""Get object from MinIO as in-memory buffer."""
assert isinstance(path, Path)
assert isinstance(self._client, Minio)
assert isinstance(self._bucket_name, str)
try:
# make request
response = self._client.get_object(
bucket_name=self._bucket_name,
object_name=path.as_posix(),
)
assert response.status == 200
# get buffer
buffer = BytesIO()
chunk_size = 2**14
while chunk := response.read(chunk_size):
buffer.write(chunk)
buffer.seek(0)
logging.debug('got %s', path)
return buffer
except Exception:
logging.error('failed getting data from MinIO')
logging.debug(format_exc())
return None
finally:
# close connection if established
if 'response' in locals():
response.close()
response.release_conn()
def _delete(
self,
path: Path,
) -> None:
"""Delete object from MinIO."""
assert isinstance(path, Path)
assert isinstance(self._client, Minio)
assert isinstance(self._bucket_name, str)
# remove object
try:
self._client.remove_object(
bucket_name=self._bucket_name,
object_name=path.as_posix(),
)
logging.debug('deleted %s', path)
except Exception as exc:
logging.error('failed deleting %s', path)
logging.debug(format_exc())
raise exc
def _list_objects(
self,
path: Path,
) -> list[str]:
"""List objects in bucket under path."""
assert isinstance(path, Path)
assert isinstance(self._client, Minio)
assert isinstance(self._bucket_name, str)
try:
# list objects
obj_list = self._client.list_objects(
bucket_name=self._bucket_name,
prefix=path.as_posix(),
recursive=True,
)
# extract info
name_list = [obj.object_name for obj in obj_list]
logging.debug('got %s objects matching %s', len(name_list), path)
return name_list
except Exception as exc:
logging.error('failed listing objects under %s', path)
logging.debug(format_exc())
raise exc
@@ -0,0 +1,145 @@
"""Mongo implementation."""
from __future__ import annotations
import logging
import os
import traceback
from pymongo import MongoClient
from pymongo.collection import Collection
from pymongo.database import Database
from pymongo.errors import ServerSelectionTimeoutError
from shared.utils import check_env
from ..interfaces import DatabaseInterface
class MongoImplementation(DatabaseInterface):
"""MongoDB basic CRUD implementation."""
def __init__(self):
# ensure necessary env vars available
var_list = {
'MONGO_ENDPOINT',
'MONGO_DB',
'MONGO_COLLECTION',
}
check_env(var_list)
# prepare internal variables
self.client: MongoClient | None = None
self.db: Database | None = None
self.collection: Collection | None = None
def connect(self) -> None:
"""Connect to Mongo server."""
# prepare arguments
mongo_endpoint = str(os.getenv('MONGO_ENDPOINT'))
mongo_database = str(os.getenv('MONGO_DB'))
mongo_collection = str(os.getenv('MONGO_COLLECTION'))
# connect client
client = MongoClient(mongo_endpoint)
database = client[mongo_database]
collection = database[mongo_collection]
# set unique index on 'name'
collection.create_index(keys='name', unique=True)
# persist state
self._client = client
self._database = database
self._collection = collection
def close(self) -> None:
"""Close connection to Mongo server."""
self._client.close()
self._client = None
self._database = None
self._collection = None
def connected(self) -> bool:
"""Check connection to Mongo."""
if self._client is None:
return False
try:
# trigger fetch data
_ = self._client.server_info()
res = True
except ServerSelectionTimeoutError:
res = False
logging.debug(res)
return res
def __enter__(self) -> MongoImplementation:
self.connect()
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
if any(
(
exc_type is not None,
exc_val is not None,
exc_tb is not None,
),
):
logging.error('error while exiting context')
traceback.print_exception(exc_type, exc_val, exc_tb)
self.close()
def _save(
self,
data: dict,
query: dict,
) -> None:
"""Save document in Mongo."""
assert isinstance(data, dict)
assert isinstance(query, dict)
assert 'name' in data
assert self.connected()
self._collection.update_one(
filter=query,
update={
'$set': data.copy(),
},
upsert=True,
)
logging.debug('Save %s', data)
def _get(
self,
query: dict,
) -> dict | None:
"""Get document from Mongo."""
assert isinstance(query, dict)
assert self.connected()
doc = self._collection.find_one(query, projection={'_id': False})
logging.debug('Found %s', doc)
return doc
def _delete(
self,
query: dict,
) -> None:
"""Remove document from Mongo."""
assert isinstance(query, dict)
assert self.connected()
doc = self._collection.delete_one(query)
logging.debug('Deleted %s', doc)
def _list_documents(self, key='name') -> list[str]:
"""List documents in Mongo."""
assert isinstance(key, str)
assert len(key) > 0
# build query
doc_list = list(
self._collection.find(
filter={},
projection={
'_id': False,
key: True,
},
),
)
# extract values
value_list = [doc[key] for doc in doc_list]
logging.debug('Got %s document(s)', len(doc_list))
return value_list
@@ -0,0 +1,4 @@
from .database_interface import DatabaseInterface
from .image_interface import ImageInterface
from .model_interface import ModelInterface
from .visual_communication_interface import VisualCommunicationInterface
@@ -0,0 +1,28 @@
"""Definition of DatabaseInterface class."""
from abc import ABC, abstractmethod
class DatabaseInterface(ABC):
"""Interface base class adding 'connect', 'close' and context
functionalities."""
@abstractmethod
def connect(self):
raise NotImplementedError()
@abstractmethod
def close(self):
raise NotImplementedError()
@abstractmethod
def connected(self) -> bool:
raise NotImplementedError()
@abstractmethod
def __enter__(self):
raise NotImplementedError()
@abstractmethod
def __exit__(self, exc_type, exc_val, exc_tb):
raise NotImplementedError()
@@ -0,0 +1,26 @@
"""Definition of ImageInterface."""
from abc import abstractmethod
from ..dto import ImageData
from .database_interface import DatabaseInterface
class ImageInterface(DatabaseInterface):
"""Image interface class."""
@abstractmethod
def get_data(self, image_name: str) -> ImageData | None:
raise NotImplementedError()
@abstractmethod
def put_data(self, data: ImageData) -> None:
raise NotImplementedError()
@abstractmethod
def remove_data(self, image_name: str) -> None:
raise NotImplementedError()
@abstractmethod
def list_names(self) -> list[str]:
raise NotImplementedError()
@@ -0,0 +1,26 @@
"""Definition of ModelInterface."""
from abc import abstractmethod
from ..dto import ModelData
from .database_interface import DatabaseInterface
class ModelInterface(DatabaseInterface):
"""Model interface class."""
@abstractmethod
def get_data(self, object_name: str) -> ModelData | None:
raise NotImplementedError()
@abstractmethod
def put_data(self, data: ModelData) -> None:
raise NotImplementedError()
@abstractmethod
def remove_data(self, object_name: str) -> None:
raise NotImplementedError()
@abstractmethod
def list_names(self) -> list[str]:
raise NotImplementedError()
@@ -0,0 +1,26 @@
"""Definition of VisualCommunicationInterface."""
from abc import abstractmethod
from ..dto import VisualCommunicationData
from .database_interface import DatabaseInterface
class VisualCommunicationInterface(DatabaseInterface):
"""Visual communication interface class."""
@abstractmethod
def get_data(self, name: str) -> VisualCommunicationData | None:
raise NotImplementedError()
@abstractmethod
def put_data(self, data: VisualCommunicationData) -> None:
raise NotImplementedError()
@abstractmethod
def remove_data(self, name: str) -> None:
raise NotImplementedError()
@abstractmethod
def list_names(self) -> list[str]:
raise NotImplementedError()
@@ -0,0 +1,77 @@
"""Definition of ModelRepository class."""
from __future__ import annotations
from pathlib import Path
from .dto import HexadecimalString, ModelData
from .implementations import MinioImplementation
from .interfaces import ModelInterface
class ModelRepository(ModelInterface, MinioImplementation):
"""Model repository class that handles CRUD functionality for ModelData."""
def __enter__(self) -> ModelRepository:
self.connect()
return self
@staticmethod
def _prefix() -> Path:
"""Object name prefix."""
return Path('models')
@classmethod
def _build_object_name(cls, data: ModelData) -> str:
"""Build object name from data."""
return f'{data.class_name}-{data.buffer_checksum}'
def get_data(self, object_name: str) -> ModelData | None:
"""Get model data."""
assert isinstance(object_name, str)
# build path
path = self._prefix() / object_name
# get object from bucket
buffer = self._get(path)
# handle if no data found
if not buffer:
return None
# extract info
class_name, buffer_checksum_str = object_name.split('-')
# convert data
buffer_checksum = HexadecimalString(buffer_checksum_str)
# instantiate data
data = ModelData(
buffer=buffer,
buffer_checksum=buffer_checksum,
class_name=class_name,
)
return data
def put_data(self, data: ModelData) -> None:
"""Put model data."""
assert isinstance(data, ModelData)
# build object name
object_name = self._build_object_name(data)
# build path
path = self._prefix() / object_name
# put object in bucket
self._put(path, data.buffer)
def remove_data(self, object_name: str) -> None:
"""Remove model data."""
assert isinstance(object_name, str)
# build path
path = self._prefix() / object_name
# remove object
self._delete(path)
def list_names(self) -> list[str]:
"""List names of all models."""
# build path
path = self._prefix()
# list object paths
obj_path_list = self._list_objects(path)
# strip prefix
name_list = [obj_path.split('/')[-1] for obj_path in obj_path_list]
return name_list
@@ -0,0 +1,56 @@
"""Definition of VisualCommunicationRepository class."""
from __future__ import annotations
from .dto import VisualCommunicationData
from .implementations import MongoImplementation
from .interfaces import VisualCommunicationInterface
class VisualCommunicationRepository(VisualCommunicationInterface, MongoImplementation):
"""Visual communication repository class that handles CRUD functionality
for VisualCommunicationData."""
def __enter__(self) -> VisualCommunicationRepository:
self.connect()
return self
def get_data(self, name: str) -> VisualCommunicationData | None:
"""Get visual communication data."""
assert isinstance(name, str)
assert len(name) > 0
# build query
query = {'name': name}
# get document from mongo
doc = self._get(query)
# handle if no data found
if doc is None:
return None
# instantiate object
data = VisualCommunicationData(**doc)
return data
def put_data(self, data: VisualCommunicationData) -> None:
"""Put visual communication data."""
assert isinstance(data, VisualCommunicationData)
# convert to dict
data = data.model_dump(mode='json')
# build query
query = {'name': data['name']}
# save to mongo
self._save(data, query)
def remove_data(self, name: str) -> None:
"""Remove visual communication data."""
assert isinstance(name, str)
assert len(name) > 0
# build query
query = {'name': name}
# remove document from mongo
self._delete(query)
def list_names(self) -> list[str]:
"""List names of all documents."""
# list documents
name_list = self._list_documents(key='name')
return name_list
@@ -0,0 +1,444 @@
"""Integration tests configuration."""
import os
import random
from collections.abc import Iterator
from io import BytesIO
from pathlib import Path
import minio
import pytest
from PIL import Image
from pymongo import MongoClient
from testcontainers.minio import MinioContainer
from testcontainers.mongodb import MongoDbContainer
from model.src.models import VisualCommunicationModel
from shared.repositories import (
HexadecimalString,
ImageData,
ImageRepository,
ModelData,
ModelRepository,
VisualCommunicationData,
VisualCommunicationRepository,
VisualCommunicationValues,
)
from shared.repositories.src.implementations import (
MinioImplementation,
MongoImplementation,
)
# set random seed for reproducibility
random.seed(13)
# define test static variables
MINIO_PORT = 9000
MINIO_ACCESS_KEY = 'minioadmin'
MINIO_SECRET_KEY = 'minioadmin'
MONGO_PORT = 27017
MONGO_USERNAME = 'admin'
MONGO_PASSWORD = 'password'
env_var_map = {
'MINIO_ENDPOINT': f'localhost:{MINIO_PORT}', # updated when container is running
'MINIO_ACCESS_KEY': MINIO_ACCESS_KEY,
'MINIO_SECRET_KEY': MINIO_SECRET_KEY,
'MINIO_BUCKET_NAME': 'test-bucket',
'MINIO_OBJECT_NAME': 'test-object',
'MINIO_IMAGE_NAME': 'test-image',
'MONGO_ENDPOINT': (
f'mongodb://{MONGO_USERNAME}:{MONGO_PASSWORD}' f'@localhost:{MONGO_PORT}/'
), # updated when container is running
'MONGO_DB': 'test-db',
'MONGO_COLLECTION': 'test-collection',
}
container_map = {
'minio': MinioContainer(
port=MINIO_PORT,
access_key=MINIO_ACCESS_KEY,
secret_key=MINIO_SECRET_KEY,
),
'mongo': MongoDbContainer(
port=MONGO_PORT,
username=MONGO_USERNAME,
password=MONGO_PASSWORD,
dbname=env_var_map['MONGO_DB'],
),
}
@pytest.fixture(scope='session', autouse=True)
def setup_infrastructure(request: pytest.FixtureRequest) -> None:
"""Prepare infrastructure for integration test."""
# prepare infrastructure
for container in container_map.values():
container.start()
# update env var map
minio_host = container_map['minio'].get_container_host_ip()
minio_port = container_map['minio'].get_exposed_port(MINIO_PORT)
env_var_map['MINIO_ENDPOINT'] = f'{minio_host}:{minio_port}'
mongo_host = container_map['mongo'].get_container_host_ip()
mongo_port = container_map['mongo'].get_exposed_port(MONGO_PORT)
env_var_map['MONGO_ENDPOINT'] = (
f'mongodb://{MONGO_USERNAME}:{MONGO_PASSWORD}@{mongo_host}:{mongo_port}/'
)
# ensure cleanup
def cleanup_infrastructure():
for container in container_map.values():
container.stop()
request.addfinalizer(cleanup_infrastructure)
@pytest.fixture(scope='session', autouse=True)
def populate_env(
request: pytest.FixtureRequest,
setup_infrastructure,
) -> None:
"""Populate environment with variables used for testing."""
# update env
for key, val in env_var_map.items():
os.environ[key] = val
# ensure cleanup
def cleanup_env():
for key in env_var_map:
_ = os.environ.pop(key, default=None)
request.addfinalizer(cleanup_env)
@pytest.fixture(scope='session')
def raw_minio_client(
setup_infrastructure,
populate_env,
) -> Iterator[minio.Minio]:
"""Raw Minio client fixture."""
# 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 client
client = minio.Minio(
endpoint=minio_endpoint,
access_key=minio_access_key,
secret_key=minio_secret_key,
secure=False,
)
# ensure bucket exists
if not client.bucket_exists(bucket_name=minio_bucket_name):
client.make_bucket(bucket_name=minio_bucket_name)
# expose client
yield client
# cleanup
object_list = client.list_objects(minio_bucket_name, recursive=True)
for obj in object_list:
client.remove_object(
bucket_name=obj.bucket_name,
object_name=obj.object_name,
)
client.remove_bucket(minio_bucket_name)
@pytest.fixture
def minio_client(
setup_infrastructure,
populate_env,
) -> Iterator[MinioImplementation]:
"""MinioImplementation fixture."""
# instantiate and connect client
minio_client = MinioImplementation()
minio_client.connect()
# expose client
yield minio_client
# cleanup
object_name_list = minio_client._list_objects(Path('*'))
for name in object_name_list:
minio_client._delete(Path(name))
minio_client.close()
@pytest.fixture
def buffer() -> Iterator[BytesIO]:
"""Bytes buffer fixture."""
# generate reproducible random data
data = random.randbytes(n=2**21) # 2 MB
# convert data
buffer = BytesIO(data)
# expose buffer
yield buffer
@pytest.fixture
def buffer_in_minio(
raw_minio_client: minio.Minio,
buffer: BytesIO,
) -> Iterator[tuple[Path, BytesIO]]:
"""Buffer in Minio fixture."""
# prepare arguments
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
minio_object_name = str(os.getenv('MINIO_OBJECT_NAME'))
# prepare for saving
num_bytes = len(buffer.getvalue())
buffer.seek(0)
# put data in bucket
raw_minio_client.put_object(
bucket_name=minio_bucket_name,
object_name=minio_object_name,
length=num_bytes,
data=buffer,
)
# expose data
yield Path(minio_object_name), buffer
# cleanup
raw_minio_client.remove_object(
bucket_name=minio_bucket_name,
object_name=minio_object_name,
)
@pytest.fixture
def image_data() -> Iterator[ImageData]:
"""Image data fixture."""
# prepare arguments
name = str(os.getenv('MINIO_IMAGE_NAME'))
image = Image.new(mode='RGB', size=(480, 480))
# instantiate data
image_data = ImageData(image=image, name=name)
# expose data
yield image_data
@pytest.fixture
def image_data_in_minio(
raw_minio_client: minio.Minio,
image_data: ImageData,
) -> Iterator[ImageData]:
"""Image data in Minio fixture."""
# prepare arguments
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
image_name = image_data.name
# build object path
object_path = ImageRepository._build_path(image_name)
# save image to buffer
buffer = BytesIO()
image_data.image.save(buffer, 'png')
# prepare for saving
num_bytes = len(buffer.getvalue())
buffer.seek(0)
# put data in bucket
raw_minio_client.put_object(
bucket_name=minio_bucket_name,
object_name=object_path.as_posix(),
length=num_bytes,
data=buffer,
)
# expose data
yield image_data
# cleanup
raw_minio_client.remove_object(
bucket_name=minio_bucket_name,
object_name=object_path.as_posix(),
)
@pytest.fixture(scope='session')
def image_repo(
setup_infrastructure,
populate_env,
) -> Iterator[ImageRepository]:
"""Image repository fixture."""
repo = ImageRepository()
repo.connect()
yield repo
repo.close()
@pytest.fixture
def model_data() -> Iterator[ModelData]:
"""Model data fixture."""
# prepare arguments
vis_com_model = VisualCommunicationModel().to('cpu')
class_name = type(vis_com_model).__name__
buffer = ModelData.model_to_buffer(vis_com_model)
buffer_checksum = HexadecimalString('77dcab1769563654a6e24f92d40f29bd')
# instantiate data
model_data = ModelData(
buffer=buffer,
buffer_checksum=buffer_checksum,
class_name=class_name,
)
# expose model
yield model_data
@pytest.fixture
def model_data_in_minio(
raw_minio_client: minio.Minio,
model_data: ModelData,
) -> Iterator[ModelData]:
"""Model data in Minio fixture."""
# prepare arguments
minio_bucket_name = str(os.getenv('MINIO_BUCKET_NAME'))
object_name = ModelRepository._build_object_name(model_data)
buffer = model_data.buffer
# build object path
object_path = ModelRepository._prefix() / object_name
# prepare for saving
num_bytes = len(buffer.getvalue())
buffer.seek(0)
# put data in bucket
raw_minio_client.put_object(
bucket_name=minio_bucket_name,
object_name=object_path.as_posix(),
length=num_bytes,
data=buffer,
)
# expose data
yield model_data
# cleanup
raw_minio_client.remove_object(
bucket_name=minio_bucket_name,
object_name=object_path.as_posix(),
)
@pytest.fixture(scope='session')
def model_repo(
setup_infrastructure,
populate_env,
) -> Iterator[ModelRepository]:
"""Model repository fixture."""
repo = ModelRepository()
repo.connect()
yield repo
repo.close()
@pytest.fixture
def raw_mongo_client(
setup_infrastructure,
populate_env,
) -> Iterator[MongoClient]:
"""Raw mongo client fixture."""
# prepare arguments
mongo_endpoint = str(os.getenv('MONGO_ENDPOINT'))
mongo_database = str(os.getenv('MONGO_DB'))
# connect client
client = MongoClient(mongo_endpoint)
_ = client[mongo_database]
# expose client
yield client
# cleanup
client.drop_database(mongo_database)
@pytest.fixture
def mongo_client(
setup_infrastructure,
populate_env,
) -> Iterator[MongoImplementation]:
"""MongoImplementation fixture."""
# instantiate and connect client
mongo_client = MongoImplementation()
mongo_client.connect()
# expose client
yield mongo_client
# cleanup
mongo_client._collection.drop()
mongo_client.close()
@pytest.fixture
def dictionary() -> Iterator[dict]:
"""Dictionary fixture."""
# prepare data
data = {
'name': 'test-dictionary',
'str_key': 'value',
'int_key': 100,
'float_key': 3.14,
'list_key': [1, 2, 3],
'dict_key': {
'nested_key': 'nested_value',
},
}
# expose data
yield data
@pytest.fixture
def dictionary_in_mongo(
raw_mongo_client: MongoClient,
dictionary: dict,
) -> Iterator[dict]:
"""Dictionary in Mongo fixture."""
# prepare arguments
database = str(os.getenv('MONGO_DB'))
collection = str(os.getenv('MONGO_COLLECTION'))
# save data
_ = raw_mongo_client[database][collection].insert_one(dictionary.copy())
# expose data
yield dictionary
# cleanup
raw_mongo_client[database][collection].delete_one(dictionary)
@pytest.fixture
def visual_communication_values() -> Iterator[VisualCommunicationValues]:
"""Visual communication values fixture."""
# instantiate with random values
visual_communication_values = VisualCommunicationValues.from_random()
# expose values
yield visual_communication_values
@pytest.fixture
def visual_communication_data(
visual_communication_values: VisualCommunicationValues,
) -> Iterator[VisualCommunicationData]:
"""Visual communication data fixture."""
# prepare arguments
name = 'test-visual-communication'
annotation = visual_communication_values
# instantiate data
visual_communication_data = VisualCommunicationData(
name=name,
annotation=annotation,
)
# expose data
yield visual_communication_data
@pytest.fixture
def visual_communication_data_in_mongo(
raw_mongo_client: MongoClient,
visual_communication_data: VisualCommunicationData,
) -> Iterator[VisualCommunicationData]:
"""Visual communication data in Mongo fixture."""
# prepare arguments
database = str(os.getenv('MONGO_DB'))
collection = str(os.getenv('MONGO_COLLECTION'))
# convert data
dictionary = visual_communication_data.model_dump(mode='dict')
# save data
_ = raw_mongo_client[database][collection].insert_one(dictionary.copy())
# expose data
yield visual_communication_data
# cleanup
raw_mongo_client[database][collection].delete_one(dictionary)
@pytest.fixture(scope='session')
def visual_communication_repo(
setup_infrastructure,
populate_env,
) -> Iterator[VisualCommunicationRepository]:
"""Visual communication repository fixture."""
repo = VisualCommunicationRepository()
repo.connect()
yield repo
repo.close()
@@ -0,0 +1,156 @@
"""Integration tests for ImageRepository class."""
import numpy as np
import pytest
from PIL import Image
from shared.repositories import ImageData, ImageRepository
def same_image(
img_a: Image.Image,
img_b: Image.Image,
) -> bool:
"""Check if two images contain the same data."""
assert isinstance(img_a, Image.Image)
assert isinstance(img_b, Image.Image)
# check if images have a comparable number of channels
if img_a.getbands() != img_b.getbands():
return False
# calculate pixel difference between images
img_a_arr = np.asarray(img_a)
img_b_arr = np.asarray(img_b)
diff = np.subtract(img_a_arr, img_b_arr)
if np.sum(diff) != 0:
return False
return True
def same_image_data(
data_a: ImageData,
data_b: ImageData,
) -> bool:
"""Check if two ImageData-objects contain the same data."""
assert isinstance(data_a, ImageData)
assert isinstance(data_b, ImageData)
# compare names
if data_a.name != data_b.name:
return False
# compare images
if not same_image(data_a.image, data_b.image):
return False
return True
def test_should_have_context_handler():
"""Test that class has a working context handler implemented."""
# ACT
with ImageRepository() as repo:
# ASSERT
assert repo.connected()
def test_should_get_image_data(
image_repo: ImageRepository,
image_data_in_minio: ImageData,
):
"""Test getting image data."""
# ARRANGE
image_name = image_data_in_minio.name
# ACT
received_image_data = image_repo.get_data(image_name)
# ASSERT
assert received_image_data is not None
assert isinstance(received_image_data, ImageData)
assert same_image_data(image_data_in_minio, received_image_data)
def test_should_get_none_when_no_image_data(
image_repo: ImageRepository,
image_data: ImageData,
):
"""Test getting None when no data is available."""
# ARRANGE
image_name = image_data.name
# ACT
received_image_data = image_repo.get_data(image_name)
# ASSERT
assert received_image_data is None
def test_should_delete_image_data(
image_repo: ImageRepository,
image_data_in_minio: ImageData,
):
"""Test deleting image data."""
# ARRANGE
image_name = image_data_in_minio.name
# ACT
image_repo.remove_data(image_name)
received_image_data = image_repo.get_data(image_name)
# ASSERT
assert received_image_data is None
def test_should_put_image_data(
image_repo: ImageRepository,
image_data: ImageData,
):
"""Test putting image data."""
# ARRANGE
image_name = image_data.name
# ACT
image_repo.put_data(image_data)
received_image_data = image_repo.get_data(image_name)
# ASSERT
assert received_image_data is not None
assert same_image_data(image_data, received_image_data)
def test_should_update_image_data(
image_repo: ImageRepository,
image_data_in_minio: ImageData,
):
"""Test updating image data."""
# ARRANGE
updated_image_data = image_data_in_minio.model_copy(
update={
'image': Image.new(mode='RGB', size=(480, 480), color='white'),
},
)
image_name = updated_image_data.name
# ACT
image_repo.put_data(updated_image_data)
received_image_data = image_repo.get_data(image_name)
# ASSERT
assert not same_image_data(image_data_in_minio, updated_image_data)
assert received_image_data is not None
assert same_image_data(updated_image_data, received_image_data)
def test_should_list_names(
image_repo: ImageRepository,
image_data_in_minio: ImageData,
):
"""Test get all image names."""
# ARRANGE
updated_image_data = image_data_in_minio.model_copy(
update={
'name': 'updated-test-image',
},
)
image_repo.put_data(updated_image_data)
expected_name_list = [
image_data_in_minio.name,
updated_image_data.name,
]
# ACT
name_list = image_repo.list_names()
# ASSERT
assert len(name_list) == 2
for name in name_list:
assert name in expected_name_list
if __name__ == '__main__':
pytest.main(['-s', '-v', __file__])
@@ -0,0 +1,133 @@
"""Integration tests related to Minio implementation."""
import logging
import os
from io import BytesIO
from pathlib import Path
import pytest
from shared.repositories.src.implementations import MinioImplementation
def same_data(
data_a: BytesIO,
data_b: BytesIO,
) -> bool:
"""Check if two BytesIO-objects contain the same data."""
assert isinstance(data_a, BytesIO)
assert isinstance(data_b, BytesIO)
# prepare for being read
data_a.seek(0)
data_b.seek(0)
# convert to bytes
data_a_bytes = data_a.read()
data_b_bytes = data_b.read()
# compare size
if len(data_a_bytes) != len(data_b_bytes):
logging.error(
'data has different length: %s and %s',
len(data_a_bytes),
len(data_b_bytes),
)
return False
# compare content
if data_a_bytes != data_b_bytes:
logging.error('data has different bytes')
return False
return True
def test_should_connect_to_minio():
"""Test connection to Minio."""
# ARRANGE
client = MinioImplementation()
# ACT
client.connect()
# ASSERT
assert client.connected()
client.close()
def test_should_have_context_handler():
"""Test that class has a working context handler implemented."""
# ACT
with MinioImplementation() as client:
# ASSERT
assert client.connected()
def test_should_get_data(
minio_client: MinioImplementation,
buffer_in_minio: tuple[Path, BytesIO],
):
"""Test getting data from Minio."""
# ARRANGE
path, buffer = buffer_in_minio
# ACT
received_buffer = minio_client._get(path)
# ASSERT
assert received_buffer is not None
assert same_data(received_buffer, buffer)
def test_should_get_none_when_no_data(
minio_client: MinioImplementation,
):
"""Test getting None when no data is available in Minio."""
# ARRANGE
nonexistent_path = Path('nonexistent-object-name')
# ACT
received_buffer = minio_client._get(nonexistent_path)
# ASSERT
assert received_buffer is None
def test_should_delete_data(
minio_client: MinioImplementation,
buffer_in_minio: tuple[Path, BytesIO],
):
"""Test deleting data from Minio."""
# ARRANGE
path, _ = buffer_in_minio
# ACT
minio_client._delete(path)
# ASSERT
received_buffer = minio_client._get(path)
assert received_buffer is None
def test_should_put_data(
minio_client: MinioImplementation,
buffer: BytesIO,
):
"""Test putting data in Minio."""
# ARRANGE
path = Path(os.getenv('MINIO_OBJECT_NAME', default=''))
# ACT
minio_client._put(path, buffer)
received_buffer = minio_client._get(path)
# ASSERT
assert received_buffer is not None
assert same_data(received_buffer, buffer)
def test_should_update_data(
minio_client: MinioImplementation,
buffer_in_minio: tuple[Path, BytesIO],
):
"""Test updating data in Minio."""
# ARRANGE
path, buffer = buffer_in_minio
updated_buffer = BytesIO(buffer.getvalue() + b'extra data')
# ACT
minio_client._put(path, updated_buffer)
received_buffer = minio_client._get(path)
# ASSERT
assert not same_data(updated_buffer, buffer)
assert received_buffer is not None
assert same_data(received_buffer, updated_buffer)
if __name__ == '__main__':
pytest.main(['-s', '-v', __file__])
@@ -0,0 +1,136 @@
"""Integration tests for ModelRepository class."""
from io import BytesIO
import pytest
from shared.repositories import ModelData, ModelRepository
def same_buffer(
buffer_a: BytesIO,
buffer_b: BytesIO,
) -> bool:
"""Check if 2 buffers contain the same data."""
assert isinstance(buffer_a, BytesIO)
assert isinstance(buffer_b, BytesIO)
# read buffers
a_values = buffer_a.getvalue()
b_values = buffer_b.getvalue()
# compare length of buffers
if len(a_values) != len(b_values):
return False
# compare content of buffers
if a_values != b_values:
return False
return True
def same_model_data(
data_a: ModelData,
data_b: ModelData,
) -> bool:
"""Check if to ModelData-objects contain the same data."""
assert isinstance(data_a, ModelData)
assert isinstance(data_b, ModelData)
# compare names
if data_a.buffer_checksum != data_b.buffer_checksum:
return False
# compare buffer
if not same_buffer(data_a.buffer, data_b.buffer):
return False
return True
def test_should_have_context_handler():
"""Test that class has a working context handler implemented."""
# ACT
with ModelRepository() as repo:
# ASSERT
assert repo.connected()
def test_should_get_model_data(
model_repo: ModelRepository,
model_data_in_minio: ModelData,
):
"""Test getting model data."""
# ARRANGE
object_name = ModelRepository._build_object_name(model_data_in_minio)
# ACT
received_model_data = model_repo.get_data(object_name)
# ASSERT
assert received_model_data is not None
assert isinstance(received_model_data, ModelData)
assert same_model_data(model_data_in_minio, received_model_data)
def test_should_get_none_when_no_model_data(
model_repo: ModelRepository,
model_data: ModelData,
):
"""Test getting None whne no data is available."""
# ARRANGE
object_name = ModelRepository._build_object_name(model_data)
# ACT
received_model_data = model_repo.get_data(object_name)
# ASSERT
assert received_model_data is None
def test_should_delete_model_data(
model_repo: ModelRepository,
model_data_in_minio: ModelData,
):
"""Test deleting model data."""
# ARRANGE
object_name = ModelRepository._build_object_name(model_data_in_minio)
# ACT
model_repo.remove_data(object_name)
received_model_data = model_repo.get_data(object_name)
# ASSERT
assert received_model_data is None
def test_should_put_model_data(
model_repo: ModelRepository,
model_data: ModelData,
):
"""Test putting model data."""
# ARRANGE
object_name = ModelRepository._build_object_name(model_data)
# ACT
model_repo.put_data(model_data)
received_model_data = model_repo.get_data(object_name)
# ASSERT
assert received_model_data is not None
assert same_model_data(model_data, received_model_data)
def test_should_list_names(
model_repo: ModelRepository,
model_data_in_minio: ModelData,
):
"""Test get all model names."""
# ARRANGE
new_buffer_checksum = ModelData.calculate_checksum(model_data_in_minio.buffer)
updated_model_data = model_data_in_minio.model_copy(
update={
'buffer_checksum': new_buffer_checksum,
},
)
model_repo.put_data(updated_model_data)
expected_name_list = [
ModelRepository._build_object_name(model_data_in_minio),
ModelRepository._build_object_name(updated_model_data),
]
# ACT
name_list = model_repo.list_names()
# ASSERT
assert len(name_list) == 2
for name in name_list:
assert name in expected_name_list
if __name__ == '__main__':
pytest.main(['-s', '-v', __file__])
@@ -0,0 +1,113 @@
"""Integration tests related to Mongo implementation."""
import pytest
from shared.repositories.src.implementations import MongoImplementation
def test_should_connect_mongo():
"""Test connecting to Mongo."""
# ARRANGE
client = MongoImplementation()
# ACT
client.connect()
# ASSERT
assert client.connected()
client.close()
def test_should_have_context_handler():
"""Test that class has a working context handler implemented."""
# ACT
with MongoImplementation() as client:
# ASSERT
assert client.connected()
def test_should_get_data(
mongo_client: MongoImplementation,
dictionary_in_mongo: dict,
):
"""Test getting data from mongo."""
# ARRANGE
name = dictionary_in_mongo['name']
query = {'name': name}
# ACT
received_dictionary = mongo_client._get(query)
# ASSERT
assert received_dictionary is not None
assert isinstance(received_dictionary, dict)
assert received_dictionary == dictionary_in_mongo
def test_should_get_none_when_no_data(
mongo_client: MongoImplementation,
dictionary: dict,
):
"""Test getting None when no data is available in mongo."""
# ARRANGE
name = dictionary['name']
query = {'name': name}
# ACT
received_dictionary = mongo_client._get(query)
# ASSERT
assert received_dictionary is None
def test_should_delete_data(
mongo_client: MongoImplementation,
dictionary_in_mongo: dict,
):
"""Test deleting data from mongo."""
# ARRANGE
name = dictionary_in_mongo['name']
query = {'name': name}
# ACT
mongo_client._delete(query)
received_dictionary = mongo_client._get(query)
# ASSERT
assert received_dictionary is None
def test_should_save_data(
mongo_client: MongoImplementation,
dictionary: dict,
):
"""Test saving data to mongo."""
# ARRANGE
name = dictionary['name']
query = {'name': name}
# ACT
mongo_client._save(
data=dictionary,
query=query,
)
# ASSERT
received_dictionary = mongo_client._get(query)
assert received_dictionary is not None
assert isinstance(received_dictionary, dict)
assert received_dictionary == dictionary
def test_should_update_data(
mongo_client: MongoImplementation,
dictionary_in_mongo: dict,
):
"""Test updating data in mongo."""
# ARRANGE
name = dictionary_in_mongo['name']
query = {'name': name}
# ACT
dictionary_in_mongo['str-key'] = 'updated-value'
mongo_client._save(
data=dictionary_in_mongo,
query=query,
)
# ASSERT
received_dictionary = mongo_client._get(query)
assert received_dictionary is not None
assert received_dictionary == dictionary_in_mongo
if __name__ == '__main__':
pytest.main(['-s', '-v', __file__])
@@ -0,0 +1,125 @@
"""Integration tests for VisualCommunicationRepository class."""
import pytest
from shared.repositories import (
VisualCommunicationData,
VisualCommunicationRepository,
VisualCommunicationValues,
)
def test_should_have_context_handler():
"""Test that class has a working context handler implemented."""
# ACT
with VisualCommunicationRepository() as repo:
# ASSERT
assert repo.connected()
def test_should_get_visual_communication_data(
visual_communication_repo: VisualCommunicationRepository,
visual_communication_data_in_mongo: VisualCommunicationData,
):
"""Test getting visual communication data."""
# ARRANGE
name = visual_communication_data_in_mongo.name
# ACT
received_visual_communication_data = visual_communication_repo.get_data(name)
# ASSERT
assert received_visual_communication_data is not None
assert isinstance(received_visual_communication_data, VisualCommunicationData)
assert received_visual_communication_data == visual_communication_data_in_mongo
def test_should_get_none_when_no_visual_communication_data(
visual_communication_repo: VisualCommunicationRepository,
visual_communication_data: VisualCommunicationData,
):
"""Test getting None when no data is available."""
# ARRANGE
name = visual_communication_data.name
# ACT
received_visual_communication_data = visual_communication_repo.get_data(name)
# ASSERT
assert received_visual_communication_data is None
def test_should_delete_image_data(
visual_communication_repo: VisualCommunicationRepository,
visual_communication_data_in_mongo: VisualCommunicationData,
):
"""Test deleting visual communication data."""
# ARRANGE
name = visual_communication_data_in_mongo.name
# ACT
visual_communication_repo.remove_data(name)
received_visual_communication_data = visual_communication_repo.get_data(name)
# ASSERT
assert received_visual_communication_data is None
def test_should_put_visual_communication_data(
visual_communication_repo: VisualCommunicationRepository,
visual_communication_data: VisualCommunicationData,
):
"""Test putting visual communication data."""
# ARRANGE
name = visual_communication_data.name
# ACT
visual_communication_repo.put_data(visual_communication_data)
received_visual_communication_data = visual_communication_repo.get_data(name)
# ASSERT
assert received_visual_communication_data is not None
assert isinstance(received_visual_communication_data, VisualCommunicationData)
assert received_visual_communication_data == visual_communication_data
def test_should_update_visual_communication_data(
visual_communication_repo: VisualCommunicationRepository,
visual_communication_data_in_mongo: VisualCommunicationData,
):
"""Test updating visual communication data."""
# ARRANGE
name = visual_communication_data_in_mongo.name
updated_annotation = VisualCommunicationValues.from_random()
updated_visual_communication = VisualCommunicationData(
name=name,
annotation=updated_annotation,
)
# ACT
visual_communication_repo.put_data(updated_visual_communication)
received_visual_communication_data = visual_communication_repo.get_data(name)
# ASSERT
assert received_visual_communication_data is not None
assert isinstance(received_visual_communication_data, VisualCommunicationData)
assert received_visual_communication_data == updated_visual_communication
def test_should_list_names(
visual_communication_repo: VisualCommunicationRepository,
visual_communication_data_in_mongo: VisualCommunicationData,
):
"""Test listing names of all documents."""
# ARRANGE
name = 'test-visual-communication-2'
annotation = VisualCommunicationValues.from_random()
second_visual_communication = VisualCommunicationData(
name=name,
annotation=annotation,
)
visual_communication_repo.put_data(second_visual_communication)
expected_name_list = [
visual_communication_data_in_mongo.name,
second_visual_communication.name,
]
# ACT
name_list = visual_communication_repo.list_names()
# ASSERT
assert len(name_list) == 2
for name in name_list:
assert name in expected_name_list
if __name__ == '__main__':
pytest.main(['-s', '-v', __file__])