From e1d9c2878f2a5bc2b628e0631b9d9e590c7fec3b Mon Sep 17 00:00:00 2001
From: Brian Pandola <bpandola@hsdp.io>
Date: Tue, 24 Apr 2018 17:30:17 -0700
Subject: [PATCH] Add support for Redshift.Waiter.ClusterRestored

* Add `restored_from_snapshot` boolean to Cluster metadata.
* Return `RestoreStatus` from describe_db_clusters if cluster was restored from a snapshot.

Fixes #1506
---
 moto/redshift/models.py              | 17 +++++++++--
 tests/test_redshift/test_redshift.py | 42 ++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/moto/redshift/models.py b/moto/redshift/models.py
index 4eb9d6b5c..4eafcfc79 100644
--- a/moto/redshift/models.py
+++ b/moto/redshift/models.py
@@ -73,7 +73,8 @@ class Cluster(TaggableResourceMixin, BaseModel):
                  preferred_maintenance_window, cluster_parameter_group_name,
                  automated_snapshot_retention_period, port, cluster_version,
                  allow_version_upgrade, number_of_nodes, publicly_accessible,
-                 encrypted, region_name, tags=None, iam_roles_arn=None):
+                 encrypted, region_name, tags=None, iam_roles_arn=None,
+                 restored_from_snapshot=False):
         super(Cluster, self).__init__(region_name, tags)
         self.redshift_backend = redshift_backend
         self.cluster_identifier = cluster_identifier
@@ -119,6 +120,7 @@ class Cluster(TaggableResourceMixin, BaseModel):
             self.number_of_nodes = 1
 
         self.iam_roles_arn = iam_roles_arn or []
+        self.restored_from_snapshot = restored_from_snapshot
 
     @classmethod
     def create_from_cloudformation_json(cls, resource_name, cloudformation_json, region_name):
@@ -242,7 +244,15 @@ class Cluster(TaggableResourceMixin, BaseModel):
                 "IamRoleArn": iam_role_arn
             } for iam_role_arn in self.iam_roles_arn]
         }
-
+        if self.restored_from_snapshot:
+            json_response['RestoreStatus'] = {
+                'Status': 'completed',
+                'CurrentRestoreRateInMegaBytesPerSecond': 123.0,
+                'SnapshotSizeInMegaBytes': 123,
+                'ProgressInMegaBytes': 123,
+                'ElapsedTimeInSeconds': 123,
+                'EstimatedTimeToCompletionInSeconds': 123
+            }
         try:
             json_response['ClusterSnapshotCopyStatus'] = self.cluster_snapshot_copy_status
         except AttributeError:
@@ -639,7 +649,8 @@ class RedshiftBackend(BaseBackend):
             "cluster_version": snapshot.cluster.cluster_version,
             "number_of_nodes": snapshot.cluster.number_of_nodes,
             "encrypted": snapshot.cluster.encrypted,
-            "tags": snapshot.cluster.tags
+            "tags": snapshot.cluster.tags,
+            "restored_from_snapshot": True
         }
         create_kwargs.update(kwargs)
         return self.create_cluster(**create_kwargs)
diff --git a/tests/test_redshift/test_redshift.py b/tests/test_redshift/test_redshift.py
index 96e3ee5b3..6e027b86c 100644
--- a/tests/test_redshift/test_redshift.py
+++ b/tests/test_redshift/test_redshift.py
@@ -818,6 +818,48 @@ def test_create_cluster_from_snapshot():
     new_cluster['Endpoint']['Port'].should.equal(1234)
 
 
+@mock_redshift
+def test_create_cluster_from_snapshot_with_waiter():
+    client = boto3.client('redshift', region_name='us-east-1')
+    original_cluster_identifier = 'original-cluster'
+    original_snapshot_identifier = 'original-snapshot'
+    new_cluster_identifier = 'new-cluster'
+
+    client.create_cluster(
+        ClusterIdentifier=original_cluster_identifier,
+        ClusterType='single-node',
+        NodeType='ds2.xlarge',
+        MasterUsername='username',
+        MasterUserPassword='password',
+    )
+    client.create_cluster_snapshot(
+        SnapshotIdentifier=original_snapshot_identifier,
+        ClusterIdentifier=original_cluster_identifier
+    )
+    response = client.restore_from_cluster_snapshot(
+        ClusterIdentifier=new_cluster_identifier,
+        SnapshotIdentifier=original_snapshot_identifier,
+        Port=1234
+    )
+    response['Cluster']['ClusterStatus'].should.equal('creating')
+
+    client.get_waiter('cluster_restored').wait(
+        ClusterIdentifier=new_cluster_identifier,
+        WaiterConfig={
+            'Delay': 1,
+            'MaxAttempts': 2,
+        }
+    )
+
+    response = client.describe_clusters(
+        ClusterIdentifier=new_cluster_identifier
+    )
+    new_cluster = response['Clusters'][0]
+    new_cluster['NodeType'].should.equal('ds2.xlarge')
+    new_cluster['MasterUsername'].should.equal('username')
+    new_cluster['Endpoint']['Port'].should.equal(1234)
+
+
 @mock_redshift
 def test_create_cluster_from_non_existent_snapshot():
     client = boto3.client('redshift', region_name='us-east-1')