def test_perform_maintenance_hosts_failed_custom_sla(self):
    with temporary_file() as fp:
      mock_options = self.make_mock_options()
      mock_options.post_drain_script = None
      mock_options.grouping = 'by_host'
      mock_options.percentage = 50
      mock_options.duration = '10m'
      mock_options.reason = 'Test overrides'
      mock_options.unsafe_hosts_filename = fp.name

      mock_api, mock_scheduler_proxy = self.create_mock_api()
      mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result()
      mock_scheduler_proxy.drainHosts.return_value = self.create_start_maintenance_result()
      mock_vector = self.create_mock_probe_hosts_vector([
          self.create_probe_hosts(self.HOSTNAMES[0], 95, False, None),
          self.create_probe_hosts(self.HOSTNAMES[1], 95, False, None),
          self.create_probe_hosts(self.HOSTNAMES[2], 95, False, None)
      ])

      with contextlib.nested(
          patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
          patch('apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
                return_value=mock_vector),
          patch('apache.aurora.client.commands.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
          patch('apache.aurora.admin.admin_util.log_admin_message'),
          patch('twitter.common.app.get_options', return_value=mock_options)) as (_, _, _, log, _):
        host_drain([self.TEST_CLUSTER])

        assert 'Test overrides' in log.call_args[0][1]
        mock_scheduler_proxy.startMaintenance.assert_called_with(Hosts(set(self.HOSTNAMES)))
  def test_perform_maintenance_hosts_no_prod_tasks(self):
    mock_options = self.make_mock_options()
    mock_options.post_drain_script = None
    mock_options.grouping = 'by_host'

    def host_status_results(hostnames):
      if isinstance(hostnames, Hosts):
        return self.create_drained_status_result(hostnames)
      return self.create_maintenance_status_result()

    mock_api, mock_scheduler_proxy = self.create_mock_api()
    mock_scheduler_proxy.maintenanceStatus.side_effect = host_status_results
    mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result()
    mock_scheduler_proxy.drainHosts.return_value = self.create_start_maintenance_result()

    def create_empty_sla_results():
      mock_vector = Mock()
      mock_vector.probe_hosts.return_value = []
      return mock_vector

    with contextlib.nested(
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
              return_value=create_empty_sla_results()),
        patch('apache.aurora.client.commands.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
        patch('twitter.common.app.get_options', return_value=mock_options)):

      host_drain([self.TEST_CLUSTER])

      mock_scheduler_proxy.startMaintenance.assert_called_with(Hosts(set(self.HOSTNAMES)))
      assert mock_scheduler_proxy.maintenanceStatus.call_count == 3
      assert mock_scheduler_proxy.drainHosts.call_count == 3
  def test_perform_maintenance_hosts_unknown_hosts_skipped(self):
    mock_options = self.make_mock_options()
    mock_options.post_drain_script = None
    mock_options.grouping = 'by_host'

    def host_status_results(hostnames):
      if isinstance(hostnames, Hosts):
        return self.create_drained_status_result(hostnames)
      return self.create_maintenance_status_result()

    mock_api, mock_scheduler_proxy = self.create_mock_api()
    mock_scheduler_proxy.maintenanceStatus.side_effect = host_status_results
    mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result(
        skip_hosts=['us-grf-20'])
    mock_scheduler_proxy.drainHosts.return_value = self.create_start_maintenance_result()
    mock_vector = self.create_mock_probe_hosts_vector([
      self.create_probe_hosts(self.HOSTNAMES[0], 95, True, None),
      self.create_probe_hosts(self.HOSTNAMES[1], 95, True, None),
      self.create_probe_hosts(self.HOSTNAMES[2], 95, True, None)
    ])

    with contextlib.nested(
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
              return_value=mock_vector),
        patch('apache.aurora.client.commands.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
        patch('twitter.common.app.get_options', return_value=mock_options)):
      host_drain([self.TEST_CLUSTER])

      mock_scheduler_proxy.startMaintenance.assert_called_with(Hosts(set(self.HOSTNAMES)))
      assert mock_scheduler_proxy.maintenanceStatus.call_count == 2
      assert mock_scheduler_proxy.drainHosts.call_count == 2
  def test_perform_maintenance_hosts_failed_default_sla(self):
    with temporary_file() as fp:
      mock_options = self.make_mock_options()
      mock_options.post_drain_script = None
      mock_options.grouping = 'by_host'
      mock_options.unsafe_hosts_filename = fp.name

      mock_api, mock_scheduler_proxy = self.create_mock_api()
      mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result()
      mock_scheduler_proxy.drainHosts.return_value = self.create_start_maintenance_result()
      mock_vector = self.create_mock_probe_hosts_vector([
          self.create_probe_hosts(self.HOSTNAMES[0], 95, False, None),
          self.create_probe_hosts(self.HOSTNAMES[1], 95, False, None),
          self.create_probe_hosts(self.HOSTNAMES[2], 95, False, None)
      ])

      with contextlib.nested(
          patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
          patch('apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
                return_value=mock_vector),
          patch('apache.aurora.client.commands.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
          patch('twitter.common.app.get_options', return_value=mock_options)):
        host_drain([self.TEST_CLUSTER])

        mock_scheduler_proxy.startMaintenance.assert_called_with(Hosts(set(self.HOSTNAMES)))
    def test_perform_maintenance_hosts_multiple_sla_groups_failure(self):
        mock_options = self.make_mock_options()
        mock_options.post_drain_script = None
        mock_options.grouping = 'by_host'
        mock_options.unsafe_hosts_filename = None

        mock_api, mock_scheduler_proxy = self.create_mock_api()
        mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result(
        )

        def create_multiple_sla_results():
            mock_vector = Mock()
            mock_vector.probe_hosts.return_value = self.HOSTNAMES
            return mock_vector

        with contextlib.nested(
                patch('apache.aurora.client.api.SchedulerProxy',
                      return_value=mock_scheduler_proxy),
                patch(
                    'apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
                    return_value=create_multiple_sla_results()),
                patch('apache.aurora.client.commands.maintenance.CLUSTERS',
                      new=self.TEST_CLUSTERS),
                patch('twitter.common.app.get_options',
                      return_value=mock_options)):

            host_drain([self.TEST_CLUSTER])

            mock_scheduler_proxy.startMaintenance.assert_called_with(
                Hosts(set(self.HOSTNAMES)))
    def test_perform_maintenance_hosts_failed_default_sla(self):
        with temporary_file() as fp:
            mock_options = self.make_mock_options()
            mock_options.post_drain_script = None
            mock_options.grouping = 'by_host'
            mock_options.unsafe_hosts_filename = fp.name

            mock_api, mock_scheduler_proxy = self.create_mock_api()
            mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result(
            )
            mock_scheduler_proxy.drainHosts.return_value = self.create_start_maintenance_result(
            )
            mock_vector = self.create_mock_probe_hosts_vector([
                self.create_probe_hosts(self.HOSTNAMES[0], 95, False, None),
                self.create_probe_hosts(self.HOSTNAMES[1], 95, False, None),
                self.create_probe_hosts(self.HOSTNAMES[2], 95, False, None)
            ])

            with contextlib.nested(
                    patch('apache.aurora.client.api.SchedulerProxy',
                          return_value=mock_scheduler_proxy),
                    patch(
                        'apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
                        return_value=mock_vector),
                    patch('apache.aurora.client.commands.maintenance.CLUSTERS',
                          new=self.TEST_CLUSTERS),
                    patch('twitter.common.app.get_options',
                          return_value=mock_options)):
                host_drain([self.TEST_CLUSTER])

                mock_scheduler_proxy.startMaintenance.assert_called_with(
                    Hosts(set(self.HOSTNAMES)))
    def test_perform_maintenance_hosts(self):
        mock_options = self.make_mock_options()
        mock_options.post_drain_script = 'callback'
        mock_options.grouping = 'by_host'

        def host_status_results(hostnames):
            if isinstance(hostnames, Hosts):
                return self.create_drained_status_result(hostnames)
            return self.create_maintenance_status_result()

        mock_api, mock_scheduler_proxy = self.create_mock_api()
        mock_callback = Mock()
        mock_scheduler_proxy.maintenanceStatus.side_effect = host_status_results
        mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result(
        )
        mock_scheduler_proxy.drainHosts.return_value = self.create_start_maintenance_result(
        )
        mock_vector = self.create_mock_probe_hosts_vector([
            self.create_probe_hosts(self.HOSTNAMES[0], 95, True, None),
            self.create_probe_hosts(self.HOSTNAMES[1], 95, True, None),
            self.create_probe_hosts(self.HOSTNAMES[2], 95, True, None)
        ])

        with contextlib.nested(
                patch('apache.aurora.client.api.SchedulerProxy',
                      return_value=mock_scheduler_proxy),
                patch(
                    'apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
                    return_value=mock_vector),
                patch('apache.aurora.client.commands.maintenance.CLUSTERS',
                      new=self.TEST_CLUSTERS),
                patch('apache.aurora.client.commands.maintenance.parse_script',
                      return_value=mock_callback),
                patch('twitter.common.app.get_options',
                      return_value=mock_options)):
            host_drain([self.TEST_CLUSTER])

            mock_scheduler_proxy.startMaintenance.assert_called_with(
                Hosts(set(self.HOSTNAMES)))
            assert mock_scheduler_proxy.maintenanceStatus.call_count == 3
            assert mock_scheduler_proxy.drainHosts.call_count == 3
            assert mock_callback.call_count == 3
  def test_perform_maintenance_hosts(self):
    mock_options = self.make_mock_options()
    mock_options.post_drain_script = 'callback'
    mock_options.grouping = 'by_host'

    def host_status_results(hostnames):
      if isinstance(hostnames, Hosts):
        return self.create_drained_status_result(hostnames)
      return self.create_maintenance_status_result()

    mock_api, mock_scheduler_proxy = self.create_mock_api()
    mock_callback = Mock()
    mock_scheduler_proxy.maintenanceStatus.side_effect = host_status_results
    mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result()
    mock_scheduler_proxy.drainHosts.return_value = self.create_start_maintenance_result()
    mock_vector = self.create_mock_probe_hosts_vector([
        self.create_probe_hosts(self.HOSTNAMES[0], 95, True, None),
        self.create_probe_hosts(self.HOSTNAMES[1], 95, True, None),
        self.create_probe_hosts(self.HOSTNAMES[2], 95, True, None)
    ])

    with contextlib.nested(
        patch('time.sleep'),
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
              return_value=mock_vector),
        patch('apache.aurora.client.commands.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
        patch('apache.aurora.client.commands.maintenance.parse_script', return_value=mock_callback),
        patch('twitter.common.app.get_options', return_value=mock_options)) as (
            mock_sleep, _, _, _, _, _):
      host_drain([self.TEST_CLUSTER])

      mock_scheduler_proxy.startMaintenance.assert_called_with(Hosts(set(self.HOSTNAMES)))
      #TODO(jsmith): Consider not mocking out sleep and instead refactoring
      assert mock_sleep.call_count == 3
      assert mock_scheduler_proxy.maintenanceStatus.call_count == 3
      assert mock_scheduler_proxy.drainHosts.call_count == 3
      assert mock_callback.call_count == 3
    def test_perform_maintenance_hosts_failed_custom_sla(self):
        with temporary_file() as fp:
            mock_options = self.make_mock_options()
            mock_options.post_drain_script = None
            mock_options.grouping = 'by_host'
            mock_options.percentage = 50
            mock_options.duration = '10m'
            mock_options.reason = 'Test overrides'
            mock_options.unsafe_hosts_filename = fp.name

            mock_api, mock_scheduler_proxy = self.create_mock_api()
            mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result(
            )
            mock_scheduler_proxy.drainHosts.return_value = self.create_start_maintenance_result(
            )
            mock_vector = self.create_mock_probe_hosts_vector([
                self.create_probe_hosts(self.HOSTNAMES[0], 95, False, None),
                self.create_probe_hosts(self.HOSTNAMES[1], 95, False, None),
                self.create_probe_hosts(self.HOSTNAMES[2], 95, False, None)
            ])

            with contextlib.nested(
                    patch('apache.aurora.client.api.SchedulerProxy',
                          return_value=mock_scheduler_proxy),
                    patch(
                        'apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
                        return_value=mock_vector),
                    patch('apache.aurora.client.commands.maintenance.CLUSTERS',
                          new=self.TEST_CLUSTERS),
                    patch('apache.aurora.admin.admin_util.log_admin_message'),
                    patch('twitter.common.app.get_options',
                          return_value=mock_options)) as (_, _, _, log, _):
                host_drain([self.TEST_CLUSTER])

                assert 'Test overrides' in log.call_args[0][1]
                mock_scheduler_proxy.startMaintenance.assert_called_with(
                    Hosts(set(self.HOSTNAMES)))
  def test_perform_maintenance_hosts_multiple_sla_groups_failure(self):
    mock_options = self.make_mock_options()
    mock_options.post_drain_script = None
    mock_options.grouping = 'by_host'
    mock_options.unsafe_hosts_filename = None

    mock_api, mock_scheduler_proxy = self.create_mock_api()
    mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result()

    def create_multiple_sla_results():
      mock_vector = Mock()
      mock_vector.probe_hosts.return_value = self.HOSTNAMES
      return mock_vector

    with contextlib.nested(
        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
        patch('apache.aurora.client.api.sla.Sla.get_domain_uptime_vector',
              return_value=create_multiple_sla_results()),
        patch('apache.aurora.client.commands.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
        patch('twitter.common.app.get_options', return_value=mock_options)):

      host_drain([self.TEST_CLUSTER])

      mock_scheduler_proxy.startMaintenance.assert_called_with(Hosts(set(self.HOSTNAMES)))