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