sdk/python/tests/test_storage_classes.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: Apache-2.0
   4
   5 import arvados
   6 import pycurl
   7
   8 import unittest
   9 import parameterized
  10 from . import arvados_testutil as tutil
  11 from .arvados_testutil import DiskCacheBase
  12
  13 @tutil.skip_sleep
  14 @parameterized.parameterized_class([{"disk_cache": True}, {"disk_cache": False}])
  15 class KeepStorageClassesTestCase(unittest.TestCase, tutil.ApiClientMock, DiskCacheBase):
  16     disk_cache = False
  17
  18     def setUp(self):
  19         self.api_client = self.mock_keep_services(count=2)
  20         self.keep_client = arvados.KeepClient(api_client=self.api_client, block_cache=self.make_block_cache(self.disk_cache))
  21         self.data = b'xyzzy'
  22         self.locator = '1271ed5ef305aadabc605b1609e24c52'
  23
  24     def tearDown(self):
  25         DiskCacheBase.tearDown(self)
  26
  27     def test_multiple_default_storage_classes_req_header(self):
  28         api_mock = self.api_client_mock()
  29         api_mock.config.return_value = {
  30             'StorageClasses': {
  31                 'foo': { 'Default': True },
  32                 'bar': { 'Default': True },
  33                 'baz': { 'Default': False }
  34             }
  35         }
  36         api_client = self.mock_keep_services(api_mock=api_mock, count=2)
  37         keep_client = arvados.KeepClient(api_client=api_client, block_cache=self.make_block_cache(self.disk_cache))
  38         resp_hdr = {
  39             'x-keep-storage-classes-confirmed': 'foo=1, bar=1',
  40             'x-keep-replicas-stored': 1
  41         }
  42         with tutil.mock_keep_responses(self.locator, 200, **resp_hdr) as mock:
  43             keep_client.put(self.data, copies=1)
  44             req_hdr = mock.responses[0]
  45             self.assertIn(
  46                 'X-Keep-Storage-Classes: bar, foo', req_hdr.getopt(pycurl.HTTPHEADER))
  47
  48     def test_storage_classes_req_header(self):
  49         self.assertEqual(
  50             self.api_client.config()['StorageClasses'],
  51             {'default': {'Default': True}})
  52         cases = [
  53             # requested, expected
  54             [['foo'], 'X-Keep-Storage-Classes: foo'],
  55             [['bar', 'foo'], 'X-Keep-Storage-Classes: bar, foo'],
  56             [[], 'X-Keep-Storage-Classes: default'],
  57             [None, 'X-Keep-Storage-Classes: default'],
  58         ]
  59         for req_classes, expected_header in cases:
  60             headers = {'x-keep-replicas-stored': 1}
  61             if req_classes is None or len(req_classes) == 0:
  62                 confirmed_hdr = 'default=1'
  63             elif len(req_classes) > 0:
  64                 confirmed_hdr = ', '.join(["{}=1".format(cls) for cls in req_classes])
  65             headers.update({'x-keep-storage-classes-confirmed': confirmed_hdr})
  66             with tutil.mock_keep_responses(self.locator, 200, **headers) as mock:
  67                 self.keep_client.put(self.data, copies=1, classes=req_classes)
  68                 req_hdr = mock.responses[0]
  69                 self.assertIn(expected_header, req_hdr.getopt(pycurl.HTTPHEADER))
  70
  71     def test_partial_storage_classes_put(self):
  72         headers = {
  73             'x-keep-replicas-stored': 1,
  74             'x-keep-storage-classes-confirmed': 'foo=1'}
  75         with tutil.mock_keep_responses(self.locator, 200, 503, **headers) as mock:
  76             with self.assertRaises(arvados.errors.KeepWriteError):
  77                 self.keep_client.put(self.data, copies=1, classes=['foo', 'bar'], num_retries=0)
  78             # 1st request, both classes pending
  79             req1_headers = mock.responses[0].getopt(pycurl.HTTPHEADER)
  80             self.assertIn('X-Keep-Storage-Classes: bar, foo', req1_headers)
  81             # 2nd try, 'foo' class already satisfied
  82             req2_headers = mock.responses[1].getopt(pycurl.HTTPHEADER)
  83             self.assertIn('X-Keep-Storage-Classes: bar', req2_headers)
  84
  85     def test_successful_storage_classes_put_requests(self):
  86         cases = [
  87             # wanted_copies, wanted_classes, confirmed_copies, confirmed_classes, expected_requests
  88             [ 1, ['foo'], 1, 'foo=1', 1],
  89             [ 1, ['foo'], 2, 'foo=2', 1],
  90             [ 2, ['foo'], 2, 'foo=2', 1],
  91             [ 2, ['foo'], 1, 'foo=1', 2],
  92             [ 1, ['foo', 'bar'], 1, 'foo=1, bar=1', 1],
  93             [ 1, ['foo', 'bar'], 2, 'foo=2, bar=2', 1],
  94             [ 2, ['foo', 'bar'], 2, 'foo=2, bar=2', 1],
  95             [ 2, ['foo', 'bar'], 1, 'foo=1, bar=1', 2],
  96             [ 1, ['foo', 'bar'], 1, None, 1],
  97             [ 1, ['foo'], 1, None, 1],
  98             [ 2, ['foo'], 2, None, 1],
  99             [ 2, ['foo'], 1, None, 2],
 100         ]
 101         for w_copies, w_classes, c_copies, c_classes, e_reqs in cases:
 102             headers = {'x-keep-replicas-stored': c_copies}
 103             if c_classes is not None:
 104                 headers.update({'x-keep-storage-classes-confirmed': c_classes})
 105             with tutil.mock_keep_responses(self.locator, 200, 200, **headers) as mock:
 106                 case_desc = 'wanted_copies={}, wanted_classes="{}", confirmed_copies={}, confirmed_classes="{}", expected_requests={}'.format(w_copies, ', '.join(w_classes), c_copies, c_classes, e_reqs)
 107                 self.assertEqual(self.locator,
 108                     self.keep_client.put(self.data, copies=w_copies, classes=w_classes),
 109                     case_desc)
 110                 self.assertEqual(e_reqs, mock.call_count, case_desc)
 111
 112     def test_failed_storage_classes_put_requests(self):
 113         cases = [
 114             # wanted_copies, wanted_classes, confirmed_copies, confirmed_classes, return_code
 115             [ 1, ['foo'], 1, 'bar=1', 200],
 116             [ 1, ['foo'], 1, None, 503],
 117             [ 2, ['foo'], 1, 'bar=1, foo=0', 200],
 118             [ 3, ['foo'], 1, 'bar=1, foo=1', 200],
 119             [ 3, ['foo', 'bar'], 1, 'bar=2, foo=1', 200],
 120         ]
 121         for w_copies, w_classes, c_copies, c_classes, return_code in cases:
 122             headers = {'x-keep-replicas-stored': c_copies}
 123             if c_classes is not None:
 124                 headers.update({'x-keep-storage-classes-confirmed': c_classes})
 125             with tutil.mock_keep_responses(self.locator, return_code, return_code, **headers):
 126                 case_desc = 'wanted_copies={}, wanted_classes="{}", confirmed_copies={}, confirmed_classes="{}"'.format(w_copies, ', '.join(w_classes), c_copies, c_classes)
 127                 with self.assertRaises(arvados.errors.KeepWriteError, msg=case_desc):
 128                     self.keep_client.put(self.data, copies=w_copies, classes=w_classes, num_retries=0)