sdk/python/tests/test_util.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: Apache-2.0
   4
   5 import itertools
   6 import os
   7 import subprocess
   8 import unittest
   9
  10 import parameterized
  11 import pytest
  12 from unittest import mock
  13
  14 import arvados
  15 import arvados.util
  16
  17 class KeysetTestHelper:
  18     def __init__(self, expect, expect_num_retries=0):
  19         self.n = 0
  20         self.expect = expect
  21         self.expect_num_retries = expect_num_retries
  22
  23     def fn(self, **kwargs):
  24         assert kwargs == self.expect[self.n][0]
  25         return self
  26
  27     def execute(self, num_retries):
  28         assert num_retries == self.expect_num_retries
  29         self.n += 1
  30         return self.expect[self.n-1][1]
  31
  32 _SELECT_FAKE_ITEM = {
  33     'uuid': 'zzzzz-zyyyz-zzzzzyyyyywwwww',
  34     'name': 'KeysetListAllTestCase.test_select mock',
  35     'created_at': '2023-08-28T12:34:56.123456Z',
  36 }
  37
  38 _FAKE_COMPUTED_PERMISSIONS_ITEM = {
  39     'user_uuid': 'zzzzz-zyyyz-zzzzzyyyyywwwww',
  40     'target_uuid': 'zzzzz-ttttt-xxxxxyyyyyzzzzz',
  41     'perm_level': 'can_write',
  42 }
  43
  44 class KeysetListAllTestCase(unittest.TestCase):
  45     def test_empty(self):
  46         ks = KeysetTestHelper([[
  47             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
  48             {"items": []}
  49         ]])
  50
  51         ls = list(arvados.util.keyset_list_all(ks.fn))
  52         self.assertEqual(ls, [])
  53
  54     def test_oneitem(self):
  55         ks = KeysetTestHelper([[
  56             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
  57             {"items": [{"created_at": "1", "uuid": "1"}]}
  58         ], [
  59             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", "=", "1"], ["uuid", ">", "1"]]},
  60             {"items": []}
  61         ],[
  62             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">", "1"]]},
  63             {"items": []}
  64         ]])
  65
  66         ls = list(arvados.util.keyset_list_all(ks.fn))
  67         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}])
  68
  69     def test_onepage2(self):
  70         ks = KeysetTestHelper([[
  71             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
  72             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}]}
  73         ], [
  74             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "2"], ["uuid", "!=", "2"]]},
  75             {"items": []}
  76         ]])
  77
  78         ls = list(arvados.util.keyset_list_all(ks.fn))
  79         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}])
  80
  81     def test_onepage3(self):
  82         ks = KeysetTestHelper([[
  83             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
  84             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}, {"created_at": "3", "uuid": "3"}]}
  85         ], [
  86             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "3"], ["uuid", "!=", "3"]]},
  87             {"items": []}
  88         ]])
  89
  90         ls = list(arvados.util.keyset_list_all(ks.fn))
  91         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}, {"created_at": "3", "uuid": "3"}])
  92
  93
  94     def test_twopage(self):
  95         ks = KeysetTestHelper([[
  96             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
  97             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}]}
  98         ], [
  99             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "2"], ["uuid", "!=", "2"]]},
 100             {"items": [{"created_at": "3", "uuid": "3"}, {"created_at": "4", "uuid": "4"}]}
 101         ], [
 102             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "4"], ["uuid", "!=", "4"]]},
 103             {"items": []}
 104         ]])
 105
 106         ls = list(arvados.util.keyset_list_all(ks.fn))
 107         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"},
 108                               {"created_at": "2", "uuid": "2"},
 109                               {"created_at": "3", "uuid": "3"},
 110                               {"created_at": "4", "uuid": "4"}
 111         ])
 112
 113     def test_repeated_key(self):
 114         ks = KeysetTestHelper([[
 115             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
 116             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}, {"created_at": "2", "uuid": "3"}]}
 117         ], [
 118             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "2"], ["uuid", "!=", "3"]]},
 119             {"items": [{"created_at": "2", "uuid": "2"}, {"created_at": "2", "uuid": "4"}]}
 120         ], [
 121             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", "=", "2"], ["uuid", ">", "4"]]},
 122             {"items": []}
 123         ], [
 124             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">", "2"]]},
 125             {"items": [{"created_at": "3", "uuid": "5"}, {"created_at": "4", "uuid": "6"}]}
 126         ], [
 127             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "4"], ["uuid", "!=", "6"]]},
 128             {"items": []}
 129         ],
 130         ])
 131
 132         ls = list(arvados.util.keyset_list_all(ks.fn))
 133         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"},
 134                               {"created_at": "2", "uuid": "2"},
 135                               {"created_at": "2", "uuid": "3"},
 136                               {"created_at": "2", "uuid": "4"},
 137                               {"created_at": "3", "uuid": "5"},
 138                               {"created_at": "4", "uuid": "6"}
 139         ])
 140
 141     def test_onepage_withfilter(self):
 142         ks = KeysetTestHelper([[
 143             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["foo", ">", "bar"]]},
 144             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}]}
 145         ], [
 146             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "2"], ["uuid", "!=", "2"], ["foo", ">", "bar"]]},
 147             {"items": []}
 148         ]])
 149
 150         ls = list(arvados.util.keyset_list_all(ks.fn, filters=[["foo", ">", "bar"]]))
 151         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}])
 152
 153     def test_onepage_desc(self):
 154         ks = KeysetTestHelper([[
 155             {"limit": 1000, "count": "none", "order": ["created_at desc", "uuid desc"], "filters": []},
 156             {"items": [{"created_at": "2", "uuid": "2"}, {"created_at": "1", "uuid": "1"}]}
 157         ], [
 158             {"limit": 1000, "count": "none", "order": ["created_at desc", "uuid desc"], "filters": [["created_at", "<=", "1"], ["uuid", "!=", "1"]]},
 159             {"items": []}
 160         ]])
 161
 162         ls = list(arvados.util.keyset_list_all(ks.fn, ascending=False))
 163         self.assertEqual(ls, [{"created_at": "2", "uuid": "2"}, {"created_at": "1", "uuid": "1"}])
 164
 165     @parameterized.parameterized.expand(
 166         (fake_item, key_fields, order_key, select)
 167         for (fake_item, key_fields) in [
 168             (_SELECT_FAKE_ITEM, ('uuid',)),
 169             (_FAKE_COMPUTED_PERMISSIONS_ITEM, ('user_uuid', 'target_uuid')),
 170         ]
 171         for order_key in fake_item
 172         if order_key != 'perm_level'
 173         for count in range(len(fake_item) + 1)
 174         for select in itertools.combinations(fake_item, count)
 175     )
 176     def test_select(self, fake_item, key_fields, order_key, select):
 177         # keyset_list_all must have both uuid and order_key to function.
 178         # Test that it selects those fields along with user-specified ones.
 179         expect_select = {*key_fields, order_key, *select}
 180         item = {
 181             key: value
 182             for key, value in fake_item.items()
 183             if key in expect_select
 184         }
 185         list_func = mock.Mock()
 186         list_func().execute = mock.Mock(
 187             side_effect=[
 188                 {'items': [item]},
 189                 {'items': []},
 190                 {'items': []},
 191             ],
 192         )
 193         list_func.reset_mock()
 194         actual = list(arvados.util.keyset_list_all(list_func, order_key, select=list(select), key_fields=key_fields))
 195         self.assertEqual(actual, [item])
 196         calls = list_func.call_args_list
 197         self.assertTrue(len(calls) >= 2, "list_func() not called enough to exhaust items")
 198         for args, kwargs in calls:
 199             self.assertEqual(set(kwargs.get('select', ())), expect_select)
 200
 201
 202 class TestIterStorageClasses:
 203     @pytest.fixture
 204     def mixed_config(self):
 205         return {'StorageClasses': {
 206             'foo': {'Default': False},
 207             'bar': {'Default': True},
 208             'baz': {'Default': True},
 209         }}
 210
 211     @pytest.fixture
 212     def nodef_config(self):
 213         return {'StorageClasses': {
 214             'foo': {'Default': False},
 215             'bar': {'Default': False},
 216         }}
 217
 218     def test_defaults(self, mixed_config):
 219         assert list(arvados.util.iter_storage_classes(mixed_config)) == ['bar', 'baz']
 220
 221     def test_custom_check(self, mixed_config):
 222         assert list(arvados.util.iter_storage_classes(mixed_config, bool)) == ['foo', 'bar', 'baz']
 223
 224     def test_default_fallback(self, nodef_config):
 225         assert list(arvados.util.iter_storage_classes(nodef_config)) == ['default']
 226
 227     def test_custom_fallback(self, nodef_config):
 228         assert list(arvados.util.iter_storage_classes(nodef_config, fallback='fb')) == ['fb']
 229
 230     def test_no_fallback(self, nodef_config):
 231         assert list(arvados.util.iter_storage_classes(nodef_config, fallback='')) == []
 232