aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ma...@apache.org
Subject [3/5] incubator-aurora git commit: Removing client v1 code.
Date Thu, 08 Jan 2015 01:07:15 GMT
http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/admin/test_admin_sla.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/admin/test_admin_sla.py b/src/test/python/apache/aurora/admin/test_admin_sla.py
new file mode 100644
index 0000000..225251b
--- /dev/null
+++ b/src/test/python/apache/aurora/admin/test_admin_sla.py
@@ -0,0 +1,410 @@
+#
+# 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 create_autospec, patch
+from twitter.common.contextutil import temporary_file
+
+from apache.aurora.admin.admin import sla_list_safe_domain, sla_probe_hosts
+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.common.aurora_job_key import AuroraJobKey
+
+from .util import AuroraClientCommandTest
+
+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 = create_autospec(spec=['exclude_filename', 'exclude_hosts', 'include_filename',
+        'include_hosts', 'override_filename', 'list_jobs', 'verbosity', 'disable_all_hooks',
+        'min_instance_count', 'grouping'], instance=True)
+
+    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 = create_autospec(spec=DomainUpTimeSlaVector, instance=True)
+    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.admin.admin.make_client',
+            new=create_autospec(spec=AuroraClientAPI)),
+        patch('apache.aurora.admin.admin.print_results'),
+        patch('apache.aurora.admin.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.admin.admin.make_client',
+              new=create_autospec(spec=AuroraClientAPI)),
+          patch('apache.aurora.admin.admin.print_results'),
+          patch('apache.aurora.admin.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.admin.admin.make_client',
+            new=create_autospec(spec=AuroraClientAPI)),
+        patch('apache.aurora.admin.admin.print_results'),
+        patch('apache.aurora.admin.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.admin.admin.make_client',
+              new=create_autospec(spec=AuroraClientAPI)),
+          patch('apache.aurora.admin.admin.print_results'),
+          patch('apache.aurora.admin.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.admin.admin.make_client',
+            new=create_autospec(spec=AuroraClientAPI)),
+        patch('apache.aurora.admin.admin.print_results'),
+        patch('apache.aurora.admin.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.admin.admin.make_client',
+              new=create_autospec(spec=AuroraClientAPI)),
+          patch('apache.aurora.admin.admin.print_results'),
+          patch('apache.aurora.admin.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.admin.admin.make_client',
+            new=create_autospec(spec=AuroraClientAPI)),
+        patch('apache.aurora.admin.admin.print_results'),
+        patch('apache.aurora.admin.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 = create_autospec(
+        spec=['hosts', 'filename', 'verbosity', 'min_instance_count', 'grouping'],
+        spec_set=False,
+        instance=True)
+    mock_options.hosts = hosts
+    mock_options.filename = filename
+    mock_options.verbosity = False
+    mock_options.grouping = grouping or DEFAULT_GROUPING
+    mock_options.min_instance_count = 1
+    return mock_options
+
+  @classmethod
+  def create_mock_vector(cls, result):
+    mock_vector = create_autospec(spec=DomainUpTimeSlaVector, instance=True)
+    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.admin.admin.make_client',
+            new=create_autospec(spec=AuroraClientAPI)),
+        patch('apache.aurora.admin.admin.print_results'),
+        patch('apache.aurora.admin.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.admin.admin.make_client',
+              new=create_autospec(spec=AuroraClientAPI)),
+          patch('apache.aurora.admin.admin.print_results'),
+          patch('apache.aurora.admin.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.'

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/admin/test_maintenance.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/admin/test_maintenance.py b/src/test/python/apache/aurora/admin/test_maintenance.py
new file mode 100644
index 0000000..cccc1e5
--- /dev/null
+++ b/src/test/python/apache/aurora/admin/test_maintenance.py
@@ -0,0 +1,324 @@
+#
+# 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 mock import Mock, patch
+from twitter.common.contextutil import temporary_file
+
+from apache.aurora.admin.host_maintenance import HostMaintenance
+from apache.aurora.admin.maintenance import host_activate, host_deactivate, host_drain, host_status
+
+from .util import AuroraClientCommandTest
+
+from gen.apache.aurora.api.ttypes import (
+    DrainHostsResult,
+    EndMaintenanceResult,
+    Hosts,
+    HostStatus,
+    MaintenanceMode,
+    MaintenanceStatusResult,
+    StartMaintenanceResult
+)
+
+
+class TestMaintenanceCommands(AuroraClientCommandTest):
+  HOSTNAMES = ['us-grf-20', 'us-jim-47', 'us-suz-01']
+
+  def make_mock_options(self):
+    mock_options = Mock()
+    mock_options.filename = None
+    mock_options.hosts = ','.join(self.HOSTNAMES)
+    mock_options.cluster = self.TEST_CLUSTER
+    mock_options.verbosity = False
+    mock_options.disable_all_hooks = False
+    mock_options.percentage = None
+    mock_options.duration = None
+    mock_options.reason = None
+    return mock_options
+
+  def create_host_statuses(self, maintenance_mode, skip_hosts=None):
+    return [HostStatus(host=hostname, mode=maintenance_mode) for hostname in self.HOSTNAMES
+            if not skip_hosts or hostname not in skip_hosts]
+
+  def create_start_maintenance_result(self, skip_hosts=None):
+    host_statuses = self.create_host_statuses(MaintenanceMode.SCHEDULED, skip_hosts)
+    response = self.create_simple_success_response()
+    response.result.startMaintenanceResult = StartMaintenanceResult(statuses=set(host_statuses))
+    return response
+
+  def create_end_maintenance_result(self):
+    host_statuses = self.create_host_statuses(MaintenanceMode.NONE)
+    response = self.create_simple_success_response()
+    response.result.endMaintenanceResult = EndMaintenanceResult(statuses=set(host_statuses))
+    return response
+
+  def create_drain_hosts_result(self):
+    host_statuses = self.create_host_statuses(MaintenanceMode.DRAINING)
+    response = self.create_simple_success_response()
+    response.result.drainHostsResult = DrainHostsResult(statuses=set(host_statuses))
+    return response
+
+  def create_maintenance_status_result(self):
+    host_statuses = self.create_host_statuses(MaintenanceMode.NONE)
+    response = self.create_simple_success_response()
+    response.result.maintenanceStatusResult = MaintenanceStatusResult(statuses=set(host_statuses))
+    return response
+
+  def create_drained_status_result(self, hosts):
+    host_statuses = [
+        HostStatus(host=hostname, mode=MaintenanceMode.DRAINED) for hostname in hosts.hostNames]
+    response = self.create_simple_success_response()
+    response.result.maintenanceStatusResult = MaintenanceStatusResult(statuses=set(host_statuses))
+    return response
+
+  def test_start_maintenance_hosts(self):
+    mock_options = self.make_mock_options()
+    mock_api, mock_scheduler_proxy = self.create_mock_api()
+    mock_scheduler_proxy.startMaintenance.return_value = self.create_start_maintenance_result()
+    with contextlib.nested(
+        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
+        patch('apache.aurora.admin.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
+        patch('twitter.common.app.get_options', return_value=mock_options)):
+      host_deactivate([self.TEST_CLUSTER])
+
+      mock_scheduler_proxy.startMaintenance.assert_called_with(Hosts(set(self.HOSTNAMES)))
+
+  def test_end_maintenance_hosts(self):
+    mock_options = self.make_mock_options()
+    mock_api, mock_scheduler_proxy = self.create_mock_api()
+    mock_scheduler_proxy.endMaintenance.return_value = self.create_end_maintenance_result()
+    mock_scheduler_proxy.maintenanceStatus.return_value = self.create_maintenance_status_result()
+    with contextlib.nested(
+        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
+        patch('apache.aurora.admin.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
+        patch('twitter.common.app.get_options', return_value=mock_options)):
+      host_activate([self.TEST_CLUSTER])
+
+      mock_scheduler_proxy.endMaintenance.assert_called_with(Hosts(set(self.HOSTNAMES)))
+      mock_scheduler_proxy.maintenanceStatus.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.admin.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
+        patch('apache.aurora.admin.maintenance.parse_script', return_value=mock_callback),
+        patch('threading._Event.wait'),
+        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_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.admin.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
+        patch('threading._Event.wait'),
+        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.admin.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
+          patch('threading._Event.wait'),
+          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_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)
+      ])
+      mock_wait = Mock()
+
+      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.admin.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
+          patch('apache.aurora.admin.admin_util.log_admin_message'),
+          patch('threading._Event.wait', return_value=mock_wait),
+          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)))
+        mock_wait.called_once_with(HostMaintenance.MAX_STATUS_WAIT)
+
+  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.admin.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
+        patch('threading._Event.wait'),
+        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_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.admin.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_reason_missing(self):
+    mock_options = self.make_mock_options()
+    mock_options.grouping = 'by_host'
+    mock_options.percentage = 50
+    mock_options.duration = '10m'
+
+    with contextlib.nested(
+        patch('twitter.common.app.get_options', return_value=mock_options)):
+      self.assertRaises(SystemExit, host_drain, [self.TEST_CLUSTER])
+
+  def test_host_maintenance_status(self):
+    mock_options = self.make_mock_options()
+    mock_api, mock_scheduler_proxy = self.create_mock_api()
+    mock_scheduler_proxy.maintenanceStatus.return_value = self.create_maintenance_status_result()
+    with contextlib.nested(
+        patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
+        patch('apache.aurora.admin.maintenance.CLUSTERS', new=self.TEST_CLUSTERS),
+        patch('twitter.common.app.get_options', return_value=mock_options)):
+      host_status([self.TEST_CLUSTER])
+
+      mock_scheduler_proxy.maintenanceStatus.assert_called_with(Hosts(set(self.HOSTNAMES)))

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/admin/util.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/admin/util.py b/src/test/python/apache/aurora/admin/util.py
new file mode 100644
index 0000000..3570407
--- /dev/null
+++ b/src/test/python/apache/aurora/admin/util.py
@@ -0,0 +1,84 @@
+#
+# 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 unittest
+from collections import defaultdict
+
+from mock import create_autospec
+
+from apache.aurora.client.api.sla import DomainUpTimeSlaVector, JobUpTimeDetails
+from apache.aurora.client.hooks.hooked_api import HookedAuroraClientAPI
+from apache.aurora.common.aurora_job_key import AuroraJobKey
+from apache.aurora.common.cluster import Cluster
+from apache.aurora.common.clusters import Clusters
+
+from ..api_util import SchedulerProxyApiSpec
+
+from gen.apache.aurora.api.ttypes import Response, ResponseCode, ResponseDetail, Result
+
+
+class AuroraClientCommandTest(unittest.TestCase):
+  @classmethod
+  def create_blank_response(cls, code, msg):
+    # TODO(wfarner): Don't use a mock here.
+    response = create_autospec(spec=Response, instance=True)
+    response.responseCode = code
+    response.result = create_autospec(spec=Result, instance=True)
+    response.details = [ResponseDetail(message=msg)]
+    return response
+
+  @classmethod
+  def create_simple_success_response(cls):
+    return cls.create_blank_response(ResponseCode.OK, 'OK')
+
+  @classmethod
+  def create_mock_api(cls):
+    """Builds up a mock API object, with a mock SchedulerProxy"""
+    mock_scheduler_proxy = create_autospec(
+        spec=SchedulerProxyApiSpec,
+        spec_set=False,
+        instance=True)
+    mock_scheduler_proxy.url = "http://something_or_other"
+    mock_scheduler_proxy.scheduler_client.return_value = mock_scheduler_proxy
+    mock_api = create_autospec(spec=HookedAuroraClientAPI)
+    mock_api.scheduler_proxy = mock_scheduler_proxy
+    return mock_api, mock_scheduler_proxy
+
+  TEST_CLUSTER = 'west'
+
+  TEST_CLUSTERS = Clusters([Cluster(
+      name='west',
+      zk='zookeeper.example.com',
+      scheduler_zk_path='/foo/bar',
+      auth_mechanism='UNAUTHENTICATED')])
+
+  @classmethod
+  def create_mock_probe_hosts_vector(cls, side_effects):
+    mock_vector = create_autospec(spec=DomainUpTimeSlaVector, instance=True)
+    mock_vector.probe_hosts.side_effect = side_effects
+    return mock_vector
+
+  @classmethod
+  def create_probe_hosts(cls, hostname, predicted, safe, safe_in):
+    hosts = defaultdict(list)
+    job = AuroraJobKey.from_path('west/role/env/job-%s' % hostname)
+    hosts[hostname].append(JobUpTimeDetails(job, predicted, safe, safe_in))
+    return [hosts]
+
+  #TODO(wfarner): Remove this, force tests to call out their flags.
+  @classmethod
+  def setup_mock_options(cls):
+    mock_options = create_autospec(spec=['verbosity'], instance=True)
+    mock_options.verbosity = 'verbose'
+    return mock_options

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/api_util.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/api_util.py b/src/test/python/apache/aurora/api_util.py
new file mode 100644
index 0000000..7059955
--- /dev/null
+++ b/src/test/python/apache/aurora/api_util.py
@@ -0,0 +1,125 @@
+#
+# 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.
+#
+
+from apache.aurora.client.api.scheduler_client import SchedulerProxy
+
+from gen.apache.aurora.api import ReadOnlyScheduler
+
+
+class SchedulerThriftApiSpec(ReadOnlyScheduler.Iface):
+  """
+  A concrete definition of the thrift API used by the client. Intended primarily to be used as a
+  spec definition for unit testing, since the client effectively augments function signatures by
+  allowing callers to omit the session argument (in SchedulerProxy). These signatures should be
+  identical to those in AuroraAdmin.Iface, with the session removed.
+  """
+
+  def setQuota(self, ownerRole, quota):
+    pass
+
+  def forceTaskState(self, taskId, status):
+    pass
+
+  def performBackup(self):
+    pass
+
+  def listBackups(self):
+    pass
+
+  def stageRecovery(self, backupId):
+    pass
+
+  def queryRecovery(self, query):
+    pass
+
+  def deleteRecoveryTasks(self, query):
+    pass
+
+  def commitRecovery(self):
+    pass
+
+  def unloadRecovery(self):
+    pass
+
+  def startMaintenance(self, hosts):
+    pass
+
+  def drainHosts(self, hosts):
+    pass
+
+  def maintenanceStatus(self, hosts):
+    pass
+
+  def endMaintenance(self, hosts):
+    pass
+
+  def snapshot(self):
+    pass
+
+  def rewriteConfigs(self, request):
+    pass
+
+  def createJob(self, description, lock):
+    pass
+
+  def scheduleCronJob(self, description, lock):
+    pass
+
+  def descheduleCronJob(self, job, lock):
+    pass
+
+  def startCronJob(self, job):
+    pass
+
+  def restartShards(self, job, shardIds, lock):
+    pass
+
+  def killTasks(self, query, lock):
+    pass
+
+  def addInstances(self, config, lock):
+    pass
+
+  def acquireLock(self, lockKey):
+    pass
+
+  def releaseLock(self, lock, validation):
+    pass
+
+  def replaceCronTemplate(self, config, lock):
+    pass
+
+  def startJobUpdate(self, request):
+    pass
+
+  def pauseJobUpdate(self, jobKey):
+    pass
+
+  def resumeJobUpdate(self, jobKey):
+    pass
+
+  def abortJobUpdate(self, jobKey):
+    pass
+
+  def pulseJobUpdate(self, updateId):
+    pass
+
+
+class SchedulerProxyApiSpec(SchedulerThriftApiSpec, SchedulerProxy):
+  """
+  A concrete definition of the API provided by SchedulerProxy.
+  """
+
+  def url(self):
+    pass

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/BUILD
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/BUILD b/src/test/python/apache/aurora/client/BUILD
index e40c7a6..c55adfe 100644
--- a/src/test/python/apache/aurora/client/BUILD
+++ b/src/test/python/apache/aurora/client/BUILD
@@ -19,7 +19,6 @@ python_test_suite(name = 'all',
     ':config',
     'src/test/python/apache/aurora/client/api:all',
     'src/test/python/apache/aurora/client/cli:all',
-    'src/test/python/apache/aurora/client/commands:all',
     'src/test/python/apache/aurora/client/hooks:all',
   ],
 )

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/api/BUILD
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/BUILD b/src/test/python/apache/aurora/client/api/BUILD
index 2c0c407..65b378b 100644
--- a/src/test/python/apache/aurora/client/api/BUILD
+++ b/src/test/python/apache/aurora/client/api/BUILD
@@ -15,7 +15,6 @@
 python_test_suite(name = 'all',
   dependencies = [
     ':api',
-    ':disambiguator',
     ':instance_watcher',
     ':job_monitor',
     ':mux',
@@ -39,16 +38,6 @@ python_tests(name = 'api',
   ],
 )
 
-python_tests(
-  name = 'disambiguator',
-  sources = ['test_disambiguator.py'],
-  dependencies = [
-    '3rdparty/python:mox',
-    '3rdparty/python:twitter.common.app',
-    'src/main/python/apache/aurora/client/api:disambiguator',
-  ],
-)
-
 python_tests(name = 'job_monitor',
   sources = ['test_job_monitor.py'],
   dependencies = [

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/api/api_util.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/api_util.py b/src/test/python/apache/aurora/client/api/api_util.py
deleted file mode 100644
index 7059955..0000000
--- a/src/test/python/apache/aurora/client/api/api_util.py
+++ /dev/null
@@ -1,125 +0,0 @@
-#
-# 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.
-#
-
-from apache.aurora.client.api.scheduler_client import SchedulerProxy
-
-from gen.apache.aurora.api import ReadOnlyScheduler
-
-
-class SchedulerThriftApiSpec(ReadOnlyScheduler.Iface):
-  """
-  A concrete definition of the thrift API used by the client. Intended primarily to be used as a
-  spec definition for unit testing, since the client effectively augments function signatures by
-  allowing callers to omit the session argument (in SchedulerProxy). These signatures should be
-  identical to those in AuroraAdmin.Iface, with the session removed.
-  """
-
-  def setQuota(self, ownerRole, quota):
-    pass
-
-  def forceTaskState(self, taskId, status):
-    pass
-
-  def performBackup(self):
-    pass
-
-  def listBackups(self):
-    pass
-
-  def stageRecovery(self, backupId):
-    pass
-
-  def queryRecovery(self, query):
-    pass
-
-  def deleteRecoveryTasks(self, query):
-    pass
-
-  def commitRecovery(self):
-    pass
-
-  def unloadRecovery(self):
-    pass
-
-  def startMaintenance(self, hosts):
-    pass
-
-  def drainHosts(self, hosts):
-    pass
-
-  def maintenanceStatus(self, hosts):
-    pass
-
-  def endMaintenance(self, hosts):
-    pass
-
-  def snapshot(self):
-    pass
-
-  def rewriteConfigs(self, request):
-    pass
-
-  def createJob(self, description, lock):
-    pass
-
-  def scheduleCronJob(self, description, lock):
-    pass
-
-  def descheduleCronJob(self, job, lock):
-    pass
-
-  def startCronJob(self, job):
-    pass
-
-  def restartShards(self, job, shardIds, lock):
-    pass
-
-  def killTasks(self, query, lock):
-    pass
-
-  def addInstances(self, config, lock):
-    pass
-
-  def acquireLock(self, lockKey):
-    pass
-
-  def releaseLock(self, lock, validation):
-    pass
-
-  def replaceCronTemplate(self, config, lock):
-    pass
-
-  def startJobUpdate(self, request):
-    pass
-
-  def pauseJobUpdate(self, jobKey):
-    pass
-
-  def resumeJobUpdate(self, jobKey):
-    pass
-
-  def abortJobUpdate(self, jobKey):
-    pass
-
-  def pulseJobUpdate(self, updateId):
-    pass
-
-
-class SchedulerProxyApiSpec(SchedulerThriftApiSpec, SchedulerProxy):
-  """
-  A concrete definition of the API provided by SchedulerProxy.
-  """
-
-  def url(self):
-    pass

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/api/test_api.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_api.py b/src/test/python/apache/aurora/client/api/test_api.py
index 7390d1c..ff1aff2 100644
--- a/src/test/python/apache/aurora/client/api/test_api.py
+++ b/src/test/python/apache/aurora/client/api/test_api.py
@@ -21,7 +21,7 @@ from apache.aurora.common.cluster import Cluster
 from apache.aurora.config import AuroraConfig
 from apache.aurora.config.schema.base import UpdateConfig
 
-from .api_util import SchedulerThriftApiSpec
+from ...api_util import SchedulerThriftApiSpec
 
 from gen.apache.aurora.api.ttypes import (
     JobConfiguration,

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/api/test_disambiguator.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_disambiguator.py b/src/test/python/apache/aurora/client/api/test_disambiguator.py
deleted file mode 100644
index 72e98d2..0000000
--- a/src/test/python/apache/aurora/client/api/test_disambiguator.py
+++ /dev/null
@@ -1,125 +0,0 @@
-#
-# 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 functools
-
-import mox
-import pytest
-
-from apache.aurora.client.api import AuroraClientAPI
-from apache.aurora.client.api.disambiguator import LiveJobDisambiguator
-from apache.aurora.common.aurora_job_key import AuroraJobKey
-from apache.aurora.common.cluster import Cluster
-
-from gen.apache.aurora.api.constants import ResponseCode
-from gen.apache.aurora.api.ttypes import (
-    GetJobsResult,
-    JobConfiguration,
-    JobKey,
-    Response,
-    ResponseDetail,
-    Result
-)
-
-TEST_CLUSTER = Cluster(name='smf1')
-
-
-class LiveJobDisambiguatorTest(mox.MoxTestBase):
-  CLUSTER = TEST_CLUSTER
-  ROLE = 'mesos'
-  ENV = 'test'
-  NAME = 'labrat'
-  JOB_PATH = 'smf1/mesos/test/labrat'
-  CONFIG_FILE = 'abc.aurora'
-
-  def setUp(self):
-    super(LiveJobDisambiguatorTest, self).setUp()
-    self._api = self.mox.CreateMock(AuroraClientAPI)
-    self._api.cluster = self.CLUSTER
-    self._options = self.mox.CreateMockAnything()
-    self._options.cluster = self.CLUSTER
-
-  def test_ambiguous_property(self):
-    assert LiveJobDisambiguator(self._api, self.ROLE, None, self.NAME).ambiguous
-    assert not LiveJobDisambiguator(self._api, self.ROLE, self.ENV, self.NAME).ambiguous
-
-  def _expect_get_jobs(self, *envs):
-    self._api.get_jobs(self.ROLE).AndReturn(Response(
-      responseCode=ResponseCode.OK,
-      details=[ResponseDetail(message='Mock OK')],
-      result=Result(getJobsResult=GetJobsResult(
-        configs=set(JobConfiguration(key=JobKey(role=self.ROLE, environment=env, name=self.NAME))
-        for env in envs)))))
-
-  def _try_disambiguate_ambiguous(self):
-    return LiveJobDisambiguator._disambiguate_or_die(self._api, self.ROLE, None, self.NAME)
-
-  def test_disambiguate_or_die_ambiguous(self):
-    self._expect_get_jobs('test')
-    self._expect_get_jobs('prod')
-    self._expect_get_jobs('devel', 'test')
-    self._expect_get_jobs()
-
-    self.mox.ReplayAll()
-
-    _, _, env1, _ = self._try_disambiguate_ambiguous()
-    assert env1 == 'test'
-
-    _, _, env2, _ = self._try_disambiguate_ambiguous()
-    assert env2 == 'prod'
-
-    with pytest.raises(SystemExit):
-      self._try_disambiguate_ambiguous()
-
-    with pytest.raises(SystemExit):
-      self._try_disambiguate_ambiguous()
-
-  def test_disambiguate_job_path_or_die_unambiguous(self):
-    key = LiveJobDisambiguator._disambiguate_or_die(self._api, self.ROLE, self.ENV, self.NAME)
-    cluster_name, role, env, name = key
-    assert cluster_name == self.CLUSTER.name
-    assert role == self.ROLE
-    assert env == self.ENV
-    assert name == self.NAME
-
-  def test_disambiguate_args_or_die_unambiguous_with_no_config(self):
-    expected = (self._api, AuroraJobKey(self.CLUSTER.name, self.ROLE, self.ENV, self.NAME), None)
-    result = LiveJobDisambiguator.disambiguate_args_or_die([self.JOB_PATH], None,
-        client_factory=lambda *_: self._api)
-    assert result == expected
-
-  def test_disambiguate_args_or_die_unambiguous_with_config(self):
-    expected = (self._api,
-        AuroraJobKey(self.CLUSTER.name, self.ROLE, self.ENV, self.NAME), self.CONFIG_FILE)
-    result = LiveJobDisambiguator.disambiguate_args_or_die([self.JOB_PATH, self.CONFIG_FILE], None,
-        client_factory=lambda *_: self._api)
-    assert result == expected
-
-  def test_disambiguate_args_or_die_ambiguous(self):
-    self._expect_get_jobs('test')
-    self._expect_get_jobs('prod', 'devel')
-    self._expect_get_jobs()
-
-    disambiguate_args_or_die = functools.partial(LiveJobDisambiguator.disambiguate_args_or_die,
-        (self.ROLE, self.NAME), self._options, lambda *_: self._api)
-
-    self.mox.ReplayAll()
-
-    disambiguate_args_or_die()
-
-    with pytest.raises(SystemExit):
-      disambiguate_args_or_die()
-
-    with pytest.raises(SystemExit):
-      disambiguate_args_or_die()

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/api/test_job_monitor.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_job_monitor.py b/src/test/python/apache/aurora/client/api/test_job_monitor.py
index 1e7a879..ccc8b55 100644
--- a/src/test/python/apache/aurora/client/api/test_job_monitor.py
+++ b/src/test/python/apache/aurora/client/api/test_job_monitor.py
@@ -18,7 +18,7 @@ from mock import create_autospec
 from apache.aurora.client.api.job_monitor import JobMonitor
 from apache.aurora.common.aurora_job_key import AuroraJobKey
 
-from .api_util import SchedulerThriftApiSpec
+from ...api_util import SchedulerThriftApiSpec
 
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/api/test_quota_check.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_quota_check.py b/src/test/python/apache/aurora/client/api/test_quota_check.py
index 2c9bef1..ddd9466 100644
--- a/src/test/python/apache/aurora/client/api/test_quota_check.py
+++ b/src/test/python/apache/aurora/client/api/test_quota_check.py
@@ -19,7 +19,7 @@ from mock import call, create_autospec, patch
 
 from apache.aurora.client.api.quota_check import CapacityRequest, print_quota, QuotaCheck
 
-from .api_util import SchedulerThriftApiSpec
+from ...api_util import SchedulerThriftApiSpec
 
 from gen.apache.aurora.api.ttypes import (
     GetQuotaResult,

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/api/test_task_util.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_task_util.py b/src/test/python/apache/aurora/client/api/test_task_util.py
index 048aff6..eda326d 100644
--- a/src/test/python/apache/aurora/client/api/test_task_util.py
+++ b/src/test/python/apache/aurora/client/api/test_task_util.py
@@ -19,7 +19,7 @@ from mock import create_autospec
 from apache.aurora.client.api.scheduler_mux import SchedulerMux
 from apache.aurora.client.api.task_util import StatusMuxHelper
 
-from ..api.api_util import SchedulerThriftApiSpec
+from ...api_util import SchedulerThriftApiSpec
 
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/cli/BUILD
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/BUILD b/src/test/python/apache/aurora/client/cli/BUILD
index 4b5817e..bbac5c8 100644
--- a/src/test/python/apache/aurora/client/cli/BUILD
+++ b/src/test/python/apache/aurora/client/cli/BUILD
@@ -16,7 +16,6 @@ python_test_suite(
   name = 'all',
   dependencies = [
     ':api',
-    ':bridge',
     ':client',
     ':command_hooks',
     ':cron',
@@ -63,7 +62,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -87,15 +85,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
-  ]
-)
-
-python_tests(
-  name = 'bridge',
-  sources = [ 'test_bridge.py' ],
-  dependencies = [
-    'src/main/python/apache/aurora/client/cli:bridge',
   ]
 )
 
@@ -108,7 +97,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -122,7 +110,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -135,7 +122,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -156,7 +142,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -171,7 +156,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -185,7 +169,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -198,7 +181,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -211,7 +193,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 
@@ -224,7 +205,6 @@ python_tests(
     '3rdparty/python:twitter.common.contextutil',
     'src/main/python/apache/aurora/client/cli',
     'src/main/python/apache/aurora/client/cli:client_lib',
-    'src/test/python/apache/aurora/client/commands:util'
   ]
 )
 

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/cli/test_bridge.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_bridge.py b/src/test/python/apache/aurora/client/cli/test_bridge.py
deleted file mode 100644
index 6c40bf7..0000000
--- a/src/test/python/apache/aurora/client/cli/test_bridge.py
+++ /dev/null
@@ -1,85 +0,0 @@
-#
-# 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 unittest
-
-from apache.aurora.client.cli.bridge import Bridge, CommandProcessor
-
-
-class CommandOne(CommandProcessor):
-  def get_commands(self):
-    return ['one', 'two', 'three']
-
-  def execute(self, args):
-    return 1
-
-
-class CommandTwo(CommandProcessor):
-  def get_commands(self):
-    return ['three', 'four', 'five']
-
-  def execute(self, args):
-    return 2
-
-
-class CommandThree(CommandProcessor):
-  def get_commands(self):
-    return ['six']
-
-  def execute(self, args):
-    return '3[%s]' % args[1]
-
-
-class TestBridgedCommandLine(unittest.TestCase):
-  def setUp(self):
-    self.one = CommandOne()
-    self.two = CommandTwo()
-    self.three = CommandThree()
-
-  def test_bridge_with_default_three(self):
-    bridge = Bridge([self.one, self.two, self.three], default=self.three)
-    assert bridge.execute(['test', 'one']) == 1
-    assert bridge.execute(['test', 'two']) == 1
-    assert bridge.execute(['test', 'three']) == 1
-    assert bridge.execute(['test', 'four']) == 2
-    assert bridge.execute(['test', 'five']) == 2
-    assert bridge.execute(['test', 'six']) == '3[six]'
-    assert bridge.execute(['test', 'seven']) == '3[seven]'
-    assert bridge.execute(['test', 'eight']) == '3[eight]'
-
-  def test_bridge_with_default_one(self):
-    bridge = Bridge([self.one, self.two, self.three], default=self.one)
-    assert bridge.execute(['test', 'one']) == 1
-    assert bridge.execute(['test', 'two']) == 1
-    assert bridge.execute(['test', 'three']) == 1
-    assert bridge.execute(['test', 'four']) == 2
-    assert bridge.execute(['test', 'five']) == 2
-    assert bridge.execute(['test', 'six']) == '3[six]'
-    assert bridge.execute(['test', 'seven']) == 1
-    assert bridge.execute(['test', 'eight']) == 1
-
-  def test_bridge_with_no_default(self):
-    bridge = Bridge([self.one, self.two, self.three])
-    assert bridge.execute(['test', 'one']) == 1
-    assert bridge.execute(['test', 'two']) == 1
-    assert bridge.execute(['test', 'three']) == 1
-    assert bridge.execute(['test', 'four']) == 2
-    assert bridge.execute(['test', 'five']) == 2
-    assert bridge.execute(['test', 'six']) == '3[six]'
-    self.assertRaises(SystemExit, bridge.execute, ['test', 'seven'])
-
-  def test_bridge_ordering(self):
-    bridge1 = Bridge([self.one, self.two, self.three])
-    bridge2 = Bridge([self.two, self.one, self.three])
-    assert bridge1.execute(['test', 'three']) == 1
-    assert bridge2.execute(['test', 'three']) == 2

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/cli/test_cron.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_cron.py b/src/test/python/apache/aurora/client/cli/test_cron.py
index c748212..19889f8 100644
--- a/src/test/python/apache/aurora/client/cli/test_cron.py
+++ b/src/test/python/apache/aurora/client/cli/test_cron.py
@@ -16,7 +16,7 @@
 
 import contextlib
 
-from mock import create_autospec, patch
+from mock import patch
 from twitter.common.contextutil import temporary_file
 
 from apache.aurora.client.cli import (
@@ -29,7 +29,6 @@ from apache.aurora.client.cli.client import AuroraCommandLine
 from apache.aurora.common.aurora_job_key import AuroraJobKey
 from apache.aurora.config import AuroraConfig
 
-from ..api.api_util import SchedulerProxyApiSpec
 from .util import AuroraClientCommandTest, FakeAuroraCommandContext
 
 from gen.apache.aurora.api.ttypes import GetJobsResult, JobConfiguration, JobKey, Result
@@ -155,7 +154,6 @@ class TestCronNoun(AuroraClientCommandTest):
 
   def test_cron_status(self):
     (_, mock_scheduler_proxy) = self.create_mock_api()
-    mock_scheduler_proxy = create_autospec(spec=SchedulerProxyApiSpec)
     with contextlib.nested(
         patch('time.sleep'),
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/cli/util.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/util.py b/src/test/python/apache/aurora/client/cli/util.py
index b51ca2f..1fa1207 100644
--- a/src/test/python/apache/aurora/client/cli/util.py
+++ b/src/test/python/apache/aurora/client/cli/util.py
@@ -24,7 +24,7 @@ from apache.aurora.common.aurora_job_key import AuroraJobKey
 from apache.aurora.common.cluster import Cluster
 from apache.aurora.common.clusters import Clusters
 
-from ..api.api_util import SchedulerProxyApiSpec, SchedulerThriftApiSpec
+from ...api_util import SchedulerProxyApiSpec, SchedulerThriftApiSpec
 
 from gen.apache.aurora.api.ttypes import (
     AssignedTask,

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/commands/BUILD
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/BUILD b/src/test/python/apache/aurora/client/commands/BUILD
deleted file mode 100644
index 86ab65e..0000000
--- a/src/test/python/apache/aurora/client/commands/BUILD
+++ /dev/null
@@ -1,117 +0,0 @@
-#
-# 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.
-#
-
-python_test_suite(
-  name = 'all',
-  dependencies = [
-    ':admin',
-    ':core',
-    ':hooks',
-    ':maintenance',
-    ':run',
-    ':ssh',
-  ])
-
-python_tests(
-  name='core',
-  sources = [
-    'test_cancel_update.py',
-    'test_create.py',
-    'test_diff.py',
-    'test_kill.py',
-    'test_listjobs.py',
-    'test_restart.py',
-    'test_status.py',
-    'test_update.py',
-    'test_version.py'
-  ],
-  dependencies = [
-    ':util',
-    '3rdparty/python:mock',
-    '3rdparty/python:twitter.common.contextutil',
-    'src/main/python/apache/aurora/client/commands:core',
-    'api/src/main/thrift/org/apache/aurora/gen:py-thrift',
-  ])
-
-python_tests(
-  name='hooks',
-  sources = [ 'test_hooks.py' ],
-  dependencies = [
-    ':util',
-    '3rdparty/python:mock',
-    '3rdparty/python:twitter.common.contextutil',
-    'src/main/python/apache/aurora/client/commands:core',
-    'api/src/main/thrift/org/apache/aurora/gen:py-thrift',
-  ])
-
-python_tests(
-  name='admin',
-  sources = [
-    'test_admin.py',
-    'test_admin_sla.py',
-  ],
-  dependencies = [
-    ':util',
-    '3rdparty/python:mock',
-    '3rdparty/python:twitter.common.contextutil',
-    'src/main/python/apache/aurora/client/commands:admin',
-    'api/src/main/thrift/org/apache/aurora/gen:py-thrift',
-  ])
-
-python_tests(
-  name='ssh',
-  sources = [
-    'test_ssh.py',
-  ],
-  dependencies = [
-    ':util',
-    '3rdparty/python:mock',
-    '3rdparty/python:twitter.common.contextutil',
-    'src/main/python/apache/aurora/client/commands:ssh',
-    'api/src/main/thrift/org/apache/aurora/gen:py-thrift',
-  ])
-
-python_tests(
-  name='run',
-  sources = [
-    'test_run.py',
-  ],
-  dependencies = [
-    ':util',
-    '3rdparty/python:mock',
-    '3rdparty/python:twitter.common.contextutil',
-    'src/main/python/apache/aurora/client/commands:run',
-    'api/src/main/thrift/org/apache/aurora/gen:py-thrift',
-  ])
-
-python_tests(
-  name='maintenance',
-  sources = [
-    'test_maintenance.py',
-  ],
-  dependencies = [
-    ':util',
-    '3rdparty/python:mock',
-    'src/main/python/apache/aurora/client/commands:maintenance',
-  ])
-
-python_library(
-  name='util',
-  sources = [ 'util.py' ],
-  dependencies = [
-    '3rdparty/python:mock',
-    'src/main/python/apache/aurora/client/commands:core',
-    'api/src/main/thrift/org/apache/aurora/gen:py-thrift',
-  ]
-)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/commands/__init__.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/__init__.py b/src/test/python/apache/aurora/client/commands/__init__.py
deleted file mode 100644
index 0663a9a..0000000
--- a/src/test/python/apache/aurora/client/commands/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-#
-# 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.
-#

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/9a817f24/src/test/python/apache/aurora/client/commands/test_admin.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/commands/test_admin.py b/src/test/python/apache/aurora/client/commands/test_admin.py
deleted file mode 100644
index b093682..0000000
--- a/src/test/python/apache/aurora/client/commands/test_admin.py
+++ /dev/null
@@ -1,264 +0,0 @@
-#
-# 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 mock import create_autospec, patch, PropertyMock
-
-from apache.aurora.client.api import AuroraClientAPI
-from apache.aurora.client.api.scheduler_client import SchedulerClient, SchedulerProxy
-from apache.aurora.client.commands.admin import (
-    get_locks,
-    get_scheduler,
-    increase_quota,
-    query,
-    set_quota
-)
-
-from .util import AuroraClientCommandTest
-
-from gen.apache.aurora.api.ttypes import (
-    AssignedTask,
-    GetLocksResult,
-    GetQuotaResult,
-    Identity,
-    JobKey,
-    Lock,
-    LockKey,
-    ResourceAggregate,
-    Response,
-    ResponseCode,
-    ResponseDetail,
-    Result,
-    ScheduledTask,
-    ScheduleStatus,
-    ScheduleStatusResult,
-    TaskConfig,
-    TaskQuery
-)
-
-
-class TestQueryCommand(AuroraClientCommandTest):
-
-  @classmethod
-  def setup_mock_options(cls, force=False, shards=None, states='RUNNING', listformat=None):
-    mock_options = create_autospec(
-        spec=['force', 'shards', 'states', 'listformat', 'verbosity'],
-        spec_set=False,
-        instance=True)
-    mock_options.force = force
-    mock_options.shards = shards
-    mock_options.states = states
-    mock_options.listformat = listformat or '%role%/%jobName%/%instanceId% %status%'
-    mock_options.verbosity = False
-    return mock_options
-
-  @classmethod
-  def create_response(cls, tasks, response_code=ResponseCode.OK):
-    return Response(
-      responseCode=response_code,
-      details=[ResponseDetail(message='test')],
-      result=Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks)))
-
-  @classmethod
-  def create_task(cls):
-    return [ScheduledTask(
-        assignedTask=AssignedTask(
-            instanceId=0,
-            task=TaskConfig(owner=Identity(role='test_role'), jobName='test_job')),
-        status=ScheduleStatus.RUNNING
-    )]
-
-  @classmethod
-  def task_query(cls):
-    return TaskQuery(
-        role='test_role',
-        jobName='test_job',
-        instanceIds=set([0]),
-        statuses=set([ScheduleStatus.RUNNING]))
-
-  def test_query(self):
-    """Tests successful execution of the query command."""
-    mock_options = self.setup_mock_options(force=True, shards="0")
-    with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.admin.make_admin_client',
-            return_value=create_autospec(spec=AuroraClientAPI)),
-        patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS)
-    ) as (_, mock_make_admin_client, _):
-
-      api = mock_make_admin_client.return_value
-      api.query.return_value = self.create_response(self.create_task())
-
-      query([self.TEST_CLUSTER, 'test_role', 'test_job'], mock_options)
-
-      api.query.assert_called_with(self.task_query())
-
-  def test_query_fails(self):
-    """Tests failed execution of the query command."""
-    mock_options = self.setup_mock_options(shards="0")
-    with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.admin.make_admin_client',
-            return_value=create_autospec(spec=AuroraClientAPI)),
-        patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS)
-    ) as (_, mock_make_admin_client, _):
-
-      api = mock_make_admin_client.return_value
-      api.query.return_value = self.create_response(self.create_task())
-
-      try:
-        query([self.TEST_CLUSTER, 'test_role', 'test_job'], mock_options)
-      except SystemExit:
-        pass
-      else:
-        assert 'Expected exception is not raised'
-
-      api.query.assert_called_with(self.task_query())
-
-
-class TestIncreaseQuotaCommand(AuroraClientCommandTest):
-
-  @classmethod
-  def create_response(cls, quota, prod, non_prod, response_code=ResponseCode.OK):
-    return Response(
-        responseCode=response_code,
-        details=[ResponseDetail(message='test')],
-        result=Result(getQuotaResult=GetQuotaResult(
-            quota=quota, prodConsumption=prod, nonProdConsumption=non_prod))
-    )
-
-  def test_increase_quota(self):
-    """Tests successful execution of the increase_quota command."""
-    mock_options = self.setup_mock_options()
-    with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.admin.make_admin_client',
-            return_value=create_autospec(spec=AuroraClientAPI)),
-        patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS)
-    ) as (_, mock_make_admin_client, _):
-
-      api = mock_make_admin_client.return_value
-      role = 'test_role'
-      api.get_quota.return_value = self.create_response(
-          ResourceAggregate(20.0, 4000, 6000),
-          ResourceAggregate(15.0, 2000, 3000),
-          ResourceAggregate(6.0, 200, 600),
-      )
-      api.set_quota.return_value = self.create_simple_success_response()
-
-      increase_quota([self.TEST_CLUSTER, role, '4.0', '1MB', '1MB'])
-
-      api.set_quota.assert_called_with(role, 24.0, 4001, 6001)
-      assert type(api.set_quota.call_args[0][1]) == type(float())
-      assert type(api.set_quota.call_args[0][2]) == type(int())
-      assert type(api.set_quota.call_args[0][3]) == type(int())
-
-
-class TestSetQuotaCommand(AuroraClientCommandTest):
-
-  @classmethod
-  def create_response(cls, quota, prod, non_prod, response_code=None):
-    response_code = ResponseCode.OK if response_code is None else response_code
-    resp = Response(responseCode=response_code, details=[ResponseDetail(message='test')])
-    resp.result = Result(getQuotaResult=GetQuotaResult(
-      quota=quota, prodConsumption=prod, nonProdConsumption=non_prod))
-    return resp
-
-  def test_set_quota(self):
-    """Tests successful execution of the set_quota command."""
-    mock_options = self.setup_mock_options()
-    with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.admin.make_admin_client',
-              return_value=create_autospec(spec=AuroraClientAPI)),
-        patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS)
-    ) as (_, mock_make_admin_client, _):
-
-      api = mock_make_admin_client.return_value
-      role = 'test_role'
-      api.set_quota.return_value = self.create_simple_success_response()
-
-      set_quota([self.TEST_CLUSTER, role, '4.0', '10MB', '10MB'])
-
-      api.set_quota.assert_called_with(role, 4.0, 10, 10)
-      assert type(api.set_quota.call_args[0][1]) == type(float())
-      assert type(api.set_quota.call_args[0][2]) == type(int())
-      assert type(api.set_quota.call_args[0][3]) == type(int())
-
-
-class TestGetLocksCommand(AuroraClientCommandTest):
-
-  MESSAGE = 'test message'
-  USER = 'test user'
-  LOCKS = [Lock(
-    key=LockKey(job=JobKey('role', 'env', 'name')),
-    token='test token',
-    user=USER,
-    timestampMs='300',
-    message=MESSAGE)]
-
-  @classmethod
-  def create_response(cls, locks, response_code=None):
-    response_code = ResponseCode.OK if response_code is None else response_code
-    resp = Response(responseCode=response_code, details=[ResponseDetail(message='test')])
-    resp.result = Result(getLocksResult=GetLocksResult(locks=locks))
-    return resp
-
-  def test_get_locks(self):
-    """Tests successful execution of the get_locks command."""
-    mock_options = self.setup_mock_options()
-    with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.admin.make_admin_client',
-              return_value=create_autospec(spec=AuroraClientAPI)),
-        patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
-        patch('apache.aurora.client.commands.admin.print_results'),
-    ) as (_, mock_make_admin_client, _, mock_print_results):
-
-      api = mock_make_admin_client.return_value
-      api.get_locks.return_value = self.create_response(self.LOCKS)
-
-      get_locks([self.TEST_CLUSTER])
-
-      assert api.get_locks.call_count == 1
-      assert mock_print_results.call_count == 1
-      assert "'message': '%s'" % self.MESSAGE in mock_print_results.call_args[0][0][0]
-      assert "'user': '%s'" % self.USER in mock_print_results.call_args[0][0][0]
-
-
-class TestGetSchedulerCommand(AuroraClientCommandTest):
-
-  def test_get_scheduler(self):
-    """Tests successful execution of the get_scheduler command."""
-    mock_options = self.setup_mock_options()
-    mock_proxy = create_autospec(spec=SchedulerProxy, instance=True)
-    mock_scheduler_client = create_autospec(spec=SchedulerClient, instance=True)
-    mock_raw_url = PropertyMock(return_value="url")
-    mock_proxy.scheduler_client.return_value = mock_scheduler_client
-    mock_scheduler_client.raw_url = mock_raw_url
-
-    with contextlib.nested(
-        patch('twitter.common.app.get_options', return_value=mock_options),
-        patch('apache.aurora.client.commands.admin.make_admin_client',
-              return_value=create_autospec(spec=AuroraClientAPI)),
-        patch('apache.aurora.client.commands.admin.CLUSTERS', new=self.TEST_CLUSTERS),
-    ) as (_, mock_make_admin_client, _):
-
-      api = mock_make_admin_client.return_value
-      api.scheduler_proxy = PropertyMock(return_value=mock_proxy)
-
-      get_scheduler([self.TEST_CLUSTER])
-
-      mock_raw_url.assert_called_once()


Mime
View raw message