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