Add state transition for restoring objects from S3 glacier (#5202)

docs/docs/configuration/state_transition/models.rst

@@ -56,6 +56,16 @@ Advancement:
 
     Call `boto3.client("dax").describe_clusters(..)`.
 
+Service: S3 (Glacier Restoration)
+---------------
+
+**Model**: `s3::keyrestore`   :raw-html:`<br />`
+Available States:
+
+    None --> "IN_PROGRESS" --> "RESTORED"
+
+Transition type: Immediate - transitions immediately
+
 Service: Support
 ------------------
 

moto/s3/models.py

@@ -32,6 +32,8 @@ from moto.core.utils import (
     BackendDict,
 )
 from moto.cloudwatch.models import MetricDatum
+from moto.moto_api import state_manager
+from moto.moto_api._internal.managed_state_model import ManagedState
 from moto.utilities.tagging_service import TaggingService
 from moto.utilities.utils import LowercaseDict, md5_hash
 from moto.s3.exceptions import (
@@ -101,7 +103,7 @@ class FakeDeleteMarker(BaseModel):
         return self._version_id
 
 
-class FakeKey(BaseModel):
+class FakeKey(BaseModel, ManagedState):
     def __init__(
         self,
         name,
@@ -121,6 +123,14 @@ class FakeKey(BaseModel):
         lock_until=None,
         s3_backend=None,
     ):
+        ManagedState.__init__(
+            self,
+            "s3::keyrestore",
+            transitions=[
+                (None, "IN_PROGRESS"),
+                ("IN_PROGRESS", "RESTORED"),
+            ],
+        )
         self.name = name
         self.last_modified = datetime.datetime.utcnow()
         self.acl = get_canned_acl("private")
@@ -256,8 +266,13 @@ class FakeKey(BaseModel):
         if self._storage_class != "STANDARD":
             res["x-amz-storage-class"] = self._storage_class
         if self._expiry is not None:
-            rhdr = 'ongoing-request="false", expiry-date="{0}"'
-            res["x-amz-restore"] = rhdr.format(self.expiry_date)
+            if self.status == "IN_PROGRESS":
+                header = 'ongoing-request="true"'
+            else:
+                header = 'ongoing-request="false", expiry-date="{0}"'.format(
+                    self.expiry_date
+                )
+            res["x-amz-restore"] = header
 
         if self._is_versioned:
             res["x-amz-version-id"] = str(self.version_id)
@@ -1380,6 +1395,10 @@ class S3Backend(BaseBackend, CloudWatchMetricProvider):
         self.buckets = {}
         self.tagger = TaggingService()
 
+        state_manager.register_default_transition(
+            "s3::keyrestore", transition={"progression": "immediate"}
+        )
+
     @property
     def _url_module(self):
         # The urls-property can be different depending on env variables
@@ -1741,6 +1760,7 @@ class S3Backend(BaseBackend, CloudWatchMetricProvider):
                 key = key.multipart.parts[part_number]
 
         if isinstance(key, FakeKey):
+            key.advance()
             return key
         else:
             return None

tests/test_s3/test_s3.py

@@ -15,6 +15,7 @@ from botocore.handlers import disable_signing
 from freezegun import freeze_time
 import requests
 
+from moto.moto_api import state_manager
 from moto.s3.responses import DEFAULT_REGION_NAME
 from unittest import SkipTest
 import pytest
@@ -549,6 +550,38 @@ def test_restore_key():
         )
 
 
+@freeze_time("2012-01-01 12:00:00")
+@mock_s3
+def test_restore_key_transition():
+    if settings.TEST_SERVER_MODE:
+        raise SkipTest("Can't set transition directly in ServerMode")
+
+    state_manager.set_transition(
+        model_name="s3::keyrestore", transition={"progression": "manual", "times": 1}
+    )
+
+    s3 = boto3.resource("s3", region_name=DEFAULT_REGION_NAME)
+    bucket = s3.Bucket("foobar")
+    bucket.create()
+
+    key = bucket.put_object(Key="the-key", Body=b"somedata", StorageClass="GLACIER")
+    key.restore.should.equal(None)
+    key.restore_object(RestoreRequest={"Days": 1})
+
+    # first call: there should be an ongoing request
+    key.restore.should.contain('ongoing-request="true"')
+
+    # second call: request should be done
+    key.load()
+    key.restore.should.contain('ongoing-request="false"')
+
+    # third call: request should still be done
+    key.load()
+    key.restore.should.contain('ongoing-request="false"')
+
+    state_manager.unset_transition(model_name="s3::keyrestore")
+
+
 @mock_s3
 def test_cannot_restore_standard_class_object():
     s3 = boto3.resource("s3", region_name=DEFAULT_REGION_NAME)