src/test/python/apache/aurora/client/api/test_health_check.py

#
# 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

import mox

from apache.aurora.client.api.health_check import HealthCheck, StatusHealthCheck

from gen.apache.aurora.api.ttypes import AssignedTask, ScheduledTask, ScheduleStatus, TaskConfig

PENDING = ScheduleStatus.PENDING
RUNNING = ScheduleStatus.RUNNING
FAILED = ScheduleStatus.FAILED


class HealthCheckTest(unittest.TestCase):

  def setUp(self):
    self._status_health_check = StatusHealthCheck()
    self._health_check_a = mox.MockObject(HealthCheck)
    self._health_check_b = mox.MockObject(HealthCheck)

  def create_task(self, instance_id, task_id, status=RUNNING):
    assigned_task = AssignedTask(taskId=task_id,
                                 instanceId=instance_id,
                                 task=TaskConfig())
    return ScheduledTask(assignedTask=assigned_task, status=status)

  def replay(self):
    mox.Replay(self._health_check_a)
    mox.Replay(self._health_check_b)

  def verify(self):
    mox.Verify(self._health_check_a)
    mox.Verify(self._health_check_b)

  def test_simple_status_health_check(self):
    """Verify that running instances are reported healthy"""
    task_a = self.create_task(0, 'a')
    task_b = self.create_task(1, 'b')
    assert self._status_health_check.health(task_a)
    assert self._status_health_check.health(task_b)

  def test_failed_status_health_check(self):
    """Verify that the health check fails for tasks in a state other than RUNNING"""
    pending_task = self.create_task(0, 'a', status=PENDING)
    failed_task = self.create_task(1, 'b', status=FAILED)
    assert not self._status_health_check.health(pending_task)
    assert not self._status_health_check.health(failed_task)

  def test_changed_task_id(self):
    """Verifes that an instance with a different task id causes the health check to fail"""
    task_a = self.create_task(0, 'a')
    task_b = self.create_task(0, 'b')
    assert self._status_health_check.health(task_a)
    assert not self._status_health_check.health(task_b)