text: #
  # 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
  from collections import defaultdict
  
  from mock import Mock, patch
  from twitter.common.contextutil import temporary_file
  
  from apache.aurora.client.api import AuroraClientAPI
  from apache.aurora.client.api.sla import DomainUpTimeSlaVector, JobUpTimeDetails, JobUpTimeLimit
  from apache.aurora.client.base import DEFAULT_GROUPING
  from apache.aurora.client.commands.admin import sla_list_safe_domain, sla_probe_hosts
  from apache.aurora.client.commands.util import AuroraClientCommandTest
  from apache.aurora.common.aurora_job_key import AuroraJobKey
  
  MIN_INSTANCE_COUNT = 1
  
  
  class TestAdminSlaListSafeDomainCommand(AuroraClientCommandTest):
  
    @classmethod
    def setup_mock_options(cls, exclude=None, include=None, override=None,
                           exclude_list=None, include_list=None, list_jobs=False, grouping=None):
      mock_options = Mock(spec=['exclude_filename', 'exclude_hosts', 'include_filename',
          'include_hosts', 'override_filename', 'list_jobs', 'verbosity', 'disable_all_hooks',
          'min_instance_count', 'grouping'])
  
      mock_options.exclude_filename = exclude
      mock_options.exclude_hosts = exclude_list
      mock_options.include_filename = include
      mock_options.include_hosts = include_list
      mock_options.override_filename = override
      mock_options.list_jobs = list_jobs
      mock_options.verbosity = False
      mock_options.disable_all_hooks = False
      mock_options.min_instance_count = MIN_INSTANCE_COUNT
      mock_options.grouping = grouping or DEFAULT_GROUPING
      return mock_options
  
    @classmethod
    def create_hosts(cls, num_hosts, percentage, duration):
      hosts = defaultdict(list)
      for i in range(num_hosts):
        host_name = 'h%s' % i
        job = AuroraJobKey.from_path('west/role/env/job%s' % i)
        hosts[host_name].append(JobUpTimeLimit(job, percentage, duration))
      return [hosts]
  
    @classmethod
    def create_mock_vector(cls, result):
      mock_vector = Mock(spec=DomainUpTimeSlaVector)
      mock_vector.get_safe_hosts.return_value = result
      return mock_vector
  
    def test_safe_domain_no_options(self):
      """Tests successful execution of the sla_list_safe_domain command without extra options."""
      mock_options = self.setup_mock_options()
      mock_vector = self.create_mock_vector(self.create_hosts(3, 80, 100))
      with contextlib.nested(
          patch('apache.aurora.client.commands.admin.AuroraClientAPI',
              new=Mock(spec=AuroraClientAPI)),
          patch('apache.aurora.client.commands.admin.print_results'),
          patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
          patch('twitter.common.app.get_options', return_value=mock_options)
      ) as (
          mock_api,
          mock_print_results,
          test_clusters,
          mock_options):
  
        mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
        sla_list_safe_domain(['west', '50', '100s'])
  
        mock_vector.get_safe_hosts.assert_called_once_with(50.0, 100.0, {}, DEFAULT_GROUPING)
        mock_print_results.assert_called_once_with(['h0', 'h1', 'h2'])
  
    def test_safe_domain_exclude_hosts(self):
      """Test successful execution of the sla_list_safe_domain command with exclude hosts option."""
      mock_vector = self.create_mock_vector(self.create_hosts(3, 80, 100))
      with temporary_file() as fp:
        fp.write('h1')
        fp.flush()
        mock_options = self.setup_mock_options(exclude=fp.name)
        with contextlib.nested(
            patch('apache.aurora.client.commands.admin.AuroraClientAPI',
                new=Mock(spec=AuroraClientAPI)),
            patch('apache.aurora.client.commands.admin.print_results'),
            patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
            patch('twitter.common.app.get_options', return_value=mock_options)
        ) as (
            mock_api,
            mock_print_results,
            test_clusters,
            mock_options):
  
          mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
  
          sla_list_safe_domain(['west', '50', '100s'])
  
          mock_vector.get_safe_hosts.assert_called_once_with(50.0, 100.0, {}, DEFAULT_GROUPING)
          mock_print_results.assert_called_once_with(['h0', 'h2'])
  
    def test_safe_domain_exclude_hosts_from_list(self):
      """Test successful execution of the sla_list_safe_domain command with exclude list option."""
      mock_vector = self.create_mock_vector(self.create_hosts(3, 80, 100))
      mock_options = self.setup_mock_options(exclude_list=','.join(['h0', 'h1']))
      with contextlib.nested(
          patch('apache.aurora.client.commands.admin.AuroraClientAPI',
              new=Mock(spec=AuroraClientAPI)),
          patch('apache.aurora.client.commands.admin.print_results'),
          patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
          patch('twitter.common.app.get_options', return_value=mock_options)
      ) as (
          mock_api,
          mock_print_results,
          test_clusters,
          mock_options):
  
        mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
  
        sla_list_safe_domain(['west', '50', '100s'])
  
        mock_vector.get_safe_hosts.assert_called_once_with(50.0, 100.0, {}, DEFAULT_GROUPING)
        mock_print_results.assert_called_once_with(['h2'])
  
    def test_safe_domain_include_hosts(self):
      """Test successful execution of the sla_list_safe_domain command with include hosts option."""
      mock_vector = self.create_mock_vector(self.create_hosts(1, 80, 100))
      hostname = 'h0'
      with temporary_file() as fp:
        fp.write(hostname)
        fp.flush()
        mock_options = self.setup_mock_options(include=fp.name)
        with contextlib.nested(
            patch('apache.aurora.client.commands.admin.AuroraClientAPI',
                new=Mock(spec=AuroraClientAPI)),
            patch('apache.aurora.client.commands.admin.print_results'),
            patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
            patch('twitter.common.app.get_options', return_value=mock_options)
        ) as (
            mock_api,
            mock_print_results,
            test_clusters,
            mock_options):
  
          mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
  
          sla_list_safe_domain(['west', '50', '100s'])
  
          mock_api.return_value.sla_get_safe_domain_vector.assert_called_once_with(
              MIN_INSTANCE_COUNT, [hostname])
          mock_vector.get_safe_hosts.assert_called_once_with(50.0, 100.0, {}, DEFAULT_GROUPING)
          mock_print_results.assert_called_once_with([hostname])
  
    def test_safe_domain_include_hosts_from_list(self):
      """Test successful execution of the sla_list_safe_domain command with include list option."""
      mock_vector = self.create_mock_vector(self.create_hosts(2, 80, 100))
      hosts = ['h0', 'h1']
      mock_options = self.setup_mock_options(include_list=','.join(hosts))
      with contextlib.nested(
          patch('apache.aurora.client.commands.admin.AuroraClientAPI',
              new=Mock(spec=AuroraClientAPI)),
          patch('apache.aurora.client.commands.admin.print_results'),
          patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
          patch('twitter.common.app.get_options', return_value=mock_options)
      ) as (
          mock_api,
          mock_print_results,
          test_clusters,
          mock_options):
  
        mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
  
        sla_list_safe_domain(['west', '50', '100s'])
  
        mock_api.return_value.sla_get_safe_domain_vector.assert_called_once_with(
            MIN_INSTANCE_COUNT, hosts)
        mock_vector.get_safe_hosts.assert_called_once_with(50.0, 100.0, {}, DEFAULT_GROUPING)
        mock_print_results.assert_called_once_with(hosts)
  
    def test_safe_domain_override_jobs(self):
      """Test successful execution of the sla_list_safe_domain command with override_jobs option."""
      mock_vector = self.create_mock_vector(self.create_hosts(3, 80, 100))
      with temporary_file() as fp:
        fp.write('west/role/env/job1 30 200s')
        fp.flush()
        mock_options = self.setup_mock_options(override=fp.name)
        with contextlib.nested(
            patch('apache.aurora.client.commands.admin.AuroraClientAPI',
                new=Mock(spec=AuroraClientAPI)),
            patch('apache.aurora.client.commands.admin.print_results'),
            patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
            patch('twitter.common.app.get_options', return_value=mock_options)
        ) as (
            mock_api,
            mock_print_results,
            test_clusters,
            mock_options):
  
          mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
  
          sla_list_safe_domain(['west', '50', '100s'])
  
          job_key = AuroraJobKey.from_path('west/role/env/job1')
          override = {job_key: JobUpTimeLimit(job_key, 30, 200)}
          mock_vector.get_safe_hosts.assert_called_once_with(50.0, 100.0, override, DEFAULT_GROUPING)
          mock_print_results.assert_called_once_with(['h0', 'h1', 'h2'])
  
    def test_safe_domain_list_jobs(self):
      """Tests successful execution of the sla_list_safe_domain command with list_jobs option."""
      mock_options = self.setup_mock_options(list_jobs=True)
      mock_vector = self.create_mock_vector(self.create_hosts(3, 50, 100))
      with contextlib.nested(
          patch('apache.aurora.client.commands.admin.AuroraClientAPI',
              new=Mock(spec=AuroraClientAPI)),
          patch('apache.aurora.client.commands.admin.print_results'),
          patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
          patch('twitter.common.app.get_options', return_value=mock_options)
      ) as (
          mock_api,
          mock_print_results,
          test_clusters,
          mock_options):
  
        mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
        sla_list_safe_domain(['west', '50', '100s'])
  
        mock_vector.get_safe_hosts.assert_called_once_with(50.0, 100.0, {}, DEFAULT_GROUPING)
        mock_print_results.assert_called_once_with([
            'h0\twest/role/env/job0\t50.00\t100',
            'h1\twest/role/env/job1\t50.00\t100',
            'h2\twest/role/env/job2\t50.00\t100'])
  
    def test_safe_domain_invalid_percentage(self):
      """Tests execution of the sla_list_safe_domain command with invalid percentage"""
      mock_options = self.setup_mock_options()
      with patch('twitter.common.app.get_options', return_value=mock_options):
        try:
          sla_list_safe_domain(['west', '0', '100s'])
        except SystemExit:
          pass
        else:
          assert 'Expected error is not raised.'
  
    def test_safe_domain_malformed_job_override(self):
      """Tests execution of the sla_list_safe_domain command with invalid job_override file"""
      with temporary_file() as fp:
        fp.write('30 200s')
        fp.flush()
        mock_options = self.setup_mock_options(override=fp.name)
        with patch('twitter.common.app.get_options', return_value=mock_options):
  
          try:
            sla_list_safe_domain(['west', '50', '100s'])
          except SystemExit:
            pass
          else:
            assert 'Expected error is not raised.'
  
    def test_safe_domain_hosts_error(self):
      """Tests execution of the sla_list_safe_domain command with both include file and list"""
      mock_options = self.setup_mock_options(include='file', include_list='list')
      with patch('twitter.common.app.get_options', return_value=mock_options):
  
        try:
          sla_list_safe_domain(['west', '50', '100s'])
        except SystemExit:
          pass
        else:
          assert 'Expected error is not raised.'
  
    def test_safe_domain_grouping_error(self):
      """Tests execution of the sla_list_safe_domain command invalid grouping"""
      mock_options = self.setup_mock_options(grouping='foo')
      with patch('twitter.common.app.get_options', return_value=mock_options):
  
        try:
          sla_list_safe_domain(['west', '50', '100s'])
        except SystemExit:
          pass
        else:
          assert 'Expected error is not raised.'
  
  
  class TestAdminSlaProbeHostsCommand(AuroraClientCommandTest):
  
    @classmethod
    def setup_mock_options(cls, hosts=None, filename=None, grouping=None):
      mock_options = Mock(spec=['hosts', 'filename', 'verbosity', 'min_instance_count', 'grouping'])
      mock_options.hosts = hosts
      mock_options.filename = filename
      mock_options.verbosity = False
      mock_options.grouping = grouping or DEFAULT_GROUPING
      return mock_options
  
    @classmethod
    def create_mock_vector(cls, result):
      mock_vector = Mock(spec=DomainUpTimeSlaVector)
      mock_vector.probe_hosts.return_value = result
      return mock_vector
  
    @classmethod
    def create_probe_hosts(cls, num_hosts, predicted, safe, safe_in):
      hosts = defaultdict(list)
      for i in range(num_hosts):
        host_name = 'h%s' % i
        job = AuroraJobKey.from_path('west/role/env/job%s' % i)
        hosts[host_name].append(JobUpTimeDetails(job, predicted, safe, safe_in))
      return [hosts]
  
    def test_probe_hosts_with_list(self):
      """Tests successful execution of the sla_probe_hosts command with host list."""
      hosts = ['h0', 'h1']
      mock_options = self.setup_mock_options(hosts=','.join(hosts))
      mock_vector = self.create_mock_probe_hosts_vector([self.create_probe_hosts(2, 80, True, 0)])
      with contextlib.nested(
          patch('apache.aurora.client.commands.admin.AuroraClientAPI',
              new=Mock(spec=AuroraClientAPI)),
          patch('apache.aurora.client.commands.admin.print_results'),
          patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
          patch('twitter.common.app.get_options', return_value=mock_options)
      ) as (
          mock_api,
          mock_print_results,
          test_clusters,
          options):
  
        mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
        sla_probe_hosts(['west', '90', '200s'])
  
        mock_api.return_value.sla_get_safe_domain_vector.assert_called_once_with(
            mock_options.min_instance_count, hosts)
        mock_vector.probe_hosts.assert_called_once_with(90.0, 200.0, mock_options.grouping)
        mock_print_results.assert_called_once_with([
            'h0\twest/role/env/job0\t80.00\tTrue\t0',
            'h1\twest/role/env/job1\t80.00\tTrue\t0'
        ])
  
    def test_probe_hosts_with_file(self):
      """Tests successful execution of the sla_probe_hosts command with host filename."""
      mock_vector = self.create_mock_probe_hosts_vector([self.create_probe_hosts(1, 80, False, None)])
      with temporary_file() as fp:
        fp.write('h0')
        fp.flush()
        mock_options = self.setup_mock_options(filename=fp.name)
        with contextlib.nested(
            patch('apache.aurora.client.commands.admin.AuroraClientAPI',
                new=Mock(spec=AuroraClientAPI)),
            patch('apache.aurora.client.commands.admin.print_results'),
            patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
            patch('twitter.common.app.get_options', return_value=mock_options)
        ) as (
            mock_api,
            mock_print_results,
            test_clusters,
            options):
  
          mock_api.return_value.sla_get_safe_domain_vector.return_value = mock_vector
          sla_probe_hosts(['west', '90', '200s'])
  
          mock_api.return_value.sla_get_safe_domain_vector.assert_called_once_with(
            mock_options.min_instance_count, ['h0'])
          mock_vector.probe_hosts.assert_called_once_with(90.0, 200.0, mock_options.grouping)
          mock_print_results.assert_called_once_with([
              'h0\twest/role/env/job0\t80.00\tFalse\tn/a'
          ])
  
    def test_probe_hosts_error(self):
      """Tests execution of the sla_probe_hosts command with both host and filename provided."""
      with temporary_file() as fp:
        fp.write('h0')
        fp.flush()
        mock_options = self.setup_mock_options(hosts='h0', filename=fp.name)
        with patch('twitter.common.app.get_options', return_value=mock_options):
  
          try:
            sla_probe_hosts(['west', '50', '100s'])
          except SystemExit:
            pass
          else:
            assert 'Expected error is not raised.'
  
    def test_probe_grouping_error(self):
      """Tests execution of the sla_probe_hosts command with invalid grouping."""
      mock_options = self.setup_mock_options(hosts='h0', grouping='foo')
      with patch('twitter.common.app.get_options', return_value=mock_options):
  
        try:
          sla_probe_hosts(['west', '50', '100s'])
        except SystemExit:
          pass
        else:
          assert 'Expected error is not raised.'
  