Added retrieve_inventory_job to Vault.

boto · Jun 5, 2013 · 33de290 · 33de290
1 parent efd9af3
commit 33de290
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 2 deletions.
diff --git a/boto/glacier/vault.py b/boto/glacier/vault.py
@@ -315,8 +315,8 @@ def retrieve_inventory(self, sns_topic=None,
             sends notification when the job is completed and the output
             is ready for you to download.
 
-        :rtype: :class:`boto.glacier.job.Job`
-        :return: A Job object representing the retrieval job.
+        :rtype: str
+        :return: The ID of the job
         """
         job_data = {'Type': 'inventory-retrieval'}
         if sns_topic is not None:
@@ -327,6 +327,25 @@ def retrieve_inventory(self, sns_topic=None,
         response = self.layer1.initiate_job(self.name, job_data)
         return response['JobId']
 
+    def retrieve_inventory_job(self, **kwargs):
+        """
+        Identical to ``retrieve_inventory``, but returns a ``Job`` instance
+        instead of just the job ID.
+
+        :type description: str
+        :param description: An optional description for the job.
+
+        :type sns_topic: str
+        :param sns_topic: The Amazon SNS topic ARN where Amazon Glacier
+            sends notification when the job is completed and the output
+            is ready for you to download.
+
+        :rtype: :class:`boto.glacier.job.Job`
+        :return: A Job object representing the retrieval job.
+        """
+        job_id = self.retrieve_inventory(**kwargs)
+        return self.get_job(job_id)
+
     def delete_archive(self, archive_id):
         """
         This operation deletes an archive from the vault.

diff --git a/tests/unit/glacier/test_vault.py b/tests/unit/glacier/test_vault.py
@@ -27,6 +27,8 @@
 from mock import ANY
 
 from boto.glacier import vault
+from boto.glacier.job import Job
+from boto.glacier.response import GlacierResponse
 
 
 class TestVault(unittest.TestCase):
@@ -96,6 +98,55 @@ def test_part_size_needs_to_be_adjusted(self):
         self.vault.create_archive_writer.assert_called_with(
             description=mock.ANY, part_size=expected_part_size)
 
+    def test_retrieve_inventory(self):
+        class FakeResponse(object):
+            status = 202
+
+            def getheader(self, key, default=None):
+                if key == 'x-amz-job-id':
+                    return 'HkF9p6'
+                elif key == 'Content-Type':
+                    return 'application/json'
+
+                return 'something'
+
+            def read(self, amt=None):
+                return """{
+  "Action": "ArchiveRetrieval",
+  "ArchiveId": "NkbByEejwEggmBz2fTHgJrg0XBoDfjP4q6iu87-EXAMPLEArchiveId",
+  "ArchiveSizeInBytes": 16777216,
+  "ArchiveSHA256TreeHash": "beb0fe31a1c7ca8c6c04d574ea906e3f97",
+  "Completed": false,
+  "CreationDate": "2012-05-15T17:21:39.339Z",
+  "CompletionDate": "2012-05-15T17:21:43.561Z",
+  "InventorySizeInBytes": null,
+  "JobDescription": "My ArchiveRetrieval Job",
+  "JobId": "HkF9p6",
+  "RetrievalByteRange": "0-16777215",
+  "SHA256TreeHash": "beb0fe31a1c7ca8c6c04d574ea906e3f97b31fd",
+  "SNSTopic": "arn:aws:sns:us-east-1:012345678901:mytopic",
+  "StatusCode": "InProgress",
+  "StatusMessage": "Operation in progress.",
+  "VaultARN": "arn:aws:glacier:us-east-1:012345678901:vaults/examplevault"
+}"""
+
+        raw_resp = FakeResponse()
+        init_resp = GlacierResponse(raw_resp, [('x-amz-job-id', 'JobId')])
+        raw_resp_2 = FakeResponse()
+        desc_resp = GlacierResponse(raw_resp_2, [])
+
+        with mock.patch.object(self.vault.layer1, 'initiate_job',
+                               return_value=init_resp):
+            with mock.patch.object(self.vault.layer1, 'describe_job',
+                                   return_value=desc_resp):
+                # The old/back-compat variant of the call.
+                self.assertEqual(self.vault.retrieve_inventory(), 'HkF9p6')
+
+                # The variant the returns a full ``Job`` object.
+                job = self.vault.retrieve_inventory_job()
+                self.assertTrue(isinstance(job, Job))
+                self.assertEqual(job.id, 'HkF9p6')
+
 
 class TestConcurrentUploads(unittest.TestCase):