Ability to terminate jobs given a specific prefix

Author: mattBrzezinski

src/AWSBatch.jl

@@ -14,6 +14,7 @@ export BatchJob, ComputeEnvironment, BatchEnvironmentError, BatchJobError
 export JobQueue, JobDefinition, JobState, LogEvent
 export run_batch, describe, status, status_reason, wait, log_events, isregistered, register, deregister
 export list_job_queues, list_job_definitions, create_compute_environment, create_job_queue
+export terminate_jobs
 
 
 const logger = getlogger(@__MODULE__)
@@ -29,6 +30,7 @@ include("job_queue.jl")
 include("job_state.jl")
 include("job_definition.jl")
 include("batch_job.jl")
+include("utilities.jl")
 
 
 """

src/utilities.jl

@@ -0,0 +1,59 @@
+_incomplete_job(job) = !(job["status"] in ["SUCCEEDED", "FAILED"])
+_suffix_asterisk(prefix) = endswith(prefix, "*") ? prefix : string(prefix, "*")
+
+function _get_job_ids(job_queue, prefix)
+    # Check if the prefix provided ends w/ '*", if not append it
+    prefix = _suffix_asterisk(prefix)
+
+    resp = @mock Batch.list_jobs(
+        Dict(
+            "jobQueue" => job_queue,
+            "filters" => [Dict("name" => "JOB_NAME", "values" => [prefix])],
+        ),
+    )
+    jobs = resp["jobSummaryList"]
+
+    while haskey(resp, "nextToken")
+        resp = @mock Batch.list_jobs(
+            Dict("jobQueue" => job_queue, "nextToken" => resp["nextToken"])
+        )
+        append!(jobs, resp["jobSummaryList"])
+    end
+
+    filter!(j -> _incomplete_job(j), jobs)
+
+    return [j["jobId"] for j in jobs]
+end
+
+"""
+    terminate_jobs()
+
+Terminate all Batch jobs with a given prefix.
+
+# Arguments
+- `job_queue::JobQueue`: JobQueue where the jobs reside
+- `prefix::AbstractString`: Prefix for the jobs
+
+# Keywords
+- `reason::AbstractString=""`: Reason to terminate the jobs
+
+# Return
+- `Array{String}`: Terminated Job Ids
+"""
+function terminate_jobs(
+    job_queue::AbstractString, prefix::AbstractString; reason::AbstractString=""
+)
+    job_ids = _get_job_ids(job_queue, prefix)
+
+    for job_id in job_ids
+        @mock Batch.terminate_job(job_id, reason)
+    end
+
+    return job_ids
+end
+
+function terminate_jobs(
+    job_queue::JobQueue, prefix::AbstractString; reason::AbstractString=""
+)
+    return terminate_jobs(job_queue.arn, prefix; reason=reason)
+end

test/runtests.jl

@@ -1,4 +1,5 @@
 using AWS
+using AWS.AWSExceptions: AWSException
 using AWSTools.CloudFormation: stack_output
 using AWSBatch
 using Dates
@@ -8,7 +9,7 @@ using Memento.TestUtils: @test_log, @test_nolog
 using Mocking
 using Test
 
-using AWS.AWSExceptions: AWSException
+@service Batch
 
 Mocking.activate()
 
@@ -81,6 +82,7 @@ include("mock.jl")
             include("job_state.jl")
             include("batch_job.jl")
             include("run_batch.jl")
+            include("utilities.jl")
         end
     else
         warn(logger, "Skipping \"local\" tests. Set `ENV[\"TESTS\"] = \"local\"` to run.")

test/utilities.jl

@@ -0,0 +1,52 @@
+list_jobs_patch = @patch function AWSBatch.Batch.list_jobs(params)
+    response = Dict{String, Any}(
+        "jobSummaryList" => [
+            Dict(
+                "jobId" => "1",
+                "status" => "SUCCEEDED"
+            ),
+            Dict(
+                "jobId" => "2",
+                "status" => "FAILED"
+            ),
+            Dict(
+                "jobId" => "3",
+                "status" => "RUNNABLE"
+            ),
+        ]
+    )
+
+    if !haskey(params, "nextToken")
+        response["nextToken"] = "nextToken"
+    end
+
+    return response
+end
+
+terminate_job_patch = @patch AWSBatch.Batch.terminate_job(a...) = Dict()
+
+@testset "_incomplete_job" begin
+    @test_throws KeyError AWSBatch._incomplete_job(Dict())
+    @test !AWSBatch._incomplete_job(Dict("status" => "SUCCEEDED"))
+    @test !AWSBatch._incomplete_job(Dict("status" => "FAILED"))
+    @test AWSBatch._incomplete_job(Dict("status" => "FOOBAR"))
+end
+
+@testset "_suffix_asterisk" begin
+    @test AWSBatch._suffix_asterisk("foobar") == "foobar*"
+    @test AWSBatch._suffix_asterisk("foobar*") == "foobar*"
+end
+
+@testset "_get_job_ids" begin
+    apply(list_jobs_patch) do
+        response = AWSBatch._get_job_ids("foo", "bar")
+        @test response == ["3", "3"]
+    end
+end
+
+@testset "terminate_jobs" begin
+    apply([list_jobs_patch, terminate_job_patch]) do
+        response = terminate_jobs("foo", "bar")
+        @test response == ["3", "3"]
+    end
+end