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

import pytest
from mock import Mock, patch

from apache.aurora.client.cli import EXIT_INVALID_PARAMETER, EXIT_OK, Context
from apache.aurora.client.cli.client import AuroraCommandLine
from apache.aurora.client.cli.task import ScpCommand

from .util import AuroraClientCommandTest, FakeAuroraCommandContext, mock_verb_options

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


class TestRunCommand(AuroraClientCommandTest):

  @classmethod
  def create_status_response(cls):
    resp = cls.create_simple_success_response()
    resp.result = Result(
        scheduleStatusResult=ScheduleStatusResult(tasks=cls.create_scheduled_tasks()))
    return resp

  @classmethod
  def create_failed_status_response(cls):
    return cls.create_blank_response(ResponseCode.INVALID_REQUEST, 'No tasks found for query')

  @classmethod
  def create_mock_process(cls):
    process = Mock()
    process.communicate.return_value = ["hello", "world"]
    return process

  def test_successful_run(self):
    """Test the run command."""
    self.generic_test_successful_run(['task', 'run', 'west/bozo/test/hello', 'ls'], None)

  def test_successful_run_with_instances(self):
    """Test the run command."""
    self.generic_test_successful_run(['task', 'run', 'west/bozo/test/hello/1-3', 'ls'], [1, 2, 3])

  def test_successful_run_with_ssh_options(self):
    self.generic_test_successful_run(
        ['task', 'run', '--ssh-options=-v -k', 'west/bozo/test/hello', 'ls'],
        None,
        ssh_options=['-v', '-k'])

  def generic_test_successful_run(self, cmd_args, instances, ssh_options=None):
    """Common structure of all successful run tests.
    Params:
      cmd_args: the arguments to pass to the aurora command line to run this test.
      instances: the list of instances that should be passed to a status query.
         (The status query is the only visible difference between a sharded
         run, and an all-instances run in the test.)
    """
    # Calls api.check_status, which calls scheduler_proxy.getJobs
    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
    mock_scheduler_proxy.getTasksStatus.return_value = self.create_status_response()
    sandbox_args = {'slave_root': '/slaveroot', 'slave_run_directory': 'slaverun'}
    with contextlib.nested(
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('apache.aurora.client.api.command_runner.'
              'InstanceDistributedCommandRunner.sandbox_args',
            return_value=sandbox_args),
        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
            mock_scheduler_proxy_class,
            mock_runner_args_patch,
            mock_subprocess):
      cmd = AuroraCommandLine()
      assert cmd.execute(cmd_args) == EXIT_OK
      # The status command sends a getTasksStatus query to the scheduler,
      # and then prints the result. The use of shards, above, should change
      # this query - that's the focus of the instances test.
      mock_scheduler_proxy.getTasksStatus.assert_called_with(TaskQuery(
          jobKeys=[JobKey(role='bozo', environment='test', name='hello')],
          statuses=set([ScheduleStatus.RUNNING, ScheduleStatus.KILLING, ScheduleStatus.RESTARTING,
              ScheduleStatus.PREEMPTING, ScheduleStatus.PARTITIONED, ScheduleStatus.DRAINING]),
          instanceIds=instances),
          retry=True)

      # The mock status call returns 3 three ScheduledTasks, so three commands should have been run
      assert mock_subprocess.call_count == 3
      expected = ['ssh', '-n', '-q']
      expected += ssh_options if ssh_options else []
      expected += ['bozo@slavehost',
          'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/'
          'slaverun/sandbox;ls']
      mock_subprocess.assert_called_with(expected, stderr=-2, stdout=-1)


class TestSshCommand(AuroraClientCommandTest):
  MOCKED_PID = 12312

  @classmethod
  def create_mock_process(cls):
    process = Mock()
    process.pid = cls.MOCKED_PID
    process.wait.return_value = 0
    return process

  @classmethod
  def create_status_response(cls):
    resp = cls.create_simple_success_response()
    resp.result = Result(
        scheduleStatusResult=ScheduleStatusResult(tasks=cls.create_scheduled_tasks()))
    return resp

  @classmethod
  def create_nojob_status_response(cls):
    resp = cls.create_simple_success_response()
    resp.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=[]))
    return resp

  @classmethod
  def create_failed_status_response(cls):
    return cls.create_blank_response(ResponseCode.INVALID_REQUEST, 'No tasks found for query')

  def test_successful_ssh_with_pidfile(self):
    """Test the ssh command with a saved PID file."""
    with tempfile.NamedTemporaryFile() as pid_file:
      (mock_api, mock_scheduler_proxy) = self.create_mock_api()
      mock_scheduler_proxy.getTasksStatus.return_value = self.create_status_response()
      sandbox_args = {'slave_root': '/slaveroot', 'slave_run_directory': 'slaverun'}
      with contextlib.nested(
          patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
          patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
              return_value=sandbox_args),
          patch('subprocess.Popen', return_value=self.create_mock_process())) as (
              mock_scheduler_proxy_class,
              mock_runner_args_patch,
              mock_subprocess):
        cmd = AuroraCommandLine()
        cmd.execute([
          'task',
          'ssh',
          '--pid-file={}'.format(pid_file.name),
          '--ssh-options=-v',
          'west/bozo/test/hello/1',
          '--command=ls'
        ])

        # The status command sends a getTasksStatus query to the scheduler,
        # and then prints the result.
        mock_scheduler_proxy.getTasksStatus.assert_called_with(TaskQuery(
            jobKeys=[JobKey(role='bozo', environment='test', name='hello')],
            instanceIds=set([1]),
            statuses=set([ScheduleStatus.RUNNING, ScheduleStatus.KILLING, ScheduleStatus.RESTARTING,
                ScheduleStatus.PREEMPTING, ScheduleStatus.PARTITIONED, ScheduleStatus.DRAINING])),
            retry=True)
        mock_subprocess.assert_called_with(['ssh', '-t', '-v', 'bozo@slavehost',
            'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/'
            'slaverun/sandbox;ls'])

        pid = pid_file.read()
        assert(pid == str(self.MOCKED_PID))

  def test_successful_ssh(self):
    """Test the ssh command."""
    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
    mock_scheduler_proxy.getTasksStatus.return_value = self.create_status_response()
    sandbox_args = {'slave_root': '/slaveroot', 'slave_run_directory': 'slaverun'}
    with contextlib.nested(
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
            return_value=sandbox_args),
        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
            mock_scheduler_proxy_class,
            mock_runner_args_patch,
            mock_subprocess):
      cmd = AuroraCommandLine()
      cmd.execute(['task', 'ssh', '--ssh-options=-v', 'west/bozo/test/hello/1', '--command=ls'])

      # The status command sends a getTasksStatus query to the scheduler,
      # and then prints the result.
      mock_scheduler_proxy.getTasksStatus.assert_called_with(TaskQuery(
          jobKeys=[JobKey(role='bozo', environment='test', name='hello')],
          instanceIds=set([1]),
          statuses=set([ScheduleStatus.RUNNING, ScheduleStatus.KILLING, ScheduleStatus.RESTARTING,
              ScheduleStatus.PREEMPTING, ScheduleStatus.PARTITIONED, ScheduleStatus.DRAINING])),
          retry=True)
      mock_subprocess.assert_called_with(['ssh', '-t', '-v', 'bozo@slavehost',
          'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/'
          'slaverun/sandbox;ls'])

  def test_successful_ssh_no_instance(self):
    """Test the ssh command when the instance id is not specified."""
    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
    mock_scheduler_proxy.getTasksStatus.return_value = self.create_status_response()
    sandbox_args = {'slave_root': '/slaveroot', 'slave_run_directory': 'slaverun'}
    with contextlib.nested(
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
            return_value=sandbox_args),
        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
            mock_scheduler_proxy_class,
            mock_runner_args_patch,
            mock_subprocess):
      cmd = AuroraCommandLine()
      cmd.execute(['task', 'ssh', '--ssh-options=-v', 'west/bozo/test/hello'])

      # The status command sends a getTasksStatus query to the scheduler,
      # and then prints the result.
      mock_scheduler_proxy.getTasksStatus.assert_called_with(TaskQuery(
          jobKeys=[JobKey(role='bozo', environment='test', name='hello')],
          instanceIds=None,
          statuses=set([ScheduleStatus.RUNNING, ScheduleStatus.KILLING, ScheduleStatus.RESTARTING,
              ScheduleStatus.PREEMPTING, ScheduleStatus.PARTITIONED, ScheduleStatus.DRAINING])),
          retry=True)
      mock_subprocess.assert_called_with(['ssh', '-t', '-v', 'bozo@slavehost',
          'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/'
          'slaverun/sandbox;bash'])

  def test_ssh_job_not_found(self):
    """Test the ssh command when the jobkey parameter specifies a job that isn't running."""
    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
    mock_scheduler_proxy.getTasksStatus.return_value = self.create_nojob_status_response()
    with contextlib.nested(
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
            mock_scheduler_proxy_class,
            mock_subprocess):
      cmd = AuroraCommandLine()
      result = cmd.execute(['task', 'ssh', 'west/bozo/test/hello/1', '--command=ls'])
      assert result == EXIT_INVALID_PARAMETER
      assert mock_subprocess.call_count == 0

  def test_ssh_no_instance_command(self):
    """Test the ssh command when the jobkey parameter doesn't specify an instance."""
    (mock_api, mock_scheduler_proxy) = self.create_mock_api()
    mock_scheduler_proxy.getTasksStatus.return_value = self.create_nojob_status_response()
    with contextlib.nested(
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
            mock_scheduler_proxy_class,
            mock_subprocess):
      cmd = AuroraCommandLine()
      result = cmd.execute(['task', 'ssh', 'west/bozo/test/hello', '--command=ls'])
      assert result == EXIT_INVALID_PARAMETER
      assert mock_subprocess.call_count == 0


class TestScpCommand(AuroraClientCommandTest):

  @classmethod
  def create_status_response(cls):
    resp = cls.create_simple_success_response()
    resp.result = Result(
        scheduleStatusResult=ScheduleStatusResult(tasks=cls.create_scheduled_tasks()))
    return resp

  @classmethod
  def create_nojob_status_response(cls):
    resp = cls.create_simple_success_response()
    resp.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=[]))
    return resp

  def setUp(self):
    self._command = ScpCommand()
    self._mock_options = mock_verb_options(self._command)
    self._fake_context = FakeAuroraCommandContext()
    self._fake_context.set_options(self._mock_options)
    self._mock_api = self._fake_context.get_api('UNUSED')
    self._sandbox_args = {'slave_root': '/slaveroot', 'slave_run_directory': 'slaverun'}

  def test_successful_scp_simple(self):
    """Test the scp command."""
    self._mock_api.query.return_value = self.create_status_response()
    self._mock_options.scp_options = ['-v', '-t']
    self._mock_options.source = ['test.txt']
    self._mock_options.dest = 'west/bozo/test/hello/1:./test/dir'

    with contextlib.nested(
        patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
            return_value=self._sandbox_args),
        patch('subprocess.call', return_value=EXIT_OK)) as [
            mock_runner_args_patch,
            mock_subprocess]:
      assert self._command.execute(self._fake_context) == EXIT_OK
      mock_subprocess.assert_called_with(['scp', '-v', '-t', 'test.txt',
          'bozo@slavehost:'
           '/slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/slaverun/sandbox/'
           'test/dir'])

  def test_successful_scp_absolute_path(self):
    """Test the scp command uses absolute paths correctly."""
    self._mock_api.query.return_value = self.create_status_response()
    self._mock_options.scp_options = ['-v']
    self._mock_options.source = ['test.txt']
    self._mock_options.dest = 'west/bozo/test/hello/1:/tmp'

    with contextlib.nested(
        patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
            return_value=self._sandbox_args),
        patch('subprocess.call', return_value=EXIT_OK)) as [
            mock_runner_args_patch,
            mock_subprocess]:
      assert self._command.execute(self._fake_context) == EXIT_OK
      mock_subprocess.assert_called_with(['scp', '-v', 'test.txt', 'bozo@slavehost:/tmp'])

  def test_successful_scp_two_instances(self):
    """Test that the scp command correctly evaluates commands with two jobkeys."""
    self._mock_api.query.return_value = self.create_status_response()
    self._mock_options.scp_options = ['-v']
    self._mock_options.source = ['west/bozo/test/hello/1:test.txt']
    self._mock_options.dest = 'west/bozo/test/hello/1:/tmp'

    with contextlib.nested(
        patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
            return_value=self._sandbox_args),
        patch('subprocess.call', return_value=EXIT_OK)) as [
            mock_runner_args_patch,
            mock_subprocess]:
      assert self._command.execute(self._fake_context) == EXIT_OK
      mock_subprocess.assert_called_with(['scp', '-v',
          'bozo@slavehost:'
           '/slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/slaverun/sandbox/'
           'test.txt',
          'bozo@slavehost:/tmp'])

  def test_successful_scp_multiple_files(self):
    """Test that the scp command correctly evaluates commands with multiple files."""
    self._mock_api.query.return_value = self.create_status_response()
    self._mock_options.source = ['test.txt', 'another.txt']
    self._mock_options.dest = 'west/bozo/test/hello/1:test/dir'

    with contextlib.nested(
        patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
            return_value=self._sandbox_args),
        patch('subprocess.call', return_value=EXIT_OK)) as [
            mock_runner_args_patch,
            mock_subprocess]:
      assert self._command.execute(self._fake_context) == EXIT_OK
      mock_subprocess.assert_called_with(['scp', 'test.txt', 'another.txt',
          'bozo@slavehost:'
           '/slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/slaverun/sandbox/'
           'test/dir'])

  def test_scp_invalid_tilde_expansion(self):
    """Test the scp command fails when using tilde expansion."""
    self._mock_api.query.return_value = self.create_status_response()
    self._mock_options.scp_options = ['-v']
    self._mock_options.source = ['test.txt']
    self._mock_options.dest = 'west/bozo/test/hello/1:~/test.txt'

    with contextlib.nested(patch('subprocess.call', return_value=EXIT_OK)) as [mock_subprocess]:
      with pytest.raises(Context.CommandError) as exc:
        assert self._command.execute(self._fake_context) == EXIT_INVALID_PARAMETER
      assert(ScpCommand.TILDE_USAGE_ERROR_MSG % '~/test.txt' in exc.value.message)
      assert mock_subprocess.call_count == 0

    # Test another tilde expansion form
    self._mock_options.source = ['test.txt']
    self._mock_options.dest = 'west/bozo/test/hello/1:~'

    with contextlib.nested(patch('subprocess.call', return_value=EXIT_OK)) as [mock_subprocess]:
      with pytest.raises(Context.CommandError) as exc:
        assert self._command.execute(self._fake_context) == EXIT_INVALID_PARAMETER
      assert(ScpCommand.TILDE_USAGE_ERROR_MSG % '~' in exc.value.message)
      assert mock_subprocess.call_count == 0

  def test_scp_bad_jobkey_no_instance(self):
    """Test the scp command fails when instance id is not specified."""
    self._mock_api.query.return_value = self.create_status_response()
    self._mock_options.source = ['test.txt']
    self._mock_options.dest = 'west/bozo/test/hello:test/dir'

    with contextlib.nested(patch('subprocess.call', return_value=EXIT_OK)) as [mock_subprocess]:
      with pytest.raises(Context.CommandError) as exc:
        assert self._command.execute(self._fake_context) == EXIT_INVALID_PARAMETER
      assert('not in the form CLUSTER/ROLE/ENV/NAME/INSTANCE' in exc.value.message)
      assert mock_subprocess.call_count == 0

  def test_scp_bad_jobkey_invalid_format(self):
    """Test the scp command fails when given general scp format."""
    self._mock_api.query.return_value = self.create_status_response()
    self._mock_options.source = ['root@192.168.0.1:test.txt']
    self._mock_options.dest = 'west/bozo/test/hello/1:test/dir'

    with contextlib.nested(patch('subprocess.call', return_value=EXIT_OK)) as [mock_subprocess]:
      with pytest.raises(Context.CommandError) as exc:
        assert self._command.execute(self._fake_context) == EXIT_INVALID_PARAMETER
      assert('not in the form CLUSTER/ROLE/ENV/NAME/INSTANCE' in exc.value.message)
      assert mock_subprocess.call_count == 0

  def test_scp_job_not_found(self):
    """Test the scp command when the jobkey parameter specifies a job that isn't running."""
    self._mock_api.query.return_value = self.create_nojob_status_response()
    self._mock_options.source = ['test.txt']
    self._mock_options.dest = 'west/bozo/test/hello/0:test/dir'

    with contextlib.nested(patch('subprocess.call', return_value=EXIT_OK)) as [mock_subprocess]:
      with pytest.raises(Context.CommandError) as exc:
        assert self._command.execute(self._fake_context) == EXIT_INVALID_PARAMETER
      assert(ScpCommand.JOB_NOT_FOUND_ERROR_MSG % ('west/bozo/test/hello', '0')
          in exc.value.message)
      assert mock_subprocess.call_count == 0