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