def test_update_document(self): mock_mongo_client = MagicMock() with patch('qilib.data_set.mongo_data_set_io.MongoClient', return_value={'qilib': { 'data_sets': mock_mongo_client }}): mongo_data_set_io = MongoDataSetIO(name='test_data_set') mongo_data_set_io.update_document( {'array_updates': ('(2,2)', { 'test': 5 })}) mock_mongo_client.update_one.called_once_with( {'name': 'test_data_set'}, { '$set': { 'array_updates': ('(2,2)', { 'test': 5 }) }, "$currentDate": { "lastModified": True } })
class MongoDataSetIOWriter(DataSetIOWriter): """ Allow a DataSet to store changes, and complete DataSet, to a mongodb.""" def __init__( self, name: Optional[str] = None, document_id: Optional[str] = None, database: str = MongoDataSetIO.DEFAULT_DATABASE_NAME, collection: str = MongoDataSetIO.DEFAULT_COLLECTION_NAME) -> None: """ Construct a new instance of MongoDataSetIOWriter. If name is provided, but not found in the database a new document is created with that name. Args: name: DataSet name. document_id: _id of the DataSet in the database. database: Name of the database. collection: Name of the collections. Raises: DocumentNotFoundError: If document_id is provided but not found in the database. """ super().__init__() self._mongo_data_set_io = MongoDataSetIO(name, document_id, database=database, collection=collection) def sync_metadata_to_storage(self, field_name: str, value: Any) -> None: """ Update or add metadata field to database. Args: field_name: Field that changed. value: The new value. """ self._is_finalized() update_data = { "{}.{}".format(DataSetIOReader.METADATA, field_name): value } self._mongo_data_set_io.update_document(update_data) def sync_data_to_storage(self, index_or_slice: Union[Tuple[int, ...], int], data: Dict[str, Any]) -> None: """ Registers a DataArray update to the database. The change is registered as a change event and is applied on finalize(). Args: index_or_slice: The indices of the DataArray to update. data: Name of the DataArray to be updated and the new value. """ self._is_finalized() update_data = {DataSetIOReader.ARRAY_UPDATES: (index_or_slice, data)} self._mongo_data_set_io.append_to_document(update_data) def sync_add_data_array_to_storage(self, data_array: DataArray) -> None: """ Add or update a DataArray in the database. Args: data_array: The DataArray to be updated or added. """ self._is_finalized() update_data = { "{}.{}".format(DataSetIOReader.DATA_ARRAYS, data_array.name): { "name": data_array.name, "label": data_array.label, "unit": data_array.unit, "is_setpoint": data_array.is_setpoint, "set_arrays": [array.name for array in data_array.set_arrays], "preset_data": MongoDataSetIO.encode_numpy_array(data_array) } } self._mongo_data_set_io.update_document(update_data) def finalize(self) -> None: """ Update the underlying DataSet and close the connection to the database.""" self._mongo_data_set_io.update_document( {DataSetIOReader.ARRAY_UPDATES: []}) self._mongo_data_set_io.finalize() self._finalized = True