Repository: googleapis/storage-testbench

Issue:
Add support for object softdelete
Need to add the ability to softdelete / restore objects as well as add the appropriate fields (Soft Delete Time, Hard Delete Time).

Discussion:

Solution patch:
diff --git a/testbench/database.py b/testbench/database.py
index f17260d..30cd028 100644
--- a/testbench/database.py
+++ b/testbench/database.py
@@ -14,6 +14,7 @@
 
 import collections
 import copy
+import datetime
 import json
 import os
 import pathlib
@@ -37,11 +38,13 @@ class Database:
         rewrites,
         retry_tests,
         supported_methods,
+        soft_deleted_objects,
     ):
         self._resources_lock = threading.RLock()
         self._buckets = buckets
         self._objects = objects
         self._live_generations = live_generations
+        self._soft_deleted_objects = soft_deleted_objects
 
         self._uploads_lock = threading.RLock()
         self._uploads = uploads
@@ -58,7 +61,7 @@ class Database:
 
     @classmethod
     def init(cls):
-        return cls({}, {}, {}, {}, {}, {}, [])
+        return cls({}, {}, {}, {}, {}, {}, [], {})
 
     def clear(self):
         """Clear all data except for the supported method list."""
@@ -66,6 +69,7 @@ class Database:
             self._buckets = {}
             self._objects = {}
             self._live_generations = {}
+            self._soft_deleted_objects = {}
         with self._uploads_lock:
             self._uploads = {}
         with self._rewrites_lock:
@@ -101,6 +105,7 @@ class Database:
             self._buckets[bucket.metadata.name] = bucket
             self._objects[bucket.metadata.name] = {}
             self._live_generations[bucket.metadata.name] = {}
+            self._soft_deleted_objects[bucket.metadata.name] = {}
 
     def list_bucket(self, project_id, prefix, context):
         with self._resources_lock:
@@ -133,6 +138,7 @@ class Database:
             del self._buckets[bucket.metadata.name]
             del self._objects[bucket.metadata.name]
             del self._live_generations[bucket.metadata.name]
+            del self._soft_deleted_objects[bucket.metadata.name]
 
     def insert_test_bucket(self):
         """Automatically create a bucket if needed.
@@ -173,6 +179,7 @@ class Database:
             request.lexicographic_end,
             request.include_trailing_delimiter,
             request.match_glob,
+            request.soft_deleted,
         )
 
     @classmethod
@@ -186,6 +193,7 @@ class Database:
         end_offset = request.args.get("endOffset")
         include_trailing_delimiter = request.args.get("includeTrailingDelimiter", False)
         match_glob = request.args.get("matchGlob", None)
+        soft_deleted = request.args.get("softDeleted", False)
         return (
             delimiter,
             prefix,
@@ -194,6 +202,7 @@ class Database:
             end_offset,
             include_trailing_delimiter,
             match_glob,
+            soft_deleted,
         )
 
     def __get_live_generation(self, bucket_name, object_name, context):
@@ -208,9 +217,66 @@ class Database:
         bucket_key = self.__bucket_key(bucket_name, context)
         self._live_generations[bucket_key].pop(object_name, None)
 
+    def __soft_delete_object(
+        self, bucket_name, object_name, blob, retention_duration, context
+    ):
+        bucket_key = self.__bucket_key(bucket_name, context)
+        if self._soft_deleted_objects[bucket_key].get(object_name) is None:
+            self._soft_deleted_objects[bucket_key][object_name] = []
+        soft_delete_time = datetime.datetime.now(datetime.timezone.utc)
+        hard_delete_time = soft_delete_time + datetime.timedelta(0, retention_duration)
+        blob.metadata.soft_delete_time.FromDatetime(soft_delete_time)
+        blob.metadata.hard_delete_time.FromDatetime(hard_delete_time)
+        self._soft_deleted_objects[bucket_key][object_name].append(blob)
+
+    def __remove_expired_objects_from_soft_delete(
+        self, bucket_name, object_name, context
+    ):
+        bucket_key = self.__bucket_key(bucket_name, context)
+        now = datetime.datetime.now()
+
+        if self._soft_deleted_objects[bucket_key].get(object_name) is not None:
+            self._soft_deleted_objects[bucket_key][object_name] = list(
+                filter(
+                    lambda blob: now < blob.metadata.hard_delete_time.ToDatetime(),
+                    self._soft_deleted_objects[bucket_key][object_name],
+                )
+            )
+
+    def __remove_restored_soft_deleted_object(
+        self, bucket_name, object_name, generation, context
+    ):
+        bucket_key = self.__bucket_key(bucket_name, context)
+        if self._soft_deleted_objects[bucket_key].get(object_name) is not None:
+            self._soft_deleted_objects[bucket_key][object_name] = list(
+                filter(
+                    lambda blob: blob.metadata.generation == generation,
+                    self._soft_deleted_objects[bucket_key][object_name],
+                )
+            )
+
+    def __get_soft_deleted_object(self, bucket_name, object_name, generation, context):
+        bucket_key = self.__bucket_key(bucket_name, context)
+        blobs = self._soft_deleted_objects[bucket_key].get(object_name, [])
+        blob = next(
+            (blob for blob in blobs if blob.metadata.generation == generation), None
+        )
+        if blob is None:
+            return testbench.error.notfound(object_name, context)
+        return blob
+
+    def __get_all_soft_deleted_objects(self, bucket_name, context):
+        bucket_key = self.__bucket_key(bucket_name, context)
+        all_soft_deleted = []
+        for soft_deleted_list in self._soft_deleted_objects[bucket_key].values():
+            all_soft_deleted.extend(soft_deleted_list)
+        all_soft_deleted.sort(key=lambda blob: blob.metadata.generation)
+        return all_soft_deleted
+
     def list_object(self, request, bucket_name, context):
         with self._resources_lock:
             bucket = self.__get_bucket_for_object(bucket_name, context)
+            bucket_with_metadata = self.get_bucket(bucket_name, context)
             (
                 delimiter,
                 prefix,
@@ -219,14 +285,29 @@ class Database:
                 end_offset,
                 include_trailing_delimiter,
                 match_glob,
+                soft_deleted,
             ) = self.__extract_list_object_request(request, context)
             items = []
             prefixes = set()
-            for obj in bucket.values():
+
+            if (
+                soft_deleted
+                and not bucket_with_metadata.metadata.HasField("soft_delete_policy")
+            ) or (soft_deleted and versions):
+                return testbench.error.invalid("bad request", context)
+
+            objects = bucket.values()
+            if soft_deleted:
+                objects = self.__get_all_soft_deleted_objects(bucket_name, context)
+
+            for obj in objects:
                 generation = obj.metadata.generation
                 name = obj.metadata.name
-                if not versions and generation != self.__get_live_generation(
-                    bucket_name, name, context
+                if (
+                    not soft_deleted
+                    and not versions
+                    and generation
+                    != self.__get_live_generation(bucket_name, name, context)
                 ):
                     continue
                 if name.find(prefix) != 0:
@@ -282,12 +363,27 @@ class Database:
         return blob, live_generation
 
     def get_object(
-        self, bucket_name, object_name, context=None, generation=None, preconditions=[]
+        self,
+        bucket_name,
+        object_name,
+        context=None,
+        generation=None,
+        preconditions=[],
+        soft_deleted=False,
     ):
         with self._resources_lock:
-            blob, _ = self.__get_object(
-                bucket_name, object_name, context, generation, preconditions
-            )
+            blob = None
+            if not soft_deleted:
+                blob, _ = self.__get_object(
+                    bucket_name, object_name, context, generation, preconditions
+                )
+            else:
+                bucket_with_metadata = self.get_bucket(bucket_name, context)
+                if not bucket_with_metadata.metadata.HasField("soft_delete_policy"):
+                    testbench.error.invalid("SoftDeletePolicyRequired", context)
+                blob = self.__get_soft_deleted_object(
+                    bucket_name, object_name, int(generation), context
+                )
             # return a snapshot copy of the blob/blob.metadata
             if blob is None:
                 return None
@@ -336,6 +432,15 @@ class Database:
             if generation == 0 or live_generation == generation:
                 self.__del_live_generation(bucket_name, object_name, context)
             bucket = self.__get_bucket_for_object(bucket_name, context)
+            bucket_with_metadata = self.get_bucket(bucket_name, context)
+            if bucket_with_metadata.metadata.HasField("soft_delete_policy"):
+                self.__soft_delete_object(
+                    bucket_name,
+                    object_name,
+                    blob,
+                    bucket_with_metadata.metadata.soft_delete_policy.retention_duration.seconds,
+                    context,
+                )
             bucket.pop("%s#%d" % (blob.metadata.name, blob.metadata.generation), None)
 
     def do_update_object(
@@ -354,6 +459,47 @@ class Database:
             )
             return update_fn(blob, live_generation)
 
+    def restore_object(
+        self,
+        bucket_name: str,
+        object_name: str,
+        generation: int,
+        preconditions=[],
+        context=None,
+    ) -> T:
+        with self._resources_lock:
+            bucket_with_metadata = self.get_bucket(bucket_name, context)
+            if not bucket_with_metadata.metadata.HasField("soft_delete_policy"):
+                testbench.error.invalid("SoftDeletePolicyRequired", context)
+            bucket = self.__get_bucket_for_object(bucket_name, context)
+            blob = bucket.get("%s#%d" % (object_name, generation), None)
+            if blob is not None:
+                testbench.error.not_soft_deleted(context)
+
+            self.__remove_expired_objects_from_soft_delete(
+                bucket_name,
+                object_name,
+                context,
+            )
+            blob = self.__get_soft_deleted_object(
+                bucket_name, object_name, generation, context
+            )
+            if blob is not None:
+                blob.metadata.create_time.FromDatetime(
+                    datetime.datetime.now(datetime.timezone.utc)
+                )
+                blob.metadata.ClearField("soft_delete_time")
+                blob.metadata.metageneration = 1
+                blob.metadata.generation = blob.metadata.generation + 1
+                if bucket_with_metadata.metadata.autoclass.enabled is True:
+                    blob.metadata.storage_class = "STANDARD"
+                self.insert_object(bucket_name, blob, context, preconditions)
+                self.__remove_restored_soft_deleted_object(
+                    bucket_name, object_name, generation, context
+                )
+
+            return blob
+
     # === UPLOAD === #
 
     def get_upload(self, upload_id, context):
diff --git a/testbench/error.py b/testbench/error.py
index 41eda4f..a4d4689 100644
--- a/testbench/error.py
+++ b/testbench/error.py
@@ -112,6 +112,14 @@ def mismatch(
     generic(_simple_json_error(msg), rest_code, grpc_code, context)
 
 
+def not_soft_deleted(
+    context, rest_code=412, grpc_code=grpc.StatusCode.FAILED_PRECONDITION
+):
+    """This error is returned when object is not soft deleted but is either live or noncurrent"""
+    msg = "objectNotSoftDeleted"
+    generic(_simple_json_error(msg), rest_code, grpc_code, context)
+
+
 def notchanged(msg, context, rest_code=304, grpc_code=grpc.StatusCode.ABORTED):
     """Error returned when if*NotMatch or If-None-Match pre-conditions fail."""
     generic(
diff --git a/testbench/grpc_server.py b/testbench/grpc_server.py
index ee085f1..8fa6d68 100644
--- a/testbench/grpc_server.py
+++ b/testbench/grpc_server.py
@@ -689,6 +689,14 @@ class StorageServicer(storage_pb2_grpc.StorageServicer):
     def __get_bucket(self, bucket_name, context) -> storage_pb2.Bucket:
         return self.db.get_bucket(bucket_name, context).metadata
 
+    @retry_test(method="storage.objects.restore")
+    def RestoreObject(self, request, context):
+        preconditions = testbench.common.make_grpc_preconditions(request)
+        blob = self.db.restore_object(
+            request.bucket, request.object, request.generation, preconditions, context
+        )
+        return blob.metadata
+
     @retry_test(method="storage.objects.insert")
     def WriteObject(self, request_iterator, context):
         upload, is_resumable = gcs.upload.Upload.init_write_object_grpc(
diff --git a/testbench/rest_server.py b/testbench/rest_server.py
index 9490f24..9630840 100644
--- a/testbench/rest_server.py
+++ b/testbench/rest_server.py
@@ -568,14 +568,20 @@ def object_delete(bucket_name, object_name):
 @gcs.route("/b/<bucket_name>/o/<path:object_name>")
 @retry_test(method="storage.objects.get")
 def object_get(bucket_name, object_name):
+    soft_deleted = flask.request.args.get("softDeleted", False, bool)
+    media = flask.request.args.get("alt", None)
+    generation = flask.request.args.get("generation", None)
+    if (soft_deleted and generation is None) or (soft_deleted and media == "media"):
+        return testbench.error.invalid("invalid request", None)
+
     blob = db.get_object(
         bucket_name,
         object_name,
-        generation=flask.request.args.get("generation", None),
+        generation=generation,
         preconditions=testbench.common.make_json_preconditions(flask.request),
         context=None,
+        soft_deleted=soft_deleted,
     )
-    media = flask.request.args.get("alt", None)
     if media is None or media == "json":
         projection = testbench.common.extract_projection(flask.request, "noAcl", None)
         fields = flask.request.args.get("fields", None)
@@ -773,6 +779,21 @@ def objects_rewrite(src_bucket_name, src_object_name, dst_bucket_name, dst_objec
     return response
 
 
+@gcs.route("/b/<bucket_name>/o/<path:object_name>/restore", methods=["POST"])
+@retry_test(method="storage.objects.restore")
+def object_restore(bucket_name, object_name):
+    if flask.request.args.get("generation") is None:
+        return testbench.error.invalid("generation", None)
+    blob = db.restore_object(
+        bucket_name,
+        object_name,
+        int(flask.request.args.get("generation")),
+        testbench.common.make_json_preconditions(flask.request),
+    )
+    projection = testbench.common.extract_projection(flask.request, "noAcl", None)
+    return testbench.common.filter_response_rest(blob.rest_metadata(), projection, None)
+
+
 # === OBJECT ACCESS CONTROL === #
 
 


Test patch:
diff --git a/tests/test_database.py b/tests/test_database.py
index 0f06f43..11ee72d 100644
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -330,6 +330,115 @@ class TestDatabaseObject(unittest.TestCase):
             )
         self.assertEqual(rest.exception.code, 404)
 
+    def test_restore_object_no_soft_delete_policy(self):
+        with self.assertRaises(testbench.error.RestException) as rest:
+            _, _, _ = self.database.restore_object(
+                "bucket-name",
+                "object-name",
+                12345678,
+            )
+        self.assertEqual(rest.exception.code, 400)
+
+    def test_restore_object_not_soft_deleted(self):
+        request = testbench.common.FakeRequest(
+            args={},
+            data=json.dumps(
+                {
+                    "name": "sd-bucket-name",
+                    "softDeletePolicy": {"retentionDurationSeconds": 7 * 24 * 60 * 60},
+                }
+            ),
+        )
+        sd_bucket, _ = gcs.bucket.Bucket.init(request, None)
+        self.database.insert_bucket(sd_bucket, None)
+
+        request = testbench.common.FakeRequest(
+            args={"name": "object-name"}, data=b"12345678", headers={}, environ={}
+        )
+        blob, _ = gcs.object.Object.init_media(request, sd_bucket.metadata)
+        self.database.insert_object("sd-bucket-name", blob, context=None)
+
+        get_result = self.database.get_object(
+            "sd-bucket-name",
+            "object-name",
+            context=None,
+        )
+
+        with self.assertRaises(testbench.error.RestException) as rest:
+            _, _, _ = self.database.restore_object(
+                "sd-bucket-name",
+                "object-name",
+                get_result.metadata.generation,
+            )
+        self.assertEqual(rest.exception.code, 412)
+
+    def test_restore_object_generation_not_soft_deleted(self):
+        request = testbench.common.FakeRequest(
+            args={},
+            data=json.dumps(
+                {
+                    "name": "sd-bucket-name",
+                    "softDeletePolicy": {"retentionDurationSeconds": 7 * 24 * 60 * 60},
+                }
+            ),
+        )
+        sd_bucket, _ = gcs.bucket.Bucket.init(request, None)
+        self.database.insert_bucket(sd_bucket, None)
+
+        request = testbench.common.FakeRequest(
+            args={"name": "object-name"}, data=b"12345678", headers={}, environ={}
+        )
+        blob, _ = gcs.object.Object.init_media(request, sd_bucket.metadata)
+        self.database.insert_object("sd-bucket-name", blob, context=None)
+
+        get_result = self.database.get_object(
+            "sd-bucket-name",
+            "object-name",
+            context=None,
+        )
+
+        self.database.delete_object("sd-bucket-name", "object-name")
+
+        with self.assertRaises(testbench.error.RestException) as rest:
+            blob = self.database.restore_object(
+                "sd-bucket-name", "object-name", get_result.metadata.generation + 1
+            )
+        self.assertEqual(rest.exception.code, 404)
+
+    def test_restore_object_standard_storage(self):
+        request = testbench.common.FakeRequest(
+            args={},
+            data=json.dumps(
+                {
+                    "name": "sd-bucket-name",
+                    "softDeletePolicy": {"retentionDurationSeconds": 7 * 24 * 60 * 60},
+                    "autoclass": {"enabled": True, "terminalStorageClass": "NEARLINE"},
+                }
+            ),
+        )
+        sd_bucket, _ = gcs.bucket.Bucket.init(request, None)
+        self.database.insert_bucket(sd_bucket, None)
+
+        request = testbench.common.FakeRequest(
+            args={"name": "object-name"}, data=b"12345678", headers={}, environ={}
+        )
+        blob, _ = gcs.object.Object.init_media(request, sd_bucket.metadata)
+        self.database.insert_object("sd-bucket-name", blob, context=None)
+
+        get_result = self.database.get_object(
+            "sd-bucket-name",
+            "object-name",
+            context=None,
+        )
+
+        self.database.delete_object("sd-bucket-name", "object-name")
+        blob = self.database.restore_object(
+            "sd-bucket-name", "object-name", get_result.metadata.generation
+        )
+
+        self.assertNotEqual(get_result.metadata.generation, blob.metadata.generation)
+        self.assertEqual(blob.metadata.storage_class, "STANDARD")
+
 
 class TestDatabaseTemporaryResources(unittest.TestCase):
     """Test the Database class handling of uploads and rewrites."""
diff --git a/tests/test_error.py b/tests/test_error.py
index 35728b1..be3b500 100644
--- a/tests/test_error.py
+++ b/tests/test_error.py
@@ -106,6 +106,15 @@ class TestError(unittest.TestCase):
         )
         context.abort.assert_called_once_with(grpc.StatusCode.UNAVAILABLE, ANY)
 
+    def test_not_soft_deleted_error(self):
+        with self.assertRaises(error.RestException) as rest:
+            error.not_soft_deleted(None)
+        self.assertEqual(rest.exception.code, 412)
+
+        context = Mock()
+        error.not_soft_deleted(context)
+        context.abort.assert_called_once_with(grpc.StatusCode.FAILED_PRECONDITION, ANY)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_grpc_server.py b/tests/test_grpc_server.py
index 7d784f8..4b2ef2c 100755
--- a/tests/test_grpc_server.py
+++ b/tests/test_grpc_server.py
@@ -1396,6 +1396,46 @@ class TestGrpc(unittest.TestCase):
         self.grpc.WriteObject([r1], context=context)
         context.abort.assert_called_once()
 
+    def test_restore_object(self):
+        # Create a bucket with a soft delete policy
+        request = testbench.common.FakeRequest(
+            args={},
+            data=json.dumps(
+                {
+                    "name": "sd-bucket-name",
+                    "softDeletePolicy": {"retentionDurationSeconds": 7 * 24 * 60 * 60},
+                }
+            ),
+        )
+        sd_bucket, _ = gcs.bucket.Bucket.init(request, None)
+        self.db.insert_bucket(sd_bucket, None)
+
+        # Insert an object
+        media = b"The quick brown fox jumps over the lazy dog"
+        request = testbench.common.FakeRequest(
+            args={"name": "object-to-restore"}, data=media, headers={}, environ={}
+        )
+        blob, _ = gcs.object.Object.init_media(request, sd_bucket.metadata)
+        self.db.insert_object("sd-bucket-name", blob, context=None)
+        initial_generation = blob.metadata.generation
+
+        # Soft delete the object
+        self.db.delete_object("sd-bucket-name", "object-to-restore")
+
+        # Restore the soft deleted object
+        context = unittest.mock.Mock()
+        response = self.grpc.RestoreObject(
+            storage_pb2.RestoreObjectRequest(
+                bucket="projects/_/buckets/sd-bucket-name",
+                object="object-to-restore",
+                generation=initial_generation,
+            ),
+            context,
+        )
+        context.abort.assert_not_called()
+        self.assertIsNotNone(response)
+        self.assertNotEqual(initial_generation, response.generation)
+
     def test_rewrite_object(self):
         # We need a large enough payload to make sure the first rewrite does
         # not complete.  The minimum is 1 MiB
diff --git a/tests/test_testbench_object_metadata.py b/tests/test_testbench_object_metadata.py
index a870d6e..fa35920 100644
--- a/tests/test_testbench_object_metadata.py
+++ b/tests/test_testbench_object_metadata.py
@@ -248,6 +248,103 @@ class TestTestbenchObjectMetadata(unittest.TestCase):
         )
         self.assertEqual(response.status_code, 404)
 
+    def test_list_with_soft_deleted(self):
+        response = self.client.post(
+            "/storage/v1/b", data=json.dumps({"name": "bucket-name"})
+        )
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.get("/storage/v1/b/bucket-name/o?softDeleted=true")
+        self.assertEqual(response.status_code, 400)
+
+        response = self.client.post(
+            "/storage/v1/b",
+            data=json.dumps(
+                {
+                    "name": "sd-bucket-name",
+                    "softDeletePolicy": {"retentionDurationSeconds": 7 * 24 * 60 * 60},
+                }
+            ),
+        )
+        self.assertEqual(response.status_code, 200)
+
+        payload = "The quick brown fox jumps over the lazy dog"
+        response = self.client.put(
+            "/sd-bucket-name/fox.txt",
+            content_type="text/plain",
+            data=payload,
+        )
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.delete("/storage/v1/b/sd-bucket-name/o/fox.txt")
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.get(
+            "/storage/v1/b/sd-bucket-name/o?softDeleted=true&versions=true"
+        )
+        self.assertEqual(response.status_code, 400)
+
+        response = self.client.get("/storage/v1/b/sd-bucket-name/o?softDeleted=true")
+        self.assertEqual(response.status_code, 200)
+
+    def test_get_with_soft_deleted(self):
+        response = self.client.post(
+            "/storage/v1/b", data=json.dumps({"name": "bucket-name"})
+        )
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.get(
+            "/storage/v1/b/bucket-name/o/some-object?softDeleted=true"
+        )
+        self.assertEqual(response.status_code, 400)
+
+        response = self.client.get(
+            "/storage/v1/b/bucket-name/o/some-object?softDeleted=true&generation=12345678"
+        )
+        self.assertEqual(response.status_code, 400)
+
+        response = self.client.post(
+            "/storage/v1/b",
+            data=json.dumps(
+                {
+                    "name": "sd-bucket-name",
+                    "softDeletePolicy": {"retentionDurationSeconds": 7 * 24 * 60 * 60},
+                }
+            ),
+        )
+        self.assertEqual(response.status_code, 200)
+
+        payload = "The quick brown fox jumps over the lazy dog"
+        response = self.client.put(
+            "/sd-bucket-name/fox.txt",
+            content_type="text/plain",
+            data=payload,
+        )
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.get("/storage/v1/b/sd-bucket-name/o/fox.txt")
+        self.assertEqual(response.status_code, 200)
+        generation = json.loads(response.data).get("generation")
+
+        response = self.client.delete("/storage/v1/b/sd-bucket-name/o/fox.txt")
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.get(
+            "/storage/v1/b/sd-bucket-name/o/fox.txt?softDeleted=true&alt=media"
+        )
+        self.assertEqual(response.status_code, 400)
+
+        response = self.client.get(
+            "/storage/v1/b/sd-bucket-name/o/fox.txt?softDeleted=true"
+        )
+        self.assertEqual(response.status_code, 400)
+
+        response = self.client.get(
+            "/storage/v1/b/sd-bucket-name/o/fox.txt?softDeleted=true&generation="
+            + generation
+        )
+        self.assertEqual(response.status_code, 200)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_testbench_object_special.py b/tests/test_testbench_object_special.py
index f7232b1..aae7b7c 100644
--- a/tests/test_testbench_object_special.py
+++ b/tests/test_testbench_object_special.py
@@ -270,6 +270,46 @@ class TestTestbenchObjectSpecial(unittest.TestCase):
         self.assertEqual(response.status_code, 200)
         self.assertEqual(len(response.data.decode("utf-8")), len(payload))
 
+    def test_object_restore(self):
+        response = self.client.post(
+            "/storage/v1/b",
+            data=json.dumps(
+                {
+                    "name": "sd-bucket-name",
+                    "softDeletePolicy": {"retentionDurationSeconds": 7 * 24 * 60 * 60},
+                }
+            ),
+        )
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.put(
+            "/sd-bucket-name/sd-restore-obj",
+            content_type="text/plain",
+            data="The quick brown fox jumps over the lazy dog\n",
+        )
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.get("/storage/v1/b/sd-bucket-name/o/sd-restore-obj")
+        self.assertEqual(response.status_code, 200)
+        blob = json.loads(response.data)
+
+        response = self.client.delete("/storage/v1/b/sd-bucket-name/o/sd-restore-obj")
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.post(
+            "/storage/v1/b/sd-bucket-name/o/sd-restore-obj/restore?generation="
+            + blob.get("generation")
+        )
+        self.assertEqual(response.status_code, 200)
+        restored_blob = json.loads(response.data)
+        self.assertNotEqual(blob.get("generation"), restored_blob.get("generation"))
+
+    def test_object_restore_no_generation(self):
+        response = self.client.post(
+            "/storage/v1/b/sd-bucket-name/o/sd-restore-obj/restore"
+        )
+        self.assertEqual(response.status_code, 400)
+
 
 if __name__ == "__main__":
     unittest.main()