diff --git a/packages/syft/src/syft/store/kv_document_store.py b/packages/syft/src/syft/store/kv_document_store.py index 0b86685ea66..4c9ad60b122 100644 --- a/packages/syft/src/syft/store/kv_document_store.py +++ b/packages/syft/src/syft/store/kv_document_store.py @@ -470,13 +470,17 @@ def _delete( def _delete_unique_keys_for(self, obj: SyftObject) -> Result[SyftSuccess, str]: for _unique_ck in self.unique_cks: qk = _unique_ck.with_obj(obj) - self.unique_keys[qk.key].pop(qk.value, None) + unique_keys = self.unique_keys[qk.key] + unique_keys.pop(qk.value, None) + self.unique_keys[qk.key] = unique_keys return Ok(SyftSuccess(message="Deleted")) def _delete_search_keys_for(self, obj: SyftObject) -> Result[SyftSuccess, str]: for _search_ck in self.searchable_cks: qk = _search_ck.with_obj(obj) - self.searchable_keys[qk.key].pop(qk.value, None) + search_keys = self.searchable_keys[qk.key] + search_keys.pop(qk.value, None) + self.searchable_keys[qk.key] = search_keys return Ok(SyftSuccess(message="Deleted")) def _get_keys_index(self, qks: QueryKeys) -> Result[Set[Any], str]: diff --git a/packages/syft/tests/syft/stores/kv_document_store_test.py b/packages/syft/tests/syft/stores/kv_document_store_test.py index 3e66d753a96..2ef0c5794ae 100644 --- a/packages/syft/tests/syft/stores/kv_document_store_test.py +++ b/packages/syft/tests/syft/stores/kv_document_store_test.py @@ -125,6 +125,29 @@ def test_kv_store_partition_delete( assert len(kv_store_partition.all(root_verify_key).ok()) == 0 +def test_kv_store_partition_delete_and_recreate( + root_verify_key, worker, kv_store_partition: KeyValueStorePartition +) -> None: + obj = MockSyftObject(data="bogus") + for _ in range(2): + # running it multiple items ensures we can recreate it again once its delete from store. + + # Add an object + kv_store_partition.set(root_verify_key, obj, ignore_duplicates=False) + + assert len(kv_store_partition.all(root_verify_key).ok()) == 1 + + # Delete object + key = kv_store_partition.settings.store_key.with_obj(obj) + res = kv_store_partition.delete(root_verify_key, key) + + assert res.is_ok() + assert len(kv_store_partition.all(root_verify_key).ok()) == 0 + assert len(kv_store_partition.data) == len(kv_store_partition.permissions) + + assert len(kv_store_partition.all(root_verify_key).ok()) == 0 + + def test_kv_store_partition_update( root_verify_key, kv_store_partition: KeyValueStorePartition ) -> None: