17465: Adds tests confirming KeepClient storage support on Collection.

Also, made storage_classes_desired field be loaded into a Collection from the
API server.

Arvados-DCO-1.1-Signed-off-by: Lucas Di Pentima <lucas.dipentima@curii.com>

diff --git a/sdk/python/arvados/collection.py b/sdk/python/arvados/collection.py
index 846de0115e0866dc2499a429aa0dda1e0c41918e..319046dc01aebbcf4d0496531be5e86c25c8f783 100644 (file)
--- a/sdk/python/arvados/collection.py
+++ b/sdk/python/arvados/collection.py
@@ -1439,9 +1439,11 @@ class Collection(RichCollectionBase):
         self._manifest_text = self._api_response['manifest_text']
         self._portable_data_hash = self._api_response['portable_data_hash']
         # If not overriden via kwargs, we should try to load the
-        # replication_desired from the API server
+        # replication_desired and storage_classes_desired from the API server
         if self.replication_desired is None:
             self.replication_desired = self._api_response.get('replication_desired', None)
+        if self.storage_classes_desired is None:
+            self.storage_classes_desired = self._api_response.get('storage_classes_desired', None)
 
     def _populate(self):
         if self._manifest_text is None:

diff --git a/sdk/python/tests/test_collections.py b/sdk/python/tests/test_collections.py
index e1c6fc2bbca5aeb76ecc438c7b07739040a843ad..eaf68ec1f5b5afa9ab2e371d9058ae48a763bc20 100644 (file)
--- a/sdk/python/tests/test_collections.py
+++ b/sdk/python/tests/test_collections.py
@@ -908,6 +908,24 @@ class NewCollectionTestCase(unittest.TestCase, CollectionTestMixin):
         self.assertEqual(c1.manifest_text, c2.manifest_text)
         self.assertNotEqual(c1.replication_desired, c2.replication_desired)
 
+    def test_storage_classes_desired_kept_on_load(self):
+        m = '. 781e5e245d69b566979b86e28d23f2c7+10 0:10:count1.txt 0:10:count2.txt\n'
+        c1 = Collection(m, storage_classes_desired=['archival'])
+        c1.save_new()
+        loc = c1.manifest_locator()
+        c2 = Collection(loc)
+        self.assertEqual(c1.manifest_text, c2.manifest_text)
+        self.assertEqual(c1.storage_classes_desired, c2.storage_classes_desired)
+
+    def test_storage_classes_desired_not_loaded_if_provided(self):
+        m = '. 781e5e245d69b566979b86e28d23f2c7+10 0:10:count1.txt 0:10:count2.txt\n'
+        c1 = Collection(m, storage_classes_desired=['archival'])
+        c1.save_new()
+        loc = c1.manifest_locator()
+        c2 = Collection(loc, storage_classes_desired=['default'])
+        self.assertEqual(c1.manifest_text, c2.manifest_text)
+        self.assertNotEqual(c1.storage_classes_desired, c2.storage_classes_desired)
+
     def test_init_manifest(self):
         m1 = """. 5348b82a029fd9e971a811ce1f71360b+43 0:43:md5sum.txt
 . 085c37f02916da1cad16f93c54d899b7+41 0:41:md5sum.txt
@@ -1246,6 +1264,16 @@ class NewCollectionTestCaseWithServersAndTokens(run_test_server.TestCaseWithServ
     def setUp(self):
         self.keep_put = getattr(arvados.keep.KeepClient, 'put')
 
+    @mock.patch('arvados.keep.KeepClient.put', autospec=True)
+    def test_storage_classes_desired(self, put_mock):
+        put_mock.side_effect = self.keep_put
+        c = Collection(storage_classes_desired=['default'])
+        with c.open("file.txt", 'wb') as f:
+            f.write('content')
+        c.save_new()
+        _, kwargs = put_mock.call_args
+        self.assertEqual(['default'], kwargs['classes'])
+
     @mock.patch('arvados.keep.KeepClient.put', autospec=True)
     def test_repacked_block_submission_get_permission_token(self, mocked_put):
         '''