2 files changed, 59 insertions, 35 deletions

diff --git a/storage/src/tests/persistence/common/persistenceproviderwrapper.cpp b/storage/src/tests/persistence/common/persistenceproviderwrapper.cpp
index 5f1acf6e7da..b3bd1c6a253 100644
--- a/storage/src/tests/persistence/common/persistenceproviderwrapper.cpp
+++ b/storage/src/tests/persistence/common/persistenceproviderwrapper.cpp
@@ -96,46 +96,40 @@ PersistenceProviderWrapper::getBucketInfo(const spi::Bucket& bucket) const
     return _spi.getBucketInfo(bucket);
 }
 
-spi::Result
-PersistenceProviderWrapper::put(const spi::Bucket& bucket, spi::Timestamp timestamp,
-                                document::Document::SP doc, spi::Context& context)
+void
+PersistenceProviderWrapper::putAsync(const spi::Bucket& bucket, spi::Timestamp timestamp, document::Document::SP doc,
+                                     spi::Context& context, spi::OperationComplete::UP onComplete)
 {
     LOG_SPI("put(" << bucket << ", " << timestamp << ", " << doc->getId() << ")");
-    CHECK_ERROR(spi::Result, FAIL_PUT);
-    return _spi.put(bucket, timestamp, std::move(doc), context);
+    CHECK_ERROR_ASYNC(spi::Result, FAIL_PUT, onComplete);
+    _spi.putAsync(bucket, timestamp, std::move(doc), context, std::move(onComplete));
 }
 
-spi::RemoveResult
-PersistenceProviderWrapper::remove(const spi::Bucket& bucket,
-                                   spi::Timestamp timestamp,
-                                   const spi::DocumentId& id,
-                                   spi::Context& context)
+void
+PersistenceProviderWrapper::removeAsync(const spi::Bucket& bucket, spi::Timestamp timestamp, const spi::DocumentId& id,
+                                        spi::Context& context, spi::OperationComplete::UP onComplete)
 {
     LOG_SPI("remove(" << bucket << ", " << timestamp << ", " << id << ")");
-    CHECK_ERROR(spi::RemoveResult, FAIL_REMOVE);
-    return _spi.remove(bucket, timestamp, id, context);
+    CHECK_ERROR_ASYNC(spi::RemoveResult, FAIL_REMOVE, onComplete);
+    _spi.removeAsync(bucket, timestamp, id, context, std::move(onComplete));
 }
 
-spi::RemoveResult
-PersistenceProviderWrapper::removeIfFound(const spi::Bucket& bucket,
-                                          spi::Timestamp timestamp,
-                                          const spi::DocumentId& id,
-                                          spi::Context& context)
+void
+PersistenceProviderWrapper::removeIfFoundAsync(const spi::Bucket& bucket, spi::Timestamp timestamp, const spi::DocumentId& id,
+                                               spi::Context& context, spi::OperationComplete::UP onComplete)
 {
     LOG_SPI("removeIfFound(" << bucket << ", " << timestamp << ", " << id << ")");
-    CHECK_ERROR(spi::RemoveResult, FAIL_REMOVE_IF_FOUND);
-    return _spi.removeIfFound(bucket, timestamp, id, context);
+    CHECK_ERROR_ASYNC(spi::RemoveResult, FAIL_REMOVE_IF_FOUND, onComplete);
+    _spi.removeIfFoundAsync(bucket, timestamp, id, context, std::move(onComplete));
 }
 
-spi::UpdateResult
-PersistenceProviderWrapper::update(const spi::Bucket& bucket,
-                                   spi::Timestamp timestamp,
-                                   document::DocumentUpdate::SP upd,
-                                   spi::Context& context)
+void
+PersistenceProviderWrapper::updateAsync(const spi::Bucket& bucket, spi::Timestamp timestamp, document::DocumentUpdate::SP upd,
+                                        spi::Context& context, spi::OperationComplete::UP onComplete)
 {
     LOG_SPI("update(" << bucket << ", " << timestamp << ", " << upd->getId() << ")");
-    CHECK_ERROR(spi::UpdateResult, FAIL_UPDATE);
-    return _spi.update(bucket, timestamp, std::move(upd), context);
+    CHECK_ERROR_ASYNC(spi::UpdateResult, FAIL_UPDATE, onComplete);
+    _spi.updateAsync(bucket, timestamp, std::move(upd), context, std::move(onComplete));
 }
 
 spi::GetResult
@@ -234,4 +228,26 @@ PersistenceProviderWrapper::removeEntry(const spi::Bucket& bucket,
     return _spi.removeEntry(bucket, timestamp, context);
 }
 
+spi::Result
+PersistenceProviderWrapper::initialize() {
+    return _spi.initialize();
+}
+
+spi::BucketIdListResult
+PersistenceProviderWrapper::getModifiedBuckets(spi::PersistenceProvider::BucketSpace bucketSpace) const {
+    return _spi.getModifiedBuckets(bucketSpace);
+}
+
+spi::Result
+PersistenceProviderWrapper::setClusterState(spi::PersistenceProvider::BucketSpace bucketSpace, const spi::ClusterState &state) {
+    return _spi.setClusterState(bucketSpace, state);
+}
+
+void
+PersistenceProviderWrapper::setActiveStateAsync(const spi::Bucket &bucket, spi::BucketInfo::ActiveState state,
+                                                spi::OperationComplete::UP onComplete)
+{
+    _spi.setActiveStateAsync(bucket, state, std::move(onComplete));
+}
+
 }
diff --git a/storage/src/tests/persistence/common/persistenceproviderwrapper.h b/storage/src/tests/persistence/common/persistenceproviderwrapper.h
index 64828a2a3ee..c6628814dba 100644
--- a/storage/src/tests/persistence/common/persistenceproviderwrapper.h
+++ b/storage/src/tests/persistence/common/persistenceproviderwrapper.h
@@ -21,7 +21,7 @@
 
 namespace storage {
 
-class PersistenceProviderWrapper : public spi::AbstractPersistenceProvider
+class PersistenceProviderWrapper : public spi::PersistenceProvider
 {
 public:
     enum OPERATION_FAILURE_FLAGS
@@ -47,11 +47,11 @@ public:
         // TODO: add more as needed
     };
 private:
-    spi::PersistenceProvider& _spi;
-    spi::Result _result;
-    mutable std::mutex  _lock;
+    spi::PersistenceProvider&        _spi;
+    spi::Result                      _result;
+    mutable std::mutex               _lock;
     mutable std::vector<std::string> _log;
-    uint32_t _failureMask;
+    uint32_t                         _failureMask;
     using Guard = std::lock_guard<std::mutex>;
 public:
     PersistenceProviderWrapper(spi::PersistenceProvider& spi);
@@ -88,13 +88,21 @@ public:
         _log.clear();
     }
 
+    spi::Result initialize() override;
+    spi::BucketIdListResult getModifiedBuckets(BucketSpace bucketSpace) const override;
+
+    spi::Result setClusterState(BucketSpace bucketSpace, const spi::ClusterState &state) override;
+
+    void setActiveStateAsync(const spi::Bucket &bucket, spi::BucketInfo::ActiveState state,
+                             spi::OperationComplete::UP up) override;
+
     spi::Result createBucket(const spi::Bucket&, spi::Context&) override;
     spi::BucketIdListResult listBuckets(BucketSpace bucketSpace) const override;
     spi::BucketInfoResult getBucketInfo(const spi::Bucket&) const override;
-    spi::Result put(const spi::Bucket&, spi::Timestamp, spi::DocumentSP, spi::Context&) override;
-    spi::RemoveResult remove(const spi::Bucket&, spi::Timestamp, const spi::DocumentId&, spi::Context&) override;
-    spi::RemoveResult removeIfFound(const spi::Bucket&, spi::Timestamp, const spi::DocumentId&, spi::Context&) override;
-    spi::UpdateResult update(const spi::Bucket&, spi::Timestamp, spi::DocumentUpdateSP, spi::Context&) override;
+    void putAsync(const spi::Bucket&, spi::Timestamp, spi::DocumentSP, spi::Context&, spi::OperationComplete::UP) override;
+    void removeAsync(const spi::Bucket&, spi::Timestamp, const spi::DocumentId&, spi::Context&, spi::OperationComplete::UP) override;
+    void removeIfFoundAsync(const spi::Bucket&, spi::Timestamp, const spi::DocumentId&, spi::Context&, spi::OperationComplete::UP) override;
+    void updateAsync(const spi::Bucket&, spi::Timestamp, spi::DocumentUpdateSP, spi::Context&, spi::OperationComplete::UP) override;
     spi::GetResult get(const spi::Bucket&, const document::FieldSet&, const spi::DocumentId&, spi::Context&) const override;
 
     spi::CreateIteratorResult