#
# 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 contextlib
import unittest

import pytest
from mock import Mock, call, patch

from apache.aurora.client.api.job_monitor import JobMonitor
from apache.aurora.client.cli import Context
from apache.aurora.client.cli.client import AuroraCommandLine
from apache.aurora.client.cli.jobs import KillAllJobCommand, KillCommand
from apache.aurora.client.cli.options import TaskInstanceKey, parse_instances

from .util import AuroraClientCommandTest, FakeAuroraCommandContext, mock_verb_options

from gen.apache.aurora.api.constants import ACTIVE_STATES
from gen.apache.aurora.api.ttypes import (
    JobKey,
    ResponseCode,
    Result,
    ScheduleStatus,
    ScheduleStatusResult,
    TaskQuery
)


class TestInstancesParser(unittest.TestCase):
  def test_parse_instances(self):
    instances = '0,1-3,5'
    x = parse_instances(instances)
    assert x == [0, 1, 2, 3, 5]

  def test_parse_none(self):
    assert parse_instances(None) is None
    assert parse_instances("") is None


class TestKillCommand(AuroraClientCommandTest):

  def test_kill_updating_error_nobatch(self):
    """Verify that the no batch code path correctly includes the lock error message."""
    command = KillCommand()
    mock_options = mock_verb_options(command)
    mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [])
    mock_options.no_batching = True

    fake_context = FakeAuroraCommandContext()
    fake_context.set_options(mock_options)
    fake_context.add_expected_query_result(
      self.create_query_call_result(), job_key=self.TEST_JOBKEY)
    mock_api = fake_context.get_api('test')
    mock_api.kill_job.return_value = AuroraClientCommandTest.create_blank_response(
      ResponseCode.JOB_UPDATING_ERROR, "Error.")

    with pytest.raises(Context.CommandError):
      command.execute(fake_context)

    mock_api.kill_job.assert_called_once_with(
        self.TEST_JOBKEY,
        mock_options.instance_spec.instance,
        config=None,
        message=None)

    self.assert_lock_message(fake_context)

  def test_kill_updating_error_batches(self):
    """Verify that the batch kill path short circuits and includes the lock error message."""
    command = KillCommand()
    mock_options = mock_verb_options(command)
    mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [1])
    mock_options.no_batching = False

    fake_context = FakeAuroraCommandContext()
    fake_context.set_options(mock_options)

    fake_context.add_expected_query_result(
      self.create_query_call_result(), job_key=self.TEST_JOBKEY)
    fake_context.add_expected_query_result(
      AuroraClientCommandTest.create_query_call_result(
        AuroraClientCommandTest.create_scheduled_task(1, ScheduleStatus.RUNNING)))

    mock_api = fake_context.get_api('test')
    mock_api.kill_job.return_value = AuroraClientCommandTest.create_blank_response(
      ResponseCode.JOB_UPDATING_ERROR, "Error.")

    with pytest.raises(Context.CommandError):
      command.execute(fake_context)

    mock_api.kill_job.assert_called_once_with(
        self.TEST_JOBKEY,
        mock_options.instance_spec.instance,
        config=None,
        message=None)

    self.assert_lock_message(fake_context)

  def test_kill_inactive_instance_spec(self):
    """Verify the instance spec is validated in a batched kill."""
    command = KillCommand()
    mock_options = mock_verb_options(command)
    mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [1])
    mock_options.no_batching = False
    mock_options.strict = True

    fake_context = FakeAuroraCommandContext()
    fake_context.set_options(mock_options)

    fake_context.add_expected_query_result(
      self.create_query_call_result(), job_key=self.TEST_JOBKEY)
    fake_context.add_expected_query_result(AuroraClientCommandTest.create_empty_task_result())

    with pytest.raises(Context.CommandError) as e:
      command.execute(fake_context)
    assert e.value.message == "Invalid instance parameter: [1]"

  def test_kill_batched_queries_active_instances(self):
    """Verify that the batch kill operates on active instances only."""
    command = KillCommand()
    mock_options = mock_verb_options(command)
    mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [1])
    mock_options.no_batching = False

    fake_context = FakeAuroraCommandContext()
    fake_context.set_options(mock_options)

    fake_context.add_expected_query_result(
      self.create_query_call_result(), job_key=self.TEST_JOBKEY)
    fake_context.add_expected_query_result(AuroraClientCommandTest.create_empty_task_result())

    command.execute(fake_context)
    message = "No tasks to kill found for job %s" % self.TEST_JOBKEY.to_path()
    assert fake_context.get_err()[0] == message

  def test_kill_opens_url(self):
    """Verify the kill commands opens the job page if requested"""
    command = KillCommand()
    mock_options = mock_verb_options(command)
    mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [1])
    mock_options.open_browser = True

    fake_context = FakeAuroraCommandContext()
    fake_context.set_options(mock_options)

    fake_context.add_expected_query_result(
      self.create_query_call_result(), job_key=self.TEST_JOBKEY)
    fake_context.add_expected_query_result(AuroraClientCommandTest.create_empty_task_result())

    command.execute(fake_context)

    assert self.mock_webbrowser.mock_calls == [
        call("http://something_or_other/scheduler/bozo/test/hello")
    ]

  def test_kill_nobatch_passes_config(self):
    """Verify that the no batch code path correctly passes job config to the api."""
    command = KillCommand()
    config = self.get_valid_config()

    mock_options = mock_verb_options(command)
    mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [])
    mock_options.no_batching = True
    mock_options.config = config

    fake_context = FakeAuroraCommandContext()
    fake_context.set_options(mock_options)
    fake_context.add_config(config)

    mock_api = fake_context.get_api('test')
    fake_context.add_expected_query_result(
      self.create_query_call_result(), job_key=self.TEST_JOBKEY)
    fake_context.add_expected_query_result(AuroraClientCommandTest.create_empty_task_result())
    mock_api.kill_job.return_value = self.create_simple_success_response()

    command.execute(fake_context)

    assert mock_api.kill_job.mock_calls == [
        call(self.TEST_JOBKEY, mock_options.instance_spec.instance, config=config, message=None)
    ]

  def test_kill_batched_passes_config(self):
    """Verify that the batched code path correctly passes job config to the api."""
    command = KillCommand()
    config = self.get_valid_config()

    mock_options = mock_verb_options(command)
    mock_options.instance_spec = TaskInstanceKey(self.TEST_JOBKEY, [1])
    mock_options.no_batching = False
    mock_options.config = config

    fake_context = FakeAuroraCommandContext()
    fake_context.set_options(mock_options)
    fake_context.add_config(config)

    fake_context.add_expected_query_result(
      self.create_query_call_result(), job_key=self.TEST_JOBKEY)
    fake_context.add_expected_query_result(
      AuroraClientCommandTest.create_query_call_result(
        AuroraClientCommandTest.create_scheduled_task(1, ScheduleStatus.RUNNING)))

    mock_api = fake_context.get_api('test')
    mock_api.kill_job.return_value = self.create_simple_success_response()
    command.execute(fake_context)

    assert mock_api.kill_job.mock_calls == [
        call(self.TEST_JOBKEY, mock_options.instance_spec.instance, config=config, message=None)
    ]


class TestKillAllCommand(AuroraClientCommandTest):
  def test_killall_nobatch_passes_config(self):
    """Verify that the no batch code path correctly passes job config to the api."""
    command = KillAllJobCommand()
    config = self.get_valid_config()

    mock_options = mock_verb_options(command)
    mock_options.jobspec = self.TEST_JOBKEY
    mock_options.no_batching = True
    mock_options.config = config

    fake_context = FakeAuroraCommandContext()
    fake_context.set_options(mock_options)
    fake_context.add_config(config)

    mock_api = fake_context.get_api('test')
    fake_context.add_expected_query_result(
      self.create_query_call_result(), job_key=self.TEST_JOBKEY)
    fake_context.add_expected_query_result(AuroraClientCommandTest.create_empty_task_result())
    mock_api.kill_job.return_value = self.create_simple_success_response()

    command.execute(fake_context)

    assert mock_api.kill_job.mock_calls == [
        call(self.TEST_JOBKEY, None, config=config, message=None)
    ]


class TestClientKillCommand(AuroraClientCommandTest):
  @classmethod
  def get_expected_task_query(cls, instances=None):
    instance_ids = frozenset(instances) if instances is not None else None
    return TaskQuery(taskIds=None,
                     instanceIds=instance_ids,
                     jobKeys=[JobKey(role=cls.TEST_ROLE,
                                     environment=cls.TEST_ENV,
                                     name=cls.TEST_JOB)])

  @classmethod
  def get_monitor_mock(cls, result=True):
    mock_monitor = Mock(spec=JobMonitor)
    mock_monitor.wait_until.return_value = result
    return mock_monitor

  @classmethod
  def assert_kill_calls(cls, api, instance_range=None, instances=None, message=None):
    if instances:
      kill_calls = [call(cls.TEST_JOBKEY, instances, config=None, message=message)]
    else:
      kill_calls = [
          call(cls.TEST_JOBKEY, [i], config=None, message=message)
          for i in instance_range
      ]
    assert api.kill_job.mock_calls == kill_calls

  @classmethod
  def assert_wait_calls(cls, mock_monitor, terminal, instance_range=None, instances=None):
    if instances:
      wait_calls = [call(terminal, instances=instances, with_timeout=True)]
    else:
      wait_calls = [call(terminal, instances=[i], with_timeout=True) for i in instance_range]
    assert mock_monitor.wait_until.mock_calls == wait_calls

  @classmethod
  def assert_kill_call_no_instances(cls, api):
    assert api.kill_job.mock_calls == [call(cls.TEST_JOBKEY, None, config=None, message=None)]

  @classmethod
  def assert_query(cls, fake_api, times=1):
    calls = [call(TaskQuery(jobKeys=[cls.TEST_JOBKEY.to_thrift()], statuses=ACTIVE_STATES))
      for _ in range(times)]
    assert fake_api.query_no_configs.mock_calls == calls

  def test_killall_job(self):
    """Test killall client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    mock_monitor = self.get_monitor_mock()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor)) as (_, m):

      api = mock_context.get_api('west')
      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'killall', '--no-batching', self.TEST_JOBSPEC])

      self.assert_kill_call_no_instances(api)
      assert mock_monitor.wait_until.mock_calls == [
          call(m.terminal, instances=None, with_timeout=True)]

  def test_killall_job_batched(self):
    """Test killall command with batching."""
    mock_context = FakeAuroraCommandContext()
    mock_monitor = self.get_monitor_mock()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor)) as (_, m):

      api = mock_context.get_api('west')
      api.kill_job.return_value = self.create_simple_success_response()
      mock_context.add_expected_query_result(
          self.create_query_call_result(), job_key=self.TEST_JOBKEY)

      cmd = AuroraCommandLine()
      cmd.execute(['job', 'killall', self.TEST_JOBSPEC])

      self.assert_kill_calls(api, instance_range=range(20), message=None)
      self.assert_wait_calls(mock_monitor, m.terminal, instance_range=range(20))
      self.assert_query(api)

  def test_kill_job_with_instances_nobatching(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    mock_monitor = self.get_monitor_mock()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor)) as (_, m):
      api = mock_context.get_api('west')
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)
      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'kill', '--no-batching', self.get_instance_spec('0,2,4-6')])

      instances = [0, 2, 4, 5, 6]
      self.assert_kill_calls(api, instances=instances, message=None)
      self.assert_wait_calls(mock_monitor, m.terminal, instances=instances)
      self.assert_query(api)

  def test_kill_job_with_invalid_instances_strict(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context)):
      api = mock_context.get_api('west')
      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'kill', '--no-batching', '--strict',
           self.get_instance_spec('0,2,4-6,11-20')])

      assert api.kill_job.call_count == 0

  def test_kill_job_with_invalid_instances_nonstrict(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    mock_monitor = self.get_monitor_mock()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor)) as (_, m):
      api = mock_context.get_api('west')
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)
      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'kill', '--no-batching', self.get_instance_spec('0,2,4-6,11-13')])
      instances = [0, 2, 4, 5, 6, 11, 12, 13]
      self.assert_kill_calls(api, instances=instances, message=None)
      self.assert_wait_calls(mock_monitor, m.terminal, instances=instances)
      self.assert_query(api)

  def test_kill_job_with_instances_batched(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    mock_monitor = self.get_monitor_mock()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor)) as (_, m):
      api = mock_context.get_api('west')
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)

      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'kill', self.get_instance_spec('0-6')])

      self.assert_kill_calls(api, instance_range=range(7))
      self.assert_wait_calls(mock_monitor, m.terminal, instance_range=range(7))
      self.assert_query(api, times=2)

  def test_kill_job_with_instances_batched_maxerrors(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    mock_monitor = self.get_monitor_mock(result=False)
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor)):
      api = mock_context.get_api('west')
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)

      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'kill', '--max-total-failures=1', self.get_instance_spec('0-4')])

      # We should have aborted after the second batch.
      self.assert_kill_calls(api, instance_range=range(2), message=None)
      self.assert_query(api, times=2)

  def test_kill_job_with_empty_instances_batched(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context)):
      api = mock_context.get_api('west')
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)
      # set up an empty instance list in the getTasksWithoutConfigs response
      status_response = self.create_simple_success_response()
      status_response.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=[]))
      mock_context.add_expected_query_result(status_response)
      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'kill', self.get_instance_spec('0,2,4-13')])

      assert api.kill_job.call_count == 0
      self.assert_query(api, times=2)

  def test_killall_job_output(self):
    """Test kill output."""
    mock_context = FakeAuroraCommandContext()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=self.get_monitor_mock())):
      api = mock_context.get_api('west')
      api.kill_job.return_value = self.create_simple_success_response()

      cmd = AuroraCommandLine()
      cmd.execute(['job', 'killall', '--no-batching', self.TEST_JOBSPEC])

      assert mock_context.get_out() == ['Job killall succeeded']
      assert mock_context.get_err() == []

  def test_kill_job_with_instances_batched_output(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=self.get_monitor_mock())):
      api = mock_context.get_api('west')
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)
      mock_context.add_expected_query_result(self.create_query_call_result())
      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'kill', '--batch-size=5', self.get_instance_spec('0,2,4-6')])

    assert mock_context.get_out() == ['Successfully killed instances [0, 2, 4, 5, 6]',
        'Job kill succeeded']
    assert mock_context.get_err() == []
    self.assert_query(api, times=2)

  def test_kill_job_with_instances_batched_maxerrors_output(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    mock_monitor = self.get_monitor_mock(result=False)
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor)):
      api = mock_context.get_api('west')
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)
      mock_context.add_expected_query_result(self.create_query_call_result())
      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      cmd.execute(['job', 'kill', '--max-total-failures=1', '--batch-size=5',
          self.get_instance_spec('0,2,4-13')])

      assert mock_context.get_out() == []
      assert mock_context.get_err() == [
         'Instances [0, 2, 4, 5, 6] were not killed in time',
         'Instances [7, 8, 9, 10, 11] were not killed in time',
         'Exceeded maximum number of errors while killing instances']
      self.assert_query(api, times=2)

  def test_kill_job_with_message(self):
    """Test kill client-side API logic."""
    mock_context = FakeAuroraCommandContext()
    mock_monitor = self.get_monitor_mock()
    with contextlib.nested(
        patch('apache.aurora.client.cli.jobs.Job.create_context', return_value=mock_context),
        patch('apache.aurora.client.cli.jobs.JobMonitor', return_value=mock_monitor)) as (_, m):
      api = mock_context.get_api('west')
      mock_context.add_expected_query_result(
        self.create_query_call_result(), job_key=self.TEST_JOBKEY)
      api.kill_job.return_value = self.create_simple_success_response()
      cmd = AuroraCommandLine()
      message = 'Test message'
      cmd.execute(['job', 'kill', '--no-batch', '--message', message, self.get_instance_spec('0')])

      instances = [0]
      self.assert_kill_calls(api, instances=instances, message=message)
      self.assert_query(api)