| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| import unittest |
| from math import ceil |
| |
| import mox |
| |
| from apache.aurora.client.api.health_check import HealthCheck |
| from apache.aurora.client.api.instance_watcher import InstanceWatcher |
| |
| from ...api_util import SchedulerProxyApiSpec |
| |
| from gen.apache.aurora.api.ttypes import ( |
| AssignedTask, |
| JobKey, |
| Response, |
| ResponseCode, |
| ResponseDetail, |
| Result, |
| ScheduledTask, |
| ScheduleStatus, |
| ScheduleStatusResult, |
| TaskConfig, |
| TaskQuery |
| ) |
| |
| |
| class FakeClock(object): |
| def __init__(self): |
| self._now_seconds = 0.0 |
| |
| def time(self): |
| return self._now_seconds |
| |
| def sleep(self, seconds): |
| self._now_seconds += seconds |
| |
| |
| class FakeEvent(object): |
| def __init__(self, clock): |
| self._clock = clock |
| self._is_set = False |
| |
| def wait(self, seconds): |
| self._clock.sleep(seconds) |
| |
| def is_set(self): |
| return self._is_set |
| |
| def set(self): |
| self._is_set = True |
| |
| |
| def find_expected_cycles(period, sleep_secs): |
| return ceil(period / sleep_secs) + 1 |
| |
| |
| class InstanceWatcherTest(unittest.TestCase): |
| WATCH_INSTANCES = range(3) |
| WATCH_SECS = 50 |
| EXPECTED_CYCLES = find_expected_cycles(WATCH_SECS, 3.0) |
| |
| def setUp(self): |
| self._clock = FakeClock() |
| self._event = FakeEvent(self._clock) |
| self._scheduler = mox.MockObject(SchedulerProxyApiSpec) |
| self._job_key = JobKey(role='mesos', name='jimbob', environment='test') |
| self._health_check = mox.MockObject(HealthCheck) |
| self._watcher = InstanceWatcher(self._scheduler, |
| self._job_key, |
| self.WATCH_SECS, |
| health_check_interval_seconds=3, |
| clock=self._clock, |
| terminating_event=self._event) |
| |
| def get_tasks_status_query(self, instance_ids): |
| query = TaskQuery() |
| query.jobKeys = set([self._job_key]) |
| query.statuses = set([ScheduleStatus.RUNNING]) |
| query.instanceIds = set(instance_ids) |
| return query |
| |
| def create_task(self, instance_id): |
| return ScheduledTask( |
| assignedTask=AssignedTask(instanceId=instance_id, task=TaskConfig())) |
| |
| def expect_get_statuses(self, instance_ids=WATCH_INSTANCES, num_calls=EXPECTED_CYCLES): |
| tasks = [self.create_task(instance_id) for instance_id in instance_ids] |
| response = Response( |
| responseCode=ResponseCode.OK, |
| details=[ResponseDetail(message='test')], |
| result=Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks))) |
| |
| query = self.get_tasks_status_query(instance_ids) |
| for _ in range(int(num_calls)): |
| self._scheduler.getTasksWithoutConfigs(query, retry=True).AndReturn(response) |
| |
| def expect_io_error_in_get_statuses(self, instance_ids=WATCH_INSTANCES, |
| num_calls=EXPECTED_CYCLES): |
| |
| query = self.get_tasks_status_query(instance_ids) |
| for _ in range(int(num_calls)): |
| self._scheduler.getTasksWithoutConfigs(query, retry=True).AndRaise(IOError('oops')) |
| |
| def mock_health_check(self, task, status): |
| self._health_check.health(task).InAnyOrder().AndReturn(status) |
| |
| def expect_health_check(self, instance, status, num_calls=EXPECTED_CYCLES): |
| num_calls = num_calls if status else 1 |
| for _ in range(int(num_calls)): |
| self.mock_health_check(self.create_task(instance), status) |
| |
| def assert_watch_result(self, expected_failed_instances, instances_to_watch=WATCH_INSTANCES): |
| instances_returned = self._watcher.watch(instances_to_watch, self._health_check) |
| assert set(expected_failed_instances) == instances_returned, ( |
| 'Expected instances (%s) : Returned instances (%s)' % ( |
| expected_failed_instances, instances_returned)) |
| |
| def replay_mocks(self): |
| mox.Replay(self._scheduler) |
| mox.Replay(self._health_check) |
| |
| def verify_mocks(self): |
| mox.Verify(self._scheduler) |
| mox.Verify(self._health_check) |
| |
| def test_successful_watch(self): |
| """All instances are healthy immediately""" |
| self.expect_get_statuses() |
| self.expect_health_check(0, True) |
| self.expect_health_check(1, True) |
| self.expect_health_check(2, True) |
| self.replay_mocks() |
| self.assert_watch_result([]) |
| self.verify_mocks() |
| |
| def test_single_instance_failure(self): |
| """One failed instance in a batch of instances""" |
| self.expect_get_statuses() |
| self.expect_health_check(0, False) |
| self.expect_health_check(1, True) |
| self.expect_health_check(2, True) |
| self.replay_mocks() |
| self.assert_watch_result([0]) |
| self.verify_mocks() |
| |
| def test_all_instance_failure(self): |
| """All failed instance in a batch of instances""" |
| self.expect_get_statuses(num_calls=1) |
| self.expect_health_check(0, False) |
| self.expect_health_check(1, False) |
| self.expect_health_check(2, False) |
| self.replay_mocks() |
| self.assert_watch_result([0, 1, 2]) |
| self.verify_mocks() |
| |
| def test_watch_period_failure(self): |
| """Instances are reported unhealthy before watch_secs expires""" |
| self.expect_get_statuses() |
| self.expect_health_check(0, True, num_calls=self.EXPECTED_CYCLES - 1) |
| self.expect_health_check(1, True, num_calls=self.EXPECTED_CYCLES - 1) |
| self.expect_health_check(2, True, num_calls=self.EXPECTED_CYCLES - 1) |
| self.expect_health_check(0, False) |
| self.expect_health_check(1, False) |
| self.expect_health_check(2, False) |
| self.replay_mocks() |
| self.assert_watch_result([0, 1, 2]) |
| self.verify_mocks() |
| |
| def test_single_watch_period_failure(self): |
| """One instance is reported unhealthy before watch_secs expires""" |
| self.expect_get_statuses() |
| self.expect_health_check(0, True) |
| self.expect_health_check(1, True) |
| self.expect_health_check(2, True, num_calls=self.EXPECTED_CYCLES - 1) |
| self.expect_health_check(2, False) |
| self.replay_mocks() |
| self.assert_watch_result([2]) |
| self.verify_mocks() |
| |
| def test_terminated_exits_immediately(self): |
| """Terminated instance watched should bail out immediately.""" |
| self.replay_mocks() |
| self._watcher.terminate() |
| result = self._watcher.watch([], self._health_check) |
| assert result is None, ('Expected instances None : Returned instances (%s)' % result) |
| self.verify_mocks() |
