You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@aurora.apache.org by ma...@apache.org on 2016/01/22 23:43:32 UTC
aurora git commit: Deprecating TaskQuery in killTasks.
Repository: aurora
Updated Branches:
refs/heads/master 4b3d7bca9 -> 02beb1f12
Deprecating TaskQuery in killTasks.
Bugs closed: AURORA-1583
Reviewed at https://reviews.apache.org/r/42666/
Project: http://git-wip-us.apache.org/repos/asf/aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/02beb1f1
Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/02beb1f1
Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/02beb1f1
Branch: refs/heads/master
Commit: 02beb1f123d9884d43916a02a967174a8fc077cf
Parents: 4b3d7bc
Author: Maxim Khutornenko <ma...@apache.org>
Authored: Fri Jan 22 14:40:27 2016 -0800
Committer: Maxim Khutornenko <ma...@apache.org>
Committed: Fri Jan 22 14:40:27 2016 -0800
----------------------------------------------------------------------
NEWS | 2 +
.../thrift/org/apache/aurora/gen/api.thrift | 4 +-
.../thrift/SchedulerThriftInterface.java | 42 +++++++++-----
.../thrift/aop/AnnotatedAuroraAdmin.java | 4 +-
.../python/apache/aurora/client/api/__init__.py | 8 +--
.../http/api/security/HttpSecurityIT.java | 60 +++++++++++++++-----
.../ShiroAuthorizingParamInterceptorTest.java | 32 ++++++++---
.../thrift/SchedulerThriftInterfaceTest.java | 30 +++++++---
.../thrift/aop/AnnotatedAuroraAdminTest.java | 11 ++--
src/test/python/apache/aurora/api_util.py | 2 +-
.../aurora/client/api/test_scheduler_client.py | 10 +++-
11 files changed, 145 insertions(+), 60 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/NEWS
----------------------------------------------------------------------
diff --git a/NEWS b/NEWS
index c9ff025..f2798f6 100644
--- a/NEWS
+++ b/NEWS
@@ -30,6 +30,8 @@
universal set of parameters that should be used for every container that does not have parameters
explicitly configured at the job level.
- Removed scheduler flag `-extra_modules`.
+- Deprecated `TaskQuery` argument in `killTasks` thrift RPC to disallow killing tasks across
+ multiple roles. The new safer approach is using `JobKey` with `instances` instead.
0.11.0
------
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/api/src/main/thrift/org/apache/aurora/gen/api.thrift
----------------------------------------------------------------------
diff --git a/api/src/main/thrift/org/apache/aurora/gen/api.thrift b/api/src/main/thrift/org/apache/aurora/gen/api.thrift
index f099620..a93df21 100644
--- a/api/src/main/thrift/org/apache/aurora/gen/api.thrift
+++ b/api/src/main/thrift/org/apache/aurora/gen/api.thrift
@@ -1029,8 +1029,8 @@ service AuroraSchedulerManager extends ReadOnlyScheduler {
/** Restarts a batch of shards. */
Response restartShards(5: JobKey job, 3: set<i32> shardIds, 6: Lock lock)
- /** Initiates a kill on tasks. */
- Response killTasks(1: TaskQuery query, 3: Lock lock)
+ /** Initiates a kill on tasks. TODO(maxim): remove TaskQuery in AURORA-1591. */
+ Response killTasks(1: TaskQuery query, 3: Lock lock, 4: JobKey job, 5: set<i32> instances)
/**
* Adds new instances specified by the AddInstancesConfig. A job represented by the JobKey must be
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
index 69eab90..d1e3c83 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
@@ -418,24 +418,40 @@ class SchedulerThriftInterface implements AnnotatedAuroraAdmin {
}
}
- private static Query.Builder implicitKillQuery(TaskQuery mutableQuery) {
- Query.Builder query = Query.arbitrary(mutableQuery);
+ private static Query.Builder implicitKillQuery(Query.Builder query) {
// Unless statuses were specifically supplied, only attempt to kill active tasks.
return query.get().isSetStatuses() ? query : query.byStatus(ACTIVE_STATES);
}
@Override
- public Response killTasks(TaskQuery mutableQuery, Lock mutableLock) {
- requireNonNull(mutableQuery);
+ public Response killTasks(
+ @Nullable TaskQuery mutableQuery,
+ @Nullable Lock mutableLock,
+ @Nullable JobKey mutableJob,
+ @Nullable Set<Integer> instances) {
+
+ final Query.Builder query;
+ Response response = empty();
+ if (mutableQuery == null) {
+ IJobKey jobKey = JobKeys.assertValid(IJobKey.build(mutableJob));
+ if (instances == null || Iterables.isEmpty(instances)) {
+ query = implicitKillQuery(Query.jobScoped(jobKey));
+ } else {
+ query = implicitKillQuery(Query.instanceScoped(jobKey, instances));
+ }
+ } else {
+ requireNonNull(mutableQuery);
+ addMessage(response, "The TaskQuery field is deprecated.");
- if (mutableQuery.getJobName() != null && WHITESPACE.matchesAllOf(mutableQuery.getJobName())) {
- return invalidRequest(String.format("Invalid job name: '%s'", mutableQuery.getJobName()));
- }
+ if (mutableQuery.getJobName() != null && WHITESPACE.matchesAllOf(mutableQuery.getJobName())) {
+ return invalidRequest(String.format("Invalid job name: '%s'", mutableQuery.getJobName()));
+ }
- Query.Builder query = implicitKillQuery(mutableQuery);
- Preconditions.checkState(
- !query.get().isSetOwner(),
- "The owner field in a query should have been unset by Query.Builder.");
+ query = implicitKillQuery(Query.arbitrary(mutableQuery));
+ Preconditions.checkState(
+ !mutableQuery.isSetOwner(),
+ "The owner field in a query should have been unset by Query.Builder.");
+ }
return storage.write(storeProvider -> {
Iterable<IScheduledTask> tasks = storeProvider.getTaskStore().fetchTasks(query);
@@ -460,8 +476,8 @@ class SchedulerThriftInterface implements AnnotatedAuroraAdmin {
}
return tasksKilled
- ? ok()
- : addMessage(empty(), OK, NO_TASKS_TO_KILL_MESSAGE);
+ ? response.setResponseCode(OK)
+ : addMessage(response, OK, NO_TASKS_TO_KILL_MESSAGE);
});
}
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java b/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java
index f6669ef..c374343 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java
@@ -71,7 +71,9 @@ public interface AnnotatedAuroraAdmin extends AuroraAdmin.Iface {
@Override
Response killTasks(
@AuthorizingParam @Nullable TaskQuery query,
- @Nullable Lock lock) throws TException;
+ @Nullable Lock lock,
+ @AuthorizingParam @Nullable JobKey job,
+ @Nullable Set<Integer> instances) throws TException;
@Override
Response addInstances(
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/src/main/python/apache/aurora/client/api/__init__.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/api/__init__.py b/src/main/python/apache/aurora/client/api/__init__.py
index ac4e6f2..1b2ce4d 100644
--- a/src/main/python/apache/aurora/client/api/__init__.py
+++ b/src/main/python/apache/aurora/client/api/__init__.py
@@ -100,13 +100,11 @@ class AuroraClientAPI(object):
log.info("Killing tasks for job: %s" % job_key)
self._assert_valid_job_key(job_key)
- # Leave query.owner.user unset so the query doesn't filter jobs only submitted by a particular
- # user.
- query = job_key.to_thrift_query()
if instances is not None:
log.info("Instances to be killed: %s" % instances)
- query.instanceIds = frozenset([int(s) for s in instances])
- return self._scheduler_proxy.killTasks(query, lock)
+ instances = frozenset([int(s) for s in instances])
+
+ return self._scheduler_proxy.killTasks(None, lock, job_key.to_thrift(), instances)
def check_status(self, job_key):
self._assert_valid_job_key(job_key)
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/src/test/java/org/apache/aurora/scheduler/http/api/security/HttpSecurityIT.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/http/api/security/HttpSecurityIT.java b/src/test/java/org/apache/aurora/scheduler/http/api/security/HttpSecurityIT.java
index 3e811a4..dfe94d3 100644
--- a/src/test/java/org/apache/aurora/scheduler/http/api/security/HttpSecurityIT.java
+++ b/src/test/java/org/apache/aurora/scheduler/http/api/security/HttpSecurityIT.java
@@ -33,6 +33,7 @@ import com.google.inject.util.Modules;
import com.sun.jersey.api.client.ClientResponse;
import org.apache.aurora.gen.AuroraAdmin;
+import org.apache.aurora.gen.JobKey;
import org.apache.aurora.gen.Lock;
import org.apache.aurora.gen.Response;
import org.apache.aurora.gen.ResponseCode;
@@ -226,7 +227,7 @@ public class HttpSecurityIT extends AbstractJettyTest {
private void assertKillTasksFails(AuroraAdmin.Client client) throws TException {
try {
- client.killTasks(null, null);
+ client.killTasks(null, null, null, null);
fail("killTasks should fail.");
} catch (TTransportException e) {
// Expected.
@@ -235,41 +236,72 @@ public class HttpSecurityIT extends AbstractJettyTest {
@Test
public void testAuroraSchedulerManager() throws TException, ServletException, IOException {
- expect(auroraAdmin.killTasks(null, new Lock().setMessage("1"))).andReturn(OK);
- expect(auroraAdmin.killTasks(null, new Lock().setMessage("2"))).andReturn(OK);
+ expect(auroraAdmin.killTasks(null, new Lock().setMessage("1"), null, null)).andReturn(OK);
+ expect(auroraAdmin.killTasks(null, new Lock().setMessage("2"), null, null)).andReturn(OK);
- TaskQuery jobScopedQuery = Query.jobScoped(JobKeys.from("role", "env", "name")).get();
+ JobKey job = JobKeys.from("role", "env", "name").newBuilder();
+ TaskQuery jobScopedQuery = Query.jobScoped(IJobKey.build(job)).get();
TaskQuery adsScopedQuery = Query.jobScoped(ADS_STAGING_JOB).get();
- expect(auroraAdmin.killTasks(adsScopedQuery, null)).andReturn(OK);
+ expect(auroraAdmin.killTasks(adsScopedQuery, null, null, null)).andReturn(OK);
+ expect(auroraAdmin.killTasks(null, null, ADS_STAGING_JOB.newBuilder(), null)).andReturn(OK);
- expectShiroAfterAuthFilter().times(19);
+ expectShiroAfterAuthFilter().times(24);
replayAndStart();
- assertEquals(OK, getAuthenticatedClient(WFARNER).killTasks(null, new Lock().setMessage("1")));
- assertEquals(OK, getAuthenticatedClient(ROOT).killTasks(null, new Lock().setMessage("2")));
+ assertEquals(
+ OK,
+ getAuthenticatedClient(WFARNER).killTasks(null, new Lock().setMessage("1"), null, null));
+ assertEquals(
+ OK,
+ getAuthenticatedClient(ROOT).killTasks(null, new Lock().setMessage("2"), null, null));
+
assertEquals(
ResponseCode.INVALID_REQUEST,
- getAuthenticatedClient(UNPRIVILEGED).killTasks(null, null).getResponseCode());
+ getAuthenticatedClient(UNPRIVILEGED).killTasks(null, null, null, null).getResponseCode());
assertEquals(
ResponseCode.AUTH_FAILED,
getAuthenticatedClient(UNPRIVILEGED)
- .killTasks(jobScopedQuery, null)
+ .killTasks(jobScopedQuery, null, null, null)
+ .getResponseCode());
+ assertEquals(
+ ResponseCode.AUTH_FAILED,
+ getAuthenticatedClient(UNPRIVILEGED)
+ .killTasks(null, null, job, null)
.getResponseCode());
assertEquals(
ResponseCode.INVALID_REQUEST,
- getAuthenticatedClient(BACKUP_SERVICE).killTasks(null, null).getResponseCode());
+ getAuthenticatedClient(BACKUP_SERVICE).killTasks(null, null, null, null).getResponseCode());
assertEquals(
ResponseCode.AUTH_FAILED,
getAuthenticatedClient(BACKUP_SERVICE)
- .killTasks(jobScopedQuery, null)
+ .killTasks(jobScopedQuery, null, null, null)
+ .getResponseCode());
+ assertEquals(
+ ResponseCode.AUTH_FAILED,
+ getAuthenticatedClient(BACKUP_SERVICE)
+ .killTasks(null, null, job, null)
+ .getResponseCode());
+ assertEquals(
+ ResponseCode.AUTH_FAILED,
+ getAuthenticatedClient(DEPLOY_SERVICE)
+ .killTasks(jobScopedQuery, null, null, null)
.getResponseCode());
assertEquals(
ResponseCode.AUTH_FAILED,
getAuthenticatedClient(DEPLOY_SERVICE)
- .killTasks(jobScopedQuery, null)
+ .killTasks(null, null, job, null)
.getResponseCode());
- assertEquals(OK, getAuthenticatedClient(DEPLOY_SERVICE).killTasks(adsScopedQuery, null));
+ assertEquals(
+ OK,
+ getAuthenticatedClient(DEPLOY_SERVICE).killTasks(adsScopedQuery, null, null, null));
+ assertEquals(
+ OK,
+ getAuthenticatedClient(DEPLOY_SERVICE).killTasks(
+ null,
+ null,
+ ADS_STAGING_JOB.newBuilder(),
+ null));
assertKillTasksFails(getUnauthenticatedClient());
assertKillTasksFails(getAuthenticatedClient(INCORRECT));
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/src/test/java/org/apache/aurora/scheduler/http/api/security/ShiroAuthorizingParamInterceptorTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/http/api/security/ShiroAuthorizingParamInterceptorTest.java b/src/test/java/org/apache/aurora/scheduler/http/api/security/ShiroAuthorizingParamInterceptorTest.java
index 79e70fd..b5b405b 100644
--- a/src/test/java/org/apache/aurora/scheduler/http/api/security/ShiroAuthorizingParamInterceptorTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/http/api/security/ShiroAuthorizingParamInterceptorTest.java
@@ -33,7 +33,6 @@ import org.apache.aurora.gen.Response;
import org.apache.aurora.gen.ResponseCode;
import org.apache.aurora.gen.TaskQuery;
import org.apache.aurora.scheduler.base.JobKeys;
-import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.spi.Permissions.Domain;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.thrift.Responses;
@@ -71,7 +70,7 @@ public class ShiroAuthorizingParamInterceptorTest extends EasyMockTest {
subject = createMock(Subject.class);
statsProvider = createMock(StatsProvider.class);
thrift = createMock(AnnotatedAuroraAdmin.class);
- };
+ }
private void replayAndInitialize() {
expect(statsProvider.makeCounter(SHIRO_AUTHORIZATION_FAILURES))
@@ -126,23 +125,36 @@ public class ShiroAuthorizingParamInterceptorTest extends EasyMockTest {
@Test
public void testKillTasksWithWildcardPermission() throws TException {
- TaskQuery taskQuery = Query.unscoped().get();
Response response = Responses.ok();
+ // TODO(maxim): Remove wildcard (unscoped) permissions when TaskQuery is gone from killTasks
+ // AURORA-1592.
expect(subject.isPermitted(interceptor.makeWildcardPermission("killTasks")))
.andReturn(true);
- expect(thrift.killTasks(taskQuery, null))
+ expect(thrift.killTasks(new TaskQuery(), null, null, null))
.andReturn(response);
replayAndInitialize();
- assertSame(response, decoratedThrift.killTasks(taskQuery, null));
+ assertSame(response, decoratedThrift.killTasks(new TaskQuery(), null, null, null));
}
@Test
- public void testKillTasksWithoutWildcardPermission() throws TException {
- TaskQuery taskQuery = Query.unscoped().get();
+ public void testKillTasksWithTargetedPermission() throws TException {
+ expect(subject.isPermitted(interceptor.makeWildcardPermission("killTasks")))
+ .andReturn(false);
+ expect(subject.isPermitted(interceptor.makeTargetPermission("killTasks", JOB_KEY)))
+ .andReturn(false);
+
+ replayAndInitialize();
+ assertEquals(
+ ResponseCode.AUTH_FAILED,
+ decoratedThrift.killTasks(null, null, JOB_KEY.newBuilder(), null).getResponseCode());
+ }
+
+ @Test
+ public void testKillTasksInvalidJobKey() throws TException {
expect(subject.isPermitted(interceptor.makeWildcardPermission("killTasks")))
.andReturn(false);
@@ -150,7 +162,11 @@ public class ShiroAuthorizingParamInterceptorTest extends EasyMockTest {
assertEquals(
ResponseCode.INVALID_REQUEST,
- decoratedThrift.killTasks(taskQuery, null).getResponseCode());
+ decoratedThrift.killTasks(
+ null,
+ null,
+ JOB_KEY.newBuilder().setName(null),
+ null).getResponseCode());
}
@Test
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
index d12b56e..58d57b9 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterfaceTest.java
@@ -608,11 +608,13 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
control.replay();
- assertEquals(okEmptyResponse(), thrift.killTasks(query, null));
+ Response response = thrift.killTasks(query, null, null, null);
+ assertOkResponse(response);
+ assertMessageMatches(response, "The TaskQuery field is deprecated.");
}
@Test
- public void testKillQueryActive() throws Exception {
+ public void testJobScopedKillsActive() throws Exception {
Query.Builder query = Query.unscoped().byJob(JOB_KEY);
storageUtil.expectTaskFetch(query.active(), buildScheduledTask());
lockManager.validateIfLocked(LOCK_KEY, java.util.Optional.empty());
@@ -620,7 +622,19 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
control.replay();
- assertOkResponse(thrift.killTasks(query.get(), null));
+ assertOkResponse(thrift.killTasks(null, null, JOB_KEY.newBuilder(), null));
+ }
+
+ @Test
+ public void testInstanceScoped() throws Exception {
+ Query.Builder query = Query.instanceScoped(JOB_KEY, ImmutableSet.of(1)).active();
+ storageUtil.expectTaskFetch(query, buildScheduledTask());
+ lockManager.validateIfLocked(LOCK_KEY, java.util.Optional.empty());
+ expectTransitionsToKilling();
+
+ control.replay();
+
+ assertOkResponse(thrift.killTasks(null, null, JOB_KEY.newBuilder(), ImmutableSet.of(1)));
}
@Test
@@ -636,7 +650,9 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
control.replay();
- assertResponse(LOCK_ERROR, thrift.killTasks(query.get(), LOCK.newBuilder()));
+ assertResponse(
+ LOCK_ERROR,
+ thrift.killTasks(null, LOCK.newBuilder(), JOB_KEY.newBuilder(), null));
}
@Test
@@ -651,7 +667,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
control.replay();
- assertOkResponse(thrift.killTasks(query.get(), null));
+ assertOkResponse(thrift.killTasks(query.get(), null, null, null));
}
@Test
@@ -662,7 +678,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
control.replay();
- assertResponse(INVALID_REQUEST, thrift.killTasks(query, null));
+ assertResponse(INVALID_REQUEST, thrift.killTasks(query, null, null, null));
}
@Test
@@ -672,7 +688,7 @@ public class SchedulerThriftInterfaceTest extends EasyMockTest {
control.replay();
- Response response = thrift.killTasks(query.get(), null);
+ Response response = thrift.killTasks(null, null, JOB_KEY.newBuilder(), null);
assertOkResponse(response);
assertMessageMatches(response, SchedulerThriftInterface.NO_TASKS_TO_KILL_MESSAGE);
}
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/src/test/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdminTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdminTest.java b/src/test/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdminTest.java
index 065ad42..e4bc76b 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdminTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdminTest.java
@@ -25,7 +25,7 @@ import org.apache.aurora.gen.AuroraSchedulerManager;
import org.apache.aurora.scheduler.http.api.security.AuthorizingParam;
import org.junit.Test;
-import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
public class AnnotatedAuroraAdminTest {
@Test
@@ -45,11 +45,10 @@ public class AnnotatedAuroraAdminTest {
annotatedInvokable.getParameters(),
input -> input.getAnnotation(AuthorizingParam.class) != null);
- assertEquals(
- "Method " + invokable + " should have 1 " + AuthorizingParam.class.getName()
- + " annotation but " + annotatedParameters.size() + " were found.",
- 1,
- annotatedParameters.size());
+ assertFalse(
+ "Method " + invokable + " should have at least 1 " + AuthorizingParam.class.getName()
+ + " annotation but none were found.",
+ annotatedParameters.isEmpty());
}
}
}
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/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
index 5b2a538..9d44b88 100644
--- a/src/test/python/apache/aurora/api_util.py
+++ b/src/test/python/apache/aurora/api_util.py
@@ -85,7 +85,7 @@ class SchedulerThriftApiSpec(ReadOnlyScheduler.Iface):
def restartShards(self, job, shardIds, lock):
pass
- def killTasks(self, query, lock):
+ def killTasks(self, query, lock, jobKey, instances):
pass
def addInstances(self, config, lock):
http://git-wip-us.apache.org/repos/asf/aurora/blob/02beb1f1/src/test/python/apache/aurora/client/api/test_scheduler_client.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/api/test_scheduler_client.py b/src/test/python/apache/aurora/client/api/test_scheduler_client.py
index 2daa96c..4c4caaf 100644
--- a/src/test/python/apache/aurora/client/api/test_scheduler_client.py
+++ b/src/test/python/apache/aurora/client/api/test_scheduler_client.py
@@ -139,9 +139,13 @@ class TestSchedulerProxyInjection(unittest.TestCase):
self.make_scheduler_proxy().getJobs(ROLE)
def test_killTasks(self):
- self.mock_thrift_client.killTasks(IsA(TaskQuery)).AndReturn(DEFAULT_RESPONSE)
+ self.mock_thrift_client.killTasks(
+ IgnoreArg(),
+ IgnoreArg(),
+ IsA(JobKey),
+ IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
- self.make_scheduler_proxy().killTasks(TaskQuery())
+ self.make_scheduler_proxy().killTasks(None, None, JobKey(), set([0]))
def test_getQuota(self):
self.mock_thrift_client.getQuota(IgnoreArg()).AndReturn(DEFAULT_RESPONSE)
@@ -411,7 +415,7 @@ class TestSchedulerClient(unittest.TestCase):
client.get.return_value = mock_scheduler_client
proxy = scheduler_client.SchedulerProxy(Cluster(name='local'))
- proxy.killTasks(TaskQuery(), None)
+ proxy.killTasks(None, None, JobKey(), None)
assert mock_thrift_client.killTasks.call_count == 3