# -*- coding: utf-8 -*-
# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Unit tests for workqueue server and client classes."""

#pylint: disable=protected-access

from __future__ import print_function

import mock
import os
import pickle
import shutil
import tempfile
import unittest

from chromite.lib.workqueue import service
from chromite.lib.workqueue import tasks


REQUESTED = service._BaseWorkQueue._REQUESTED
PENDING = service._BaseWorkQueue._PENDING
RUNNING = service._BaseWorkQueue._RUNNING
COMPLETE = service._BaseWorkQueue._COMPLETE
ABORTING = service._BaseWorkQueue._ABORTING
STATES = service._BaseWorkQueue._STATES


class _MonitorTermination(Exception):
  """Exeception raised to stop `ProcessRequests` during testing.

  `service.WorkQueueServer.ProcessRequests` can only terminate
  if an exception is raised during its execution.  Unit tests that
  call the method raise this exception to end the calls.
  """
  pass


class _Task(object):
  """Information about an enqueued task."""
  def __init__(self, request_id, duration):
    self.request_id = request_id
    self.duration = duration
    self.start_time = None
    self.age = 0
    self.aborted = False


class _TestTaskManager(tasks.TaskManager):
  """A concrete subclass of `TaskManager` for testing purposes.

  This is a test double class used to track usage during tests
  of `WorkQueueServer.ProcessRequests()`.  It also serves a
  very limited role in is some other specific test cases.

  This version of `TaskManager` keeps track of the sequence of events
  that trigger operations on the various requests that are expected to
  make their way through the test cases.  These are the events
  supported:
    * Enqueue(task_id, duration) - enqueue a request for a given
      `task_id`.  The task will be marked ready for `Reap()` after
      `duration` ticks.
    * SetCapacity(capacity) - `self.HasCapacity()` will return true
      until `capacity` tasks have been started by `self.StartTask()`.
    * Abort(task_id) - Abort the request associated with the given
      `task_id`.
    * Execute(tick_count) - Resume processing events after `tick_count`
      ticks have elapsed.

  The `SetSequence()` method is used to initialize the sequence of
  events before a call to `ProcessRequests()`.  For example:

    sequence = [(task_manager.SetCapacity, 1),
                (task_manager.Enqueue, 'a', 1),
                (task_manager.Execute, 1),
                (task_manager.Abort, 'a')]
    task_manager.SetSequence(sequence)

  The sequence above has the following effect:
    * The `SetCapacity` event means that one enqueue request will be
      allowed to start.
    * The `Enqueue` request will actually enqueue a request for
      execution.
    * The `Execute` event will cause `ProcessRequests()` to execute for
      one tick.  This will allow task 'a' to begin executing.
    * The `Abort` event will execute 1 tick after task 'a' is enqueued,
      which will cause the request to be aborted while running..

  This class is responsible for the following assertions:
    * `ProcessRequests()` will not call `time.sleep()` without an
      intervening call to `manager.StartTick()`.
    * `ProcessRequests()` will not call without an `manager.StartTick()`
      intervening call to `time.sleep()`.
    * Calls to `time.sleep()` use the value of `self.sample_interval`.
    * Requests passed to `StartTask()` use the correct argument for the
      given request.
  """

  def __init__(self, test_cases=None):
    super(_TestTaskManager, self).__init__(None, 1e6)
    self.requested_tasks = {}
    self.running_tasks = set()
    self._test_cases = test_cases
    self._sequence_data = None
    self._delay_count = 0
    self._capacity_count = 0
    self._start_tick_allowed = True
    self._sleep_allowed = True
    self._current_tick = 0

  def __len__(self):
    return len(self.running_tasks)

  def SetSequence(self, sequence_data):
    """Initialize a sequence of events for `ProcessRequests()`."""
    self._start_tick_allowed = True
    self._sleep_allowed = True
    self._sequence_data = list(sequence_data)

  def Enqueue(self, task_id, duration):
    """Perform handling for an `Enqueue` event."""
    request_id = self._test_cases.client.EnqueueRequest(task_id)
    self.requested_tasks[task_id] = _Task(request_id, duration)

  def SetCapacity(self, capacity):
    """Perform handling for a `SetCapacity` event."""
    self._capacity_count = capacity

  def Abort(self, task_id):
    """Perform handling for an `Abort` event."""
    request_id = self.requested_tasks[task_id].request_id
    self._test_cases.client.AbortRequest(request_id)
    self.requested_tasks[task_id].aborted = True

  def Execute(self, tick_count):
    """Perform handling for an `Execute` event."""
    self._delay_count = tick_count

  def _SetTickReady(self, start_ready):
    """Set whether `StartTick()` or `Sleep()` should be called next."""
    self._start_tick_allowed = start_ready
    self._sleep_allowed = not start_ready

  def Sleep(self, interval):
    """Replacement for `time.sleep()` in `ProcessRequests()`."""
    self._test_cases.assertTrue(self._sleep_allowed,
                                'time.sleep() called twice without '
                                'intervening call to TaskManager.StartTick()')
    self._SetTickReady(True)
    self._test_cases.assertEqual(interval, self.sample_interval)
    self._current_tick += 1

  def StartTick(self):
    self._test_cases.assertTrue(self._start_tick_allowed,
                                'TaskManager.StartTick() called twice '
                                'without intervening call to time.sleep()')
    self._SetTickReady(False)
    # Execute event from our sequence until an `Execute` event says to
    # proceed.
    while self._delay_count <= 0:
      if not self._sequence_data:
        raise _MonitorTermination()
      next_event = self._sequence_data.pop(0)
      next_event[0](*next_event[1:])
    self._delay_count -= 1
    # Advance the age of all running tasks by one tick.
    for _, task_id in self.running_tasks:
      self.requested_tasks[task_id].age += 1

  def HasCapacity(self):
    return bool(self._capacity_count)

  def StartTask(self, request_id, request_data):
    if self._test_cases:
      task = self.requested_tasks[request_data]
      task.start_time = self._current_tick
      self._test_cases.assertEqual(request_id, task.request_id)
    self._capacity_count -= 1
    self.running_tasks.add((request_id, request_data))

  def TerminateTask(self, request_id):
    for entry in self.running_tasks.copy():
      if request_id == entry[0]:
        self.running_tasks.remove(entry)

  def Reap(self):
    for entry in self.running_tasks.copy():
      task = self.requested_tasks[entry[1]]
      if task.age >= task.duration:
        self.running_tasks.remove(entry)
        yield entry


class _BaseServiceTestCases(unittest.TestCase):
  """Base class for testing workqueue server and client methods."""

  def setUp(self):
    self._spool_tmp = tempfile.mkdtemp()
    spool_dir = os.path.join(self._spool_tmp, 'spool')
    self.server = service.WorkQueueServer(spool_dir)
    self.client = service.WorkQueueClient(spool_dir)

  def tearDown(self):
    shutil.rmtree(self._spool_tmp)

  def ValidateState(self, request_id, *args):
    """Check that a given request is in an expected state.

    Assert the following:
      * The spool file for `request_id` exists for all states in
        `args`.
      * No other spool files exist for `request_id`.

    Args:
      request_id: The request id to be tested.
      args: Allowed states of the request.
    """
    actual_states = set()
    for state in STATES:
      if self.server._RequestInState(request_id, state):
        actual_states.add(state)
    self.assertEqual(actual_states, set(args))


class ServiceInternalsTestCases(_BaseServiceTestCases):
  """Test cases for workqueue state transitions.

  These test cases exercise part of all of the life-cycle of a single
  request.  As tested, the full life-cycle follows these steps:
   1. Client creates a request via `EnqueueRequest()`.
   2. Server sees the request via `_GetNewRequests()`.
   3. Server starts a task for a request via `_StartRequest()`.
   4. Server reports request completion via `_CompleteRequest()`.
   5. Client removes completed request via `Wait()`.

  In methods and comments below, the first four stages are identified by
  the names `Enqueue`, `GetNew`, `Start`, and `Complete`, respectively.
  """

  def setUp(self):
    super(ServiceInternalsTestCases, self).setUp()
    self.server._CreateSpool()

  def _ValidateRequestState(self, request_id, value, expected_state):
    """Check that a given request is in an expected state.

    Assert the following:
      * The spool file associated with `request_id` exists for
        `expected_state`.
      * The spool file for the expected state contains `value`.
      * There is no spool file for `request_id` for any state
        other than `expected_state`.

    Args:
      request_id: The request id to be tested.
      value: The value expected to be stored in the request's state
        file.
      expected_state: The expected state of the request.
    """
    self.ValidateState(request_id, expected_state)
    request_path = self.server._GetRequestPathname(request_id,
                                                   expected_state)
    with open(request_path, 'r') as f:
      self.assertEqual(pickle.load(f), value)

  def _ValidateAborted(self, request_id, expected_state):
    """Check that an abort request has been properly recorded.

    The code calls `_GetAbortRequests()` and asserts that the returned
    value matches the given `request_id`.

    Prior to the call, assert the following:
      * The abort request file exists for `request_id`.
      * The request file remains in `expected_state`.

    After to the call, assert the following:
      * The abort request file no longer exists.
      * The request file remains in `expected_state`.

    Args:
      request_id: The request to be aborted.
      expected_state: The state of the request at the time of abort.
    """
    self.ValidateState(request_id, expected_state, ABORTING)
    abort_list = list(self.server._GetAbortRequests())
    self.assertEqual([request_id], abort_list)
    self.ValidateState(request_id, expected_state)

  def _ExecuteAbort(self, request_id, expected_state):
    """Execute a call to `AbortRequest()`.

    Makes a call to `self.client.AbortRequest()`, and then asserts the
    the conditions required for `_ValidateAborted()`, above.

    Args:
      request_id: The request to be aborted.
      expected_state: The state of the request at the time of abort.
    """
    self.client.AbortRequest(request_id)
    self._ValidateAborted(request_id, expected_state)

  def _ExecuteTimeout(self, request_id, expected_state):
    """Make a call to `Wait()` that is expected to time out.

    Assert the following:
      * A timeout exception is raised by `Wait()`.
      * After the call, the request meets the conditions required for
        `_ValidateAborted()`, above.

    Args:
      request_id: The request that will time out.
      expected_state: The state of the request at timeout.
    """
    with self.assertRaises(service.WorkQueueTimeout):
      self.client.Wait(request_id, 0)
    self._ValidateAborted(request_id, expected_state)

  def _ExecuteEnqueue(self, request_value):
    """Walk a request to completion of the `Enqueue` stage.

    At each stage of execution, and at the end, asserts that
    the request is in the state expected for the stage.

    Args:
      request_value: Initial value to pass to `EnqueueRequest`.
    """
    request_id = self.client.EnqueueRequest(request_value)
    self._ValidateRequestState(request_id, request_value,
                               self.client._REQUESTED)
    return request_id

  def _ExecuteGetNew(self, request_value):
    """Walk a request to completion of the `GetNew` stage.

    At each stage of execution, and at the end, asserts that the request
    is in the state expected for the stage.

    Args:
      request_value: Initial value to pass to `EnqueueRequest`.
    """
    request_id = self._ExecuteEnqueue(request_value)
    new_requests = self.server._GetNewRequests()
    self.assertEqual([request_id], new_requests)
    self._ValidateRequestState(request_id, request_value, PENDING)
    return request_id

  def _ExecuteStart(self, request_value):
    """Walk a request to completion of the `Start` stage.

    At each stage of execution, and at the end, asserts that the request
    is in the state expected for the stage.

    After the call to `_StartRequest()` assert that the expected
    side-effects occurred on the method's parameters.

    Args:
      request_value: Initial value to pass to `EnqueueRequest`.
    """
    request_id = self._ExecuteGetNew(request_value)
    manager = _TestTaskManager()
    self.server._StartRequest(request_id, manager)
    self.assertEqual(manager.running_tasks,
                     set([(request_id, request_value)]))
    self._ValidateRequestState(request_id, request_value, RUNNING)
    return request_id

  def _ExecuteComplete(self, request_value, result_value):
    """Walk a request to completion of the `Complete` stage.

    At each stage of execution, but not at the end, asserts that the
    request is in the state expected for the stage.

    Args:
      request_value: Initial value to pass to `EnqueueRequest`.
      result_value: the return value for the request.
    """
    request_id = self._ExecuteStart(request_value)
    self.server._CompleteRequest(request_id, result_value)
    return request_id

  def _ExecuteLifecycle(self, request_value, result_value):
    """Walk a request through the complete request life cycle.

    This creates a request, and processes all standard state
    transitions, ending with a call to `Wait()`.  The initial
    request object is passed in as `request_value`, and the
    result returned should be `result_value`.

    Args:
      request_value: the initial value to be passed to
        `EnqueueRequest()`.
      result_value: The value to be given the result in the call to
        `_CompleteRequest()`.
    """
    request_id = self._ExecuteComplete(request_value, result_value)
    return self.client.Wait(request_id, 0)

  def _EnqueueMultiple(self, count):
    """Make sequential calls to `EnqueueRequest()`.

    Make `count` calls to `EnqueueRequest()`, then assert the following:
      * The returned request id values are unique and in sorted order.
      * After all calls are complete, all requests are in the
        `_REQUESTED` state.
    """
    request_value = 'blurbity-blurb'
    request_ids = []
    for _ in range(0, count):
      request_ids.append(self.client.EnqueueRequest(request_value))
    for index in range(0, len(request_ids) - 1):
      self.assertLess(request_ids[index], request_ids[index+1],
                      'request id out of order at index '
                      '{:d}'.format(index))
    for request_id in request_ids:
      self._ValidateRequestState(request_id, request_value,
                                 self.client._REQUESTED)

  def testEnqueueMultiple(self):
    """Test sequential calls to `EnqueueRequest()`.

    Make multiple calls to `EnqueueRequest()` in sequence, and assert
    the conditions detailed in `_EnqueueMultiple()`, above.

    Test once with the standard value of `_REQUEST_ID_FORMAT`, and once
    with a patched value that by design forces id collisions between
    successive calls to `EnqueueRequest()`.
    """
    self._EnqueueMultiple(5)
    with mock.patch.object(self.client,
                           '_REQUEST_ID_FORMAT', new='{:.2f}'):
      self._EnqueueMultiple(5)

  def testAbortEnqueue(self):
    """Test aborting a request after `Enqueue`."""
    request_id = self._ExecuteEnqueue('To be or not to be')
    self._ExecuteAbort(request_id, self.client._REQUESTED)

  def testAbortGetNew(self):
    """Test aborting a request after `GetNew`."""
    request_id = self._ExecuteGetNew('To be or not to be')
    self._ExecuteAbort(request_id, self.client._PENDING)

  def testAbortStart(self):
    """Test aborting a request after `Start`."""
    request_id = self._ExecuteStart('To be or not to be')
    self._ExecuteAbort(request_id, self.client._RUNNING)

  def testAbortComplete(self):
    """Test aborting a request after `Complete`."""
    request_id = self._ExecuteComplete('To be or not to be',
                                       'That is the question')
    self._ExecuteAbort(request_id, self.client._COMPLETE)

  def testTimeoutEnqueue(self):
    """Test a call to `Wait()` that times out after `Enqueue`."""
    request_id = self._ExecuteEnqueue('To be or not to be')
    self._ExecuteTimeout(request_id, self.client._REQUESTED)

  def testTimeoutGetNew(self):
    """Test a call to `Wait()` that times out after `GetNew`."""
    request_id = self._ExecuteGetNew('To be or not to be')
    self._ExecuteTimeout(request_id, PENDING)

  def testTimeoutStart(self):
    """Test a call to `Wait()` that times out after `Start`."""
    request_id = self._ExecuteStart('To be or not to be')
    self._ExecuteTimeout(request_id, RUNNING)

  def testTransitionRequest(self):
    """Test `WorkQueueServer._TransitionRequest()`."""
    request_value = 'blurbity-blurb'
    request_id = self.client.EnqueueRequest(request_value)
    oldstate = REQUESTED
    newstate = PENDING
    self.server._TransitionRequest(request_id, oldstate, newstate)
    self._ValidateRequestState(request_id, request_value, newstate)

  def testGetNewMultiple(self):
    """Test `_GetNewRequests()` with multiple requests pending."""
    request_values = [str(i) for i in range(0, 4)]
    request_list = []
    for rq in request_values:
      request_list.append(self.client.EnqueueRequest(rq))
    new_requests = self.server._GetNewRequests()
    self.assertEqual(request_list, new_requests)
    for request_id, value in zip(request_list, request_values):
      self._ValidateRequestState(request_id, value,
                                 PENDING)

  def testRequestLifecycle(self):
    """Test the complete life cycle of a successful request."""
    request_value = 'blurbity-blurb'
    expected_result = 'not blurbity at all'
    actual_result = self._ExecuteLifecycle(request_value,
                                           expected_result)
    self.assertEqual(actual_result, expected_result)
    for state_dir in self.server._spool_state_dirs.itervalues():
      self.assertFalse(os.listdir(state_dir))

  def testRequestLifecycleFailure(self):
    """Test the complete life cycle of a request that fails."""
    request_value = 'blurbity-blurb'
    expected_result = Exception('still blurbity!')
    with self.assertRaises(Exception) as exc:
      self._ExecuteLifecycle(request_value,
                             expected_result)
    # Different exceptions won't compare equal even if constructed with
    # the same arguments, so compare the string values instead.
    self.assertEqual(str(exc.exception), str(expected_result))


class ServerProcessRequestsTestCases(_BaseServiceTestCases):
  """Test cases for `WorkQueueServer.ProcessRequests()`."""

  def setUp(self):
    super(ServerProcessRequestsTestCases, self).setUp()
    self._manager = _TestTaskManager(self)

  def _ValidateCompleted(self, task, task_id):
    """Assert that the given `task` completed successfully."""
    value = self.client.Wait(task.request_id, 0)
    self.assertEqual(task_id, value)
    self.assertEqual(task.age, task.duration)

  def _ValidateAborted(self, task):
    """Assert that the given `task` was aborted."""
    request_id = task.request_id
    self.ValidateState(request_id)

  def _RunSequence(self, *args):
    """Run `ProcessRequests()` with the given sequence of events."""
    self._manager.SetSequence(args)
    with self.assertRaises(_MonitorTermination):
      with mock.patch('time.sleep', wraps=self._manager.Sleep):
        self.server.ProcessRequests(self._manager)
    for task_id, task in self._manager.requested_tasks.iteritems():
      if task.aborted:
        self._ValidateAborted(task)
      else:
        self._ValidateCompleted(task, task_id)
    return self._manager.requested_tasks

  def testDoNothing(self):
    for capacity in range(2):
      task_results = self._RunSequence(
          (self._manager.SetCapacity, capacity),
          (self._manager.Execute, 2))
      self.assertFalse(task_results)

  def testImmediateCapacity(self):
    task_results = self._RunSequence(
        (self._manager.SetCapacity, 1),
        (self._manager.Enqueue, 'a', 1),
        (self._manager.Execute, 2))
    task = task_results['a']
    self.assertEqual(task.start_time, 0)
    self.assertEqual(task.age, task.duration)

  def testDelayedCapacity(self):
    task_results = self._RunSequence(
        (self._manager.SetCapacity, 0),
        (self._manager.Enqueue, 'a', 1),
        (self._manager.Execute, 1),
        (self._manager.SetCapacity, 1),
        (self._manager.Execute, 2))
    task = task_results['a']
    self.assertEqual(task.start_time, 1)
    self.assertEqual(task.age, task.duration)

  def testStagedCapacity(self):
    task_results = self._RunSequence(
        (self._manager.SetCapacity, 1),
        (self._manager.Enqueue, 'a', 1),
        (self._manager.Enqueue, 'b', 1),
        (self._manager.Execute, 1),
        (self._manager.SetCapacity, 1),
        (self._manager.Execute, 2))
    for task_id, start in [('a', 0), ('b', 1)]:
      task = task_results[task_id]
      self.assertEqual(task.start_time, start)
      self.assertEqual(task.age, task.duration)

  def testAbortRequested(self):
    task_results = self._RunSequence(
        (self._manager.SetCapacity, 1),
        (self._manager.Enqueue, 'a', 1),
        (self._manager.Abort, 'a'),
        (self._manager.Execute, 2))
    task = task_results['a']
    self.assertTrue(task.aborted)
    self.assertIsNone(task.start_time)

  def testAbortPending(self):
    task_results = self._RunSequence(
        (self._manager.SetCapacity, 0),
        (self._manager.Enqueue, 'a', 1),
        (self._manager.Execute, 1),
        (self._manager.SetCapacity, 1),
        (self._manager.Abort, 'a'),
        (self._manager.Execute, 2))
    task = task_results['a']
    self.assertTrue(task.aborted)
    self.assertIsNone(task.start_time)

  def testAbortRunning(self):
    task_results = self._RunSequence(
        (self._manager.SetCapacity, 1),
        (self._manager.Enqueue, 'a', 1),
        (self._manager.Execute, 1),
        (self._manager.Abort, 'a'),
        (self._manager.Execute, 1))
    task = task_results['a']
    self.assertTrue(task.aborted)
    self.assertEqual(task.start_time, 0)

  def testAbortComplete(self):
    task_results = self._RunSequence(
        (self._manager.SetCapacity, 1),
        (self._manager.Enqueue, 'a', 1),
        (self._manager.Execute, 2),
        (self._manager.Abort, 'a'),
        (self._manager.Execute, 1))
    task = task_results['a']
    self.assertTrue(task.aborted)
    self.assertEqual(task.age, task.duration)


if __name__ == '__main__':
  unittest.main()