diff --git a/packages/evo-blockmodels/src/evo/blockmodels/client.py b/packages/evo-blockmodels/src/evo/blockmodels/client.py index 92569984..f963cfb9 100644 --- a/packages/evo-blockmodels/src/evo/blockmodels/client.py +++ b/packages/evo-blockmodels/src/evo/blockmodels/client.py @@ -791,6 +791,7 @@ async def _update_columns( geometry_change: bool | None = None, fill_subblocks: bool | None = None, tags: dict[str, dict[str, Any]] | None = None, + update_type: models.UpdateType = models.UpdateType.replace ) -> Version: if self._cache is None: raise CacheNotConfiguredException( @@ -850,9 +851,9 @@ async def _update_columns( update_data_lite=models.UpdateDataLite( models.UpdateDataLite1( columns=columns, - update_type=models.UpdateType.replace, + update_type=update_type, geometry_change=geometry_change, - **({} if fill_subblocks is None else {"fill_subblocks": fill_subblocks}), + fill_subblocks=fill_subblocks, ) ), additional_headers=self._preview_headers(), @@ -868,6 +869,7 @@ async def update_block_model_columns( delete_columns: set[str] | None = None, units: dict[str, str] | None = None, tags: dict[str, dict[str, Any]] | None = None, + update_type: models.UpdateType = models.UpdateType.replace ) -> Version: """Add, update, or delete regular block model columns. @@ -883,11 +885,12 @@ async def update_block_model_columns( :param units: A dictionary mapping column names within `data` to units. :param tags: A dictionary mapping new column names to their tags object. Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them. + :param: update_type: Provide the type of update. Either 'replace' or 'merge' (default: replace) :raises CacheNotConfiguredException: If the cache is not configured. :return: The new version of the block model with the added columns. """ return await self._update_columns( - bm_id, data, new_columns, update_columns, delete_columns, units, geometry_change=None, tags=tags + bm_id, data, new_columns, update_columns, delete_columns, units, geometry_change=None, tags=tags, update_type=update_type, ) async def update_subblocked_columns( @@ -901,6 +904,7 @@ async def update_subblocked_columns( geometry_change: bool = False, fill_subblocks: bool | None = None, tags: dict[str, dict[str, Any]] | None = None, + update_type: models.UpdateType = models.UpdateType.replace ) -> Version: """Add, update, or delete sub-blocked block model columns. @@ -925,6 +929,7 @@ async def update_subblocked_columns( the block model's own ``fill_subblocks`` setting is used. :param tags: A dictionary mapping new column names to their tags object. Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them. + :param: update_type: Provide the type of update. Either 'replace' or 'merge' (default: replace) """ return await self._update_columns( bm_id, @@ -936,6 +941,7 @@ async def update_subblocked_columns( geometry_change=geometry_change, fill_subblocks=fill_subblocks, tags=tags, + update_type=update_type ) async def update_column_metadata( diff --git a/packages/evo-blockmodels/src/evo/blockmodels/endpoints/models.py b/packages/evo-blockmodels/src/evo/blockmodels/endpoints/models.py index 8ef4fc81..515c186a 100644 --- a/packages/evo-blockmodels/src/evo/blockmodels/endpoints/models.py +++ b/packages/evo-blockmodels/src/evo/blockmodels/endpoints/models.py @@ -3230,7 +3230,7 @@ class UpdateDataLite1(CustomBaseModel): """ Lineage of the block model update """ - update_type: UpdateType = "merge" + update_type: UpdateType = UpdateType.merge """ Behaviour of the update, for blocks that are omitted from the update file. @@ -3283,7 +3283,7 @@ class UpdateDataLite2(CustomBaseModel): """ Lineage of the block model update """ - update_type: UpdateType = "merge" + update_type: UpdateType = UpdateType.merge """ Behaviour of the update, for blocks that are omitted from the update file. diff --git a/packages/evo-blockmodels/tests/test_create.py b/packages/evo-blockmodels/tests/test_create.py index 45e75fea..a30a1242 100644 --- a/packages/evo-blockmodels/tests/test_create.py +++ b/packages/evo-blockmodels/tests/test_create.py @@ -31,7 +31,7 @@ from evo.common.data import HTTPHeaderDict, RequestMethod from evo.common.test_tools import BASE_URL, MockResponse, TestWithConnector, TestWithStorage from evo.common.utils import get_header_metadata -from utils import JobPollingRequestHandler +from utils import DEFAULT_EXPECTED_HEADERS, JobPollingRequestHandler BM_UUID = uuid.uuid4() GOOSE_UUID = uuid.uuid4() @@ -276,11 +276,7 @@ def _assert_create_request( comment=comment, fill_subblocks=fill_subblocks, ).model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) async def test_create_block_model(self) -> None: @@ -457,11 +453,7 @@ async def test_create_block_model_with_data(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(bm.id, BM_UUID) @@ -559,11 +551,7 @@ async def test_create_subblocked_model_with_data(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(bm.id, BM_UUID) diff --git a/packages/evo-blockmodels/tests/test_update.py b/packages/evo-blockmodels/tests/test_update.py index 8a7b34a6..364d75e4 100644 --- a/packages/evo-blockmodels/tests/test_update.py +++ b/packages/evo-blockmodels/tests/test_update.py @@ -27,7 +27,7 @@ from evo.common.data import HTTPHeaderDict, RequestMethod from evo.common.test_tools import BASE_URL, MockResponse, TestWithConnector, TestWithStorage from evo.common.utils import get_header_metadata -from utils import JobPollingRequestHandler +from utils import DEFAULT_EXPECTED_HEADERS, JobPollingRequestHandler BM_UUID = uuid.uuid4() GOOSE_UUID = uuid.uuid4() @@ -36,7 +36,6 @@ MODEL_USER = models.IMSUserInfo(email="test@test.com", name="Test User", id=uuid.uuid4()) USER = ServiceUser.from_model(MODEL_USER) - def _mock_version( version_id: int, version_uuid: uuid.UUID, goose_version_id: str, bbox=None, columns: Iterable[models.Column] = () ) -> models.Version: @@ -194,11 +193,7 @@ async def test_add_new_columns(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(version.bm_uuid, BM_UUID) self.assertEqual(version.version_id, 2) @@ -261,11 +256,7 @@ async def test_add_new_columns_with_tags(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) async def test_add_new_columns_with_unknown_tag_column(self) -> None: @@ -331,11 +322,7 @@ async def test_add_new_subblocked_columns(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(version.bm_uuid, BM_UUID) self.assertEqual(version.version_id, 2) @@ -422,16 +409,13 @@ async def test_update_columns(self) -> None: ), update_type=models.UpdateType.replace, geometry_change=None, + fill_subblocks=None, ) self.assert_any_request_made( method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(version.bm_uuid, BM_UUID) self.assertEqual(version.version_id, 2) @@ -509,11 +493,11 @@ async def test_update_columns_job_failed(self) -> None: @parameterized.expand( [ - (True,), - (False,), + (True, True), + (False,False), ] ) - async def test_update_subblocked_columns(self, geometry_change: bool) -> None: + async def test_update_subblocked_columns(self, geometry_change: bool, fill_subblocks: bool) -> None: self.transport.set_request_handler( UpdateRequestHandler( update_result=UPDATE_RESULT, @@ -535,6 +519,7 @@ async def test_update_subblocked_columns(self, geometry_change: bool) -> None: delete_columns={"col3"}, units={"col2": "g/t"}, geometry_change=geometry_change, + fill_subblocks=fill_subblocks ) mock_destination.upload_file.assert_called_once() @@ -554,16 +539,13 @@ async def test_update_subblocked_columns(self, geometry_change: bool) -> None: ), update_type=models.UpdateType.replace, geometry_change=geometry_change, + fill_subblocks=fill_subblocks, ) self.assert_any_request_made( method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(version.bm_uuid, BM_UUID) self.assertEqual(version.version_id, 2) @@ -625,15 +607,113 @@ async def test_update_subblocked_columns_with_fill_subblocks(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(version.bm_uuid, BM_UUID) self.assertEqual(version.version_id, 2) + async def test_update_block_model_columns_merge(self) -> None: + """update_type=merge is passed through the stack to UpdateDataLite for regular models.""" + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse( + job_status=JobStatus.COMPLETE, + payload=UPDATED_VERSION, + ), + ) + ) + with ( + mock.patch("evo.common.io.upload.StorageDestination") as mock_destination, + ): + mock_destination.upload_file = mock.AsyncMock() + await self.bms_client.update_block_model_columns( + BM_UUID, + REGULAR_DATA, + new_columns=["col2"], + update_columns={"col1"}, + delete_columns={"col3"}, + units={"col2": "g/t"}, + update_type=models.UpdateType.merge, + ) + mock_destination.upload_file.assert_called_once() + + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite( + new=[ + models.ColumnLite( + title="col2", + data_type=models.DataType.Float64, + unit_id="g/t", + ), + ], + update=["col1"], + rename=[], + delete=["col3"], + ), + update_type=models.UpdateType.merge, + geometry_change=None, + fill_subblocks=None, + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS, + ) + + async def test_update_subblocked_columns_merge(self) -> None: + """update_type=merge is passed through the stack to UpdateDataLite for sub-blocked models.""" + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse( + job_status=JobStatus.COMPLETE, + payload=UPDATED_VERSION, + ), + ) + ) + with ( + mock.patch("evo.common.io.upload.StorageDestination") as mock_destination, + ): + mock_destination.upload_file = mock.AsyncMock() + await self.bms_client.update_subblocked_columns( + BM_UUID, + SUBBLOCKED_DATA, + new_columns=["col2"], + update_columns={"col1"}, + delete_columns={"col3"}, + units={"col2": "g/t"}, + geometry_change=True, + fill_subblocks=True, + update_type=models.UpdateType.merge, + ) + mock_destination.upload_file.assert_called_once() + + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite( + new=[ + models.ColumnLite( + title="col2", + data_type=models.DataType.Float64, + unit_id="g/t", + ), + ], + update=["col1"], + rename=[], + delete=["col3"], + ), + update_type=models.UpdateType.merge, + geometry_change=True, + fill_subblocks=True, + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS, + ) + async def test_update_column_metadata(self) -> None: self.transport.set_request_handler( UpdateRequestHandler( @@ -724,11 +804,7 @@ async def test_update_column_metadata_with_tags(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) async def test_update_column_metadata_with_comment(self) -> None: @@ -771,11 +847,7 @@ async def test_update_column_metadata_with_comment(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(version.bm_uuid, BM_UUID) @@ -814,12 +886,7 @@ async def test_update_column_metadata_with_preview_sends_header(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - "API-Preview": "opt-in", - }, + headers=DEFAULT_EXPECTED_HEADERS | {"API-Preview": "opt-in"}, ) self.assertEqual(version.version_id, 2) @@ -875,11 +942,7 @@ async def test_rename_block_model_columns_with_comment(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(version.bm_uuid, BM_UUID) self.assertEqual(version.version_id, 2) @@ -952,11 +1015,7 @@ async def test_delete_block_model_columns_with_comment(self) -> None: method=RequestMethod.PATCH, path=f"{self.base_path}/block-models/{BM_UUID}/blocks", body=expected_update_body.model_dump(mode="json", exclude_unset=True), - headers={ - "Authorization": "Bearer ", - "Content-Type": "application/json", - "Accept": "application/json", - }, + headers=DEFAULT_EXPECTED_HEADERS, ) self.assertEqual(version.bm_uuid, BM_UUID) self.assertEqual(version.version_id, 2) diff --git a/packages/evo-blockmodels/tests/utils.py b/packages/evo-blockmodels/tests/utils.py index eca56a12..0e140e33 100644 --- a/packages/evo-blockmodels/tests/utils.py +++ b/packages/evo-blockmodels/tests/utils.py @@ -12,6 +12,11 @@ from evo.blockmodels.endpoints.models import JobResponse, JobStatus from evo.common.test_tools import AbstractTestRequestHandler, MockResponse +DEFAULT_EXPECTED_HEADERS = { + "Authorization": "Bearer ", + "Content-Type": "application/json", + "Accept": "application/json", +} class JobPollingRequestHandler(AbstractTestRequestHandler): def __init__(self, job_response: JobResponse, pending_request: int = 0) -> None: