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 <[email protected]> Authored: Fri Jan 22 14:40:27 2016 -0800 Committer: Maxim Khutornenko <[email protected]> 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
