services/dockercleaner/tests/test_cleaner.py

   1 #!/usr/bin/env python3
   2
   3 import collections
   4 import itertools
   5 import json
   6 import random
   7 import time
   8 import unittest
   9
  10 import docker
  11 import mock
  12
  13 from arvados_docker import cleaner
  14
  15 MAX_DOCKER_ID = (16 ** 64) - 1
  16
  17 def MockDockerId():
  18     return '{:064x}'.format(random.randint(0, MAX_DOCKER_ID))
  19
  20 def MockContainer(image_hash):
  21     return {'Id': MockDockerId(),
  22             'Image': image_hash['Id']}
  23
  24 def MockImage(*, size=0, vsize=None, tags=[]):
  25     if vsize is None:
  26         vsize = random.randint(100, 2000000)
  27     return {'Id': MockDockerId(),
  28             'ParentId': MockDockerId(),
  29             'RepoTags': list(tags),
  30             'Size': size,
  31             'VirtualSize': vsize}
  32
  33 class MockEvent(dict):
  34     ENCODING = 'utf-8'
  35     event_seq = itertools.count(1)
  36
  37     def __init__(self, status, docker_id=None, **event_data):
  38         if docker_id is None:
  39             docker_id = MockDockerId()
  40         super().__init__(self, **event_data)
  41         self['status'] = status
  42         self['id'] = docker_id
  43         self.setdefault('time', next(self.event_seq))
  44
  45     def encoded(self):
  46         return json.dumps(self).encode(self.ENCODING)
  47
  48
  49 class MockException(docker.errors.APIError):
  50     def __init__(self, status_code):
  51         response = mock.Mock(name='response')
  52         response.status_code = status_code
  53         super().__init__("mock exception", response)
  54
  55
  56 class DockerImageTestCase(unittest.TestCase):
  57     def test_used_at_sets_last_used(self):
  58         image = cleaner.DockerImage(MockImage())
  59         image.used_at(5)
  60         self.assertEqual(5, image.last_used)
  61
  62     def test_used_at_moves_forward(self):
  63         image = cleaner.DockerImage(MockImage())
  64         image.used_at(6)
  65         image.used_at(8)
  66         self.assertEqual(8, image.last_used)
  67
  68     def test_used_at_does_not_go_backward(self):
  69         image = cleaner.DockerImage(MockImage())
  70         image.used_at(4)
  71         image.used_at(2)
  72         self.assertEqual(4, image.last_used)
  73
  74
  75 class DockerImagesTestCase(unittest.TestCase):
  76     def setUp(self):
  77         self.mock_images = []
  78
  79     def setup_mock_images(self, *vsizes):
  80         self.mock_images.extend(MockImage(vsize=vsize) for vsize in vsizes)
  81
  82     def setup_images(self, *vsizes, target_size=1000000):
  83         self.setup_mock_images(*vsizes)
  84         images = cleaner.DockerImages(target_size)
  85         for image in self.mock_images:
  86             images.add_image(image)
  87         return images
  88
  89     def test_has_image(self):
  90         images = self.setup_images(None)
  91         self.assertTrue(images.has_image(self.mock_images[0]['Id']))
  92         self.assertFalse(images.has_image(MockDockerId()))
  93
  94     def test_del_image(self):
  95         images = self.setup_images(None)
  96         images.del_image(self.mock_images[0]['Id'])
  97         self.assertFalse(images.has_image(self.mock_images[0]['Id']))
  98
  99     def test_del_nonexistent_image(self):
 100         images = self.setup_images(None)
 101         images.del_image(MockDockerId())
 102         self.assertTrue(images.has_image(self.mock_images[0]['Id']))
 103
 104     def test_one_image_always_kept(self):
 105         # When crunch-job starts a job, it makes sure each compute node
 106         # has the Docker image loaded, then it runs all the tasks with
 107         # the assumption the image is on each node.  As long as that's
 108         # true, the cleaner should avoid removing every installed image:
 109         # crunch-job might be counting on the most recent one to be
 110         # available, even if it's not currently in use.
 111         images = self.setup_images(None, None, target_size=1)
 112         for use_time, image in enumerate(self.mock_images, 1):
 113             user = MockContainer(image)
 114             images.add_user(user, use_time)
 115             images.end_user(user['Id'])
 116         self.assertEqual([self.mock_images[0]['Id']],
 117                          list(images.should_delete()))
 118
 119     def test_images_under_target_not_deletable(self):
 120         # The images are used in this order.  target_size is set so it
 121         # could hold the largest image, but not after the most recently
 122         # used image is kept; then we have to fall back to the previous one.
 123         images = self.setup_images(20, 30, 40, 10, target_size=45)
 124         for use_time, image in enumerate(self.mock_images, 1):
 125             user = MockContainer(image)
 126             images.add_user(user, use_time)
 127             images.end_user(user['Id'])
 128         self.assertEqual([self.mock_images[ii]['Id'] for ii in [0, 2]],
 129                          list(images.should_delete()))
 130
 131     def test_images_in_use_not_deletable(self):
 132         images = self.setup_images(None, None, target_size=1)
 133         users = [MockContainer(image) for image in self.mock_images]
 134         images.add_user(users[0], 1)
 135         images.add_user(users[1], 2)
 136         images.end_user(users[1]['Id'])
 137         self.assertEqual([self.mock_images[1]['Id']],
 138                          list(images.should_delete()))
 139
 140     def test_image_deletable_after_unused(self):
 141         images = self.setup_images(None, None, target_size=1)
 142         users = [MockContainer(image) for image in self.mock_images]
 143         images.add_user(users[0], 1)
 144         images.add_user(users[1], 2)
 145         images.end_user(users[0]['Id'])
 146         self.assertEqual([self.mock_images[0]['Id']],
 147                          list(images.should_delete()))
 148
 149     def test_image_not_deletable_if_user_restarts(self):
 150         images = self.setup_images(None, target_size=1)
 151         user = MockContainer(self.mock_images[-1])
 152         images.add_user(user, 1)
 153         images.end_user(user['Id'])
 154         images.add_user(user, 2)
 155         self.assertEqual([], list(images.should_delete()))
 156
 157     def test_image_not_deletable_if_any_user_remains(self):
 158         images = self.setup_images(None, target_size=1)
 159         users = [MockContainer(self.mock_images[0]) for ii in range(2)]
 160         images.add_user(users[0], 1)
 161         images.add_user(users[1], 2)
 162         images.end_user(users[0]['Id'])
 163         self.assertEqual([], list(images.should_delete()))
 164
 165     def test_image_deletable_after_all_users_end(self):
 166         images = self.setup_images(None, None, target_size=1)
 167         users = [MockContainer(self.mock_images[ii]) for ii in [0, 1, 1]]
 168         images.add_user(users[0], 1)
 169         images.add_user(users[1], 2)
 170         images.add_user(users[2], 3)
 171         images.end_user(users[1]['Id'])
 172         images.end_user(users[2]['Id'])
 173         self.assertEqual([self.mock_images[-1]['Id']],
 174                          list(images.should_delete()))
 175
 176     def test_images_suggested_for_deletion_by_lru(self):
 177         images = self.setup_images(10, 10, 10, target_size=1)
 178         users = [MockContainer(image) for image in self.mock_images]
 179         images.add_user(users[0], 3)
 180         images.add_user(users[1], 1)
 181         images.add_user(users[2], 2)
 182         for user in users:
 183             images.end_user(user['Id'])
 184         self.assertEqual([self.mock_images[ii]['Id'] for ii in [1, 2]],
 185                          list(images.should_delete()))
 186
 187     def test_adding_user_without_image_does_not_implicitly_add_image(self):
 188         images = self.setup_images(10)
 189         images.add_user(MockContainer(MockImage()), 1)
 190         self.assertEqual([], list(images.should_delete()))
 191
 192     def test_nonexistent_user_removed(self):
 193         images = self.setup_images()
 194         images.end_user('nonexistent')
 195         # No exception should be raised.
 196
 197     def test_del_image_effective_with_users_present(self):
 198         images = self.setup_images(None, target_size=1)
 199         user = MockContainer(self.mock_images[0])
 200         images.add_user(user, 1)
 201         images.del_image(self.mock_images[0]['Id'])
 202         images.end_user(user['Id'])
 203         self.assertEqual([], list(images.should_delete()))
 204
 205     def setup_from_daemon(self, *vsizes, target_size=1500000):
 206         self.setup_mock_images(*vsizes)
 207         docker_client = mock.MagicMock(name='docker_client')
 208         docker_client.images.return_value = iter(self.mock_images)
 209         return cleaner.DockerImages.from_daemon(target_size, docker_client)
 210
 211     def test_images_loaded_from_daemon(self):
 212         images = self.setup_from_daemon(None, None)
 213         for image in self.mock_images:
 214             self.assertTrue(images.has_image(image['Id']))
 215
 216     def test_target_size_set_from_daemon(self):
 217         images = self.setup_from_daemon(20, 10, 5, target_size=15)
 218         user = MockContainer(self.mock_images[-1])
 219         images.add_user(user, 1)
 220         self.assertEqual([self.mock_images[0]['Id']],
 221                          list(images.should_delete()))
 222
 223
 224 class DockerImageUseRecorderTestCase(unittest.TestCase):
 225     TEST_CLASS = cleaner.DockerImageUseRecorder
 226     TEST_CLASS_INIT_KWARGS = {}
 227
 228     def setUp(self):
 229         self.images = mock.MagicMock(name='images')
 230         self.docker_client = mock.MagicMock(name='docker_client')
 231         self.events = []
 232         self.recorder = self.TEST_CLASS(self.images, self.docker_client,
 233                                         self.encoded_events, **self.TEST_CLASS_INIT_KWARGS)
 234
 235     @property
 236     def encoded_events(self):
 237         return (event.encoded() for event in self.events)
 238
 239     def test_unknown_events_ignored(self):
 240         self.events.append(MockEvent('mock!event'))
 241         self.recorder.run()
 242         # No exception should be raised.
 243
 244     def test_fetches_container_on_create(self):
 245         self.events.append(MockEvent('create'))
 246         self.recorder.run()
 247         self.docker_client.inspect_container.assert_called_with(
 248             self.events[0]['id'])
 249
 250     def test_adds_user_on_container_create(self):
 251         self.events.append(MockEvent('create'))
 252         self.recorder.run()
 253         self.images.add_user.assert_called_with(
 254             self.docker_client.inspect_container(), self.events[0]['time'])
 255
 256     def test_unknown_image_handling(self):
 257         # The use recorder should not fetch any images.
 258         self.events.append(MockEvent('create'))
 259         self.recorder.run()
 260         self.assertFalse(self.docker_client.inspect_image.called)
 261
 262     def test_unfetchable_containers_ignored(self):
 263         self.events.append(MockEvent('create'))
 264         self.docker_client.inspect_container.side_effect = MockException(404)
 265         self.recorder.run()
 266         self.assertFalse(self.images.add_user.called)
 267
 268     def test_ends_user_on_container_destroy(self):
 269         self.events.append(MockEvent('destroy'))
 270         self.recorder.run()
 271         self.images.end_user.assert_called_with(self.events[0]['id'])
 272
 273
 274 class DockerImageCleanerTestCase(DockerImageUseRecorderTestCase):
 275     TEST_CLASS = cleaner.DockerImageCleaner
 276
 277     def test_unknown_image_handling(self):
 278         # The image cleaner should fetch and record new images.
 279         self.images.has_image.return_value = False
 280         self.events.append(MockEvent('create'))
 281         self.recorder.run()
 282         self.docker_client.inspect_image.assert_called_with(
 283             self.docker_client.inspect_container()['Image'])
 284         self.images.add_image.assert_called_with(
 285             self.docker_client.inspect_image())
 286
 287     def test_unfetchable_images_ignored(self):
 288         self.images.has_image.return_value = False
 289         self.docker_client.inspect_image.side_effect = MockException(404)
 290         self.events.append(MockEvent('create'))
 291         self.recorder.run()
 292         self.docker_client.inspect_image.assert_called_with(
 293             self.docker_client.inspect_container()['Image'])
 294         self.assertFalse(self.images.add_image.called)
 295
 296     def test_deletions_after_destroy(self):
 297         delete_id = MockDockerId()
 298         self.images.should_delete.return_value = [delete_id]
 299         self.events.append(MockEvent('destroy'))
 300         self.recorder.run()
 301         self.docker_client.remove_image.assert_called_with(delete_id)
 302         self.images.del_image.assert_called_with(delete_id)
 303
 304     def test_failed_deletion_handling(self):
 305         delete_id = MockDockerId()
 306         self.images.should_delete.return_value = [delete_id]
 307         self.docker_client.remove_image.side_effect = MockException(500)
 308         self.events.append(MockEvent('destroy'))
 309         self.recorder.run()
 310         self.docker_client.remove_image.assert_called_with(delete_id)
 311         self.assertFalse(self.images.del_image.called)
 312
 313
 314 class DockerContainerCleanerTestCase(DockerImageUseRecorderTestCase):
 315     TEST_CLASS = cleaner.DockerImageCleaner
 316     TEST_CLASS_INIT_KWARGS = {'remove_containers_onexit': True}
 317
 318     @mock.patch('arvados_docker.cleaner.logger')
 319     def test_failed_container_deletion_handling(self, mockLogger):
 320         cid = MockDockerId()
 321         self.docker_client.remove_container.side_effect = MockException(500)
 322         self.events.append(MockEvent('die', docker_id=cid))
 323         self.recorder.run()
 324         self.docker_client.remove_container.assert_called_with(cid)
 325         self.assertEqual("Failed to remove container %s: %s",
 326                          mockLogger.warning.call_args[0][0])
 327         self.assertEqual(cid,
 328                          mockLogger.warning.call_args[0][1])
 329
 330
 331 class HumanSizeTestCase(unittest.TestCase):
 332     def check(self, human_str, count, exp):
 333         self.assertEqual(count * (1024 ** exp),
 334                          cleaner.human_size(human_str))
 335
 336     def test_bytes(self):
 337         self.check('1', 1, 0)
 338         self.check('82', 82, 0)
 339
 340     def test_kibibytes(self):
 341         self.check('2K', 2, 1)
 342         self.check('3k', 3, 1)
 343
 344     def test_mebibytes(self):
 345         self.check('4M', 4, 2)
 346         self.check('5m', 5, 2)
 347
 348     def test_gibibytes(self):
 349         self.check('6G', 6, 3)
 350         self.check('7g', 7, 3)
 351
 352     def test_tebibytes(self):
 353         self.check('8T', 8, 4)
 354         self.check('9t', 9, 4)
 355
 356
 357 class RunTestCase(unittest.TestCase):
 358     def setUp(self):
 359         self.args = mock.MagicMock(name='args')
 360         self.args.quota = 1000000
 361         self.docker_client = mock.MagicMock(name='docker_client')
 362
 363     def test_run(self):
 364         test_start_time = int(time.time())
 365         self.docker_client.events.return_value = []
 366         cleaner.run(self.args, self.docker_client)
 367         self.assertEqual(2, self.docker_client.events.call_count)
 368         event_kwargs = [args[1] for args in
 369                         self.docker_client.events.call_args_list]
 370         self.assertIn('since', event_kwargs[0])
 371         self.assertIn('until', event_kwargs[0])
 372         self.assertLessEqual(test_start_time, event_kwargs[0]['until'])
 373         self.assertIn('since', event_kwargs[1])
 374         self.assertEqual(event_kwargs[0]['until'], event_kwargs[1]['since'])
 375
 376
 377 @mock.patch('docker.Client', name='docker_client')
 378 @mock.patch('arvados_docker.cleaner.run', name='cleaner_run')
 379 class MainTestCase(unittest.TestCase):
 380     def test_client_api_version(self, run_mock, docker_client):
 381         cleaner.main(['--quota', '1000T'])
 382         self.assertEqual(1, docker_client.call_count)
 383         # 1.14 is the first version that's well defined, going back to
 384         # Docker 1.2, and still supported up to at least Docker 1.9.
 385         # See <https://docs.docker.com/engine/reference/api/docker_remote_api/>.
 386         self.assertEqual('1.14',
 387                          docker_client.call_args[1].get('version'))
 388         self.assertEqual(1, run_mock.call_count)
 389         self.assertIs(run_mock.call_args[0][1], docker_client())
 390
 391
 392 class ContainerRemovalTestCase(unittest.TestCase):
 393     LIFECYCLE = ['create', 'attach', 'start', 'resize', 'die', 'destroy']
 394
 395     def setUp(self):
 396         self.args = mock.MagicMock(name='args')
 397         self.docker_client = mock.MagicMock(name='docker_client')
 398         self.existingCID = MockDockerId()
 399         self.docker_client.containers.return_value = [{
 400             'Id': self.existingCID,
 401             'Status': 'Exited (0) 6 weeks ago',
 402         }, {
 403             # If docker_client.containers() returns non-exited
 404             # containers for some reason, do not remove them.
 405             'Id': MockDockerId(),
 406             'Status': 'Running',
 407         }]
 408         self.newCID = MockDockerId()
 409         self.docker_client.events.return_value = [
 410             MockEvent(e, docker_id=self.newCID).encoded()
 411             for e in self.LIFECYCLE]
 412
 413     def test_remove_onexit(self):
 414         self.args.remove_stopped_containers = 'onexit'
 415         cleaner.run(self.args, self.docker_client)
 416         self.docker_client.remove_container.assert_called_once_with(self.newCID)
 417
 418     def test_remove_always(self):
 419         self.args.remove_stopped_containers = 'always'
 420         cleaner.run(self.args, self.docker_client)
 421         self.docker_client.remove_container.assert_any_call(self.existingCID)
 422         self.docker_client.remove_container.assert_any_call(self.newCID)
 423         self.assertEqual(2, self.docker_client.remove_container.call_count)
 424
 425     def test_remove_never(self):
 426         self.args.remove_stopped_containers = 'never'
 427         cleaner.run(self.args, self.docker_client)
 428         self.assertEqual(0, self.docker_client.remove_container.call_count)
 429
 430     def test_container_exited_between_subscribe_events_and_check_existing(self):
 431         self.args.remove_stopped_containers = 'always'
 432         self.docker_client.events.return_value = [
 433             MockEvent(e, docker_id=self.existingCID).encoded()
 434             for e in ['die', 'destroy']]
 435         cleaner.run(self.args, self.docker_client)
 436         # Subscribed to events before getting the list of existing
 437         # exited containers?
 438         self.docker_client.assert_has_calls([
 439             mock.call.events(since=mock.ANY),
 440             mock.call.containers(filters={'status':'exited'})])
 441         # Asked to delete the container twice?
 442         self.docker_client.remove_container.assert_has_calls([mock.call(self.existingCID)] * 2)
 443         self.assertEqual(2, self.docker_client.remove_container.call_count)