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 class MkdirDashPTest(unittest.TestCase):
  17     def setUp(self):
  18         try:
  19             os.path.mkdir('./tmp')
  20         except:
  21             pass
  22     def tearDown(self):
  23         try:
  24             os.unlink('./tmp/bar')
  25             os.rmdir('./tmp/foo')
  26             os.rmdir('./tmp')
  27         except:
  28             pass
  29     def runTest(self):
  30         arvados.util.mkdir_dash_p('./tmp/foo')
  31         with open('./tmp/bar', 'wb') as f:
  32             f.write(b'bar')
  33         self.assertRaises(OSError, arvados.util.mkdir_dash_p, './tmp/bar')
  34
  35
  36 class RunCommandTestCase(unittest.TestCase):
  37     def test_success(self):
  38         stdout, stderr = arvados.util.run_command(['echo', 'test'],
  39                                                   stderr=subprocess.PIPE)
  40         self.assertEqual("test\n".encode(), stdout)
  41         self.assertEqual("".encode(), stderr)
  42
  43     def test_failure(self):
  44         with self.assertRaises(arvados.errors.CommandFailedError):
  45             arvados.util.run_command(['false'])
  46
  47 class KeysetTestHelper:
  48     def __init__(self, expect):
  49         self.n = 0
  50         self.expect = expect
  51
  52     def fn(self, **kwargs):
  53         if self.expect[self.n][0] != kwargs:
  54             raise Exception("Didn't match %s != %s" % (self.expect[self.n][0], kwargs))
  55         return self
  56
  57     def execute(self, num_retries):
  58         self.n += 1
  59         return self.expect[self.n-1][1]
  60
  61 _SELECT_FAKE_ITEM = {
  62     'uuid': 'zzzzz-zyyyz-zzzzzyyyyywwwww',
  63     'name': 'KeysetListAllTestCase.test_select mock',
  64     'created_at': '2023-08-28T12:34:56.123456Z',
  65 }
  66
  67 class KeysetListAllTestCase(unittest.TestCase):
  68     def test_empty(self):
  69         ks = KeysetTestHelper([[
  70             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
  71             {"items": []}
  72         ]])
  73
  74         ls = list(arvados.util.keyset_list_all(ks.fn))
  75         self.assertEqual(ls, [])
  76
  77     def test_oneitem(self):
  78         ks = KeysetTestHelper([[
  79             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
  80             {"items": [{"created_at": "1", "uuid": "1"}]}
  81         ], [
  82             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", "=", "1"], ["uuid", ">", "1"]]},
  83             {"items": []}
  84         ],[
  85             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">", "1"]]},
  86             {"items": []}
  87         ]])
  88
  89         ls = list(arvados.util.keyset_list_all(ks.fn))
  90         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}])
  91
  92     def test_onepage2(self):
  93         ks = KeysetTestHelper([[
  94             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
  95             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}]}
  96         ], [
  97             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "2"], ["uuid", "!=", "2"]]},
  98             {"items": []}
  99         ]])
 100
 101         ls = list(arvados.util.keyset_list_all(ks.fn))
 102         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}])
 103
 104     def test_onepage3(self):
 105         ks = KeysetTestHelper([[
 106             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
 107             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}, {"created_at": "3", "uuid": "3"}]}
 108         ], [
 109             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "3"], ["uuid", "!=", "3"]]},
 110             {"items": []}
 111         ]])
 112
 113         ls = list(arvados.util.keyset_list_all(ks.fn))
 114         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}, {"created_at": "3", "uuid": "3"}])
 115
 116
 117     def test_twopage(self):
 118         ks = KeysetTestHelper([[
 119             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
 120             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}]}
 121         ], [
 122             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "2"], ["uuid", "!=", "2"]]},
 123             {"items": [{"created_at": "3", "uuid": "3"}, {"created_at": "4", "uuid": "4"}]}
 124         ], [
 125             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "4"], ["uuid", "!=", "4"]]},
 126             {"items": []}
 127         ]])
 128
 129         ls = list(arvados.util.keyset_list_all(ks.fn))
 130         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"},
 131                               {"created_at": "2", "uuid": "2"},
 132                               {"created_at": "3", "uuid": "3"},
 133                               {"created_at": "4", "uuid": "4"}
 134         ])
 135
 136     def test_repeated_key(self):
 137         ks = KeysetTestHelper([[
 138             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": []},
 139             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}, {"created_at": "2", "uuid": "3"}]}
 140         ], [
 141             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "2"], ["uuid", "!=", "3"]]},
 142             {"items": [{"created_at": "2", "uuid": "2"}, {"created_at": "2", "uuid": "4"}]}
 143         ], [
 144             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", "=", "2"], ["uuid", ">", "4"]]},
 145             {"items": []}
 146         ], [
 147             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">", "2"]]},
 148             {"items": [{"created_at": "3", "uuid": "5"}, {"created_at": "4", "uuid": "6"}]}
 149         ], [
 150             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "4"], ["uuid", "!=", "6"]]},
 151             {"items": []}
 152         ],
 153         ])
 154
 155         ls = list(arvados.util.keyset_list_all(ks.fn))
 156         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"},
 157                               {"created_at": "2", "uuid": "2"},
 158                               {"created_at": "2", "uuid": "3"},
 159                               {"created_at": "2", "uuid": "4"},
 160                               {"created_at": "3", "uuid": "5"},
 161                               {"created_at": "4", "uuid": "6"}
 162         ])
 163
 164     def test_onepage_withfilter(self):
 165         ks = KeysetTestHelper([[
 166             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["foo", ">", "bar"]]},
 167             {"items": [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}]}
 168         ], [
 169             {"limit": 1000, "count": "none", "order": ["created_at asc", "uuid asc"], "filters": [["created_at", ">=", "2"], ["uuid", "!=", "2"], ["foo", ">", "bar"]]},
 170             {"items": []}
 171         ]])
 172
 173         ls = list(arvados.util.keyset_list_all(ks.fn, filters=[["foo", ">", "bar"]]))
 174         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}])
 175
 176     def test_onepage_desc(self):
 177         ks = KeysetTestHelper([[
 178             {"limit": 1000, "count": "none", "order": ["created_at desc", "uuid desc"], "filters": []},
 179             {"items": [{"created_at": "2", "uuid": "2"}, {"created_at": "1", "uuid": "1"}]}
 180         ], [
 181             {"limit": 1000, "count": "none", "order": ["created_at desc", "uuid desc"], "filters": [["created_at", "<=", "1"], ["uuid", "!=", "1"]]},
 182             {"items": []}
 183         ]])
 184
 185         ls = list(arvados.util.keyset_list_all(ks.fn, ascending=False))
 186         self.assertEqual(ls, [{"created_at": "2", "uuid": "2"}, {"created_at": "1", "uuid": "1"}])
 187
 188     @parameterized.parameterized.expand(zip(
 189         itertools.cycle(_SELECT_FAKE_ITEM),
 190         itertools.chain.from_iterable(
 191             itertools.combinations(_SELECT_FAKE_ITEM, count)
 192             for count in range(len(_SELECT_FAKE_ITEM) + 1)
 193         ),
 194     ))
 195     def test_select(self, order_key, select):
 196         # keyset_list_all must have both uuid and order_key to function.
 197         # Test that it selects those fields along with user-specified ones.
 198         expect_select = {'uuid', order_key, *select}
 199         item = {
 200             key: value
 201             for key, value in _SELECT_FAKE_ITEM.items()
 202             if key in expect_select
 203         }
 204         list_func = mock.Mock()
 205         list_func().execute = mock.Mock(
 206             side_effect=[
 207                 {'items': [item]},
 208                 {'items': []},
 209                 {'items': []},
 210             ],
 211         )
 212         list_func.reset_mock()
 213         actual = list(arvados.util.keyset_list_all(list_func, order_key, select=list(select)))
 214         self.assertEqual(actual, [item])
 215         calls = list_func.call_args_list
 216         self.assertTrue(len(calls) >= 2, "list_func() not called enough to exhaust items")
 217         for args, kwargs in calls:
 218             self.assertEqual(set(kwargs.get('select', ())), expect_select)