From 7692b64d6accf7a70f19c1eb65107612a4071d48 Mon Sep 17 00:00:00 2001
From: Dilum Aluthge <dilum@aluthge.com>
Date: Fri, 7 Feb 2025 14:52:30 -0500
Subject: [PATCH 1/2] Refactor the `SLURM_NTASKS` and `SLURM_JOB_ID`
 functionality out into separate utility functions, and add some more unit
 tests to increase code coverage

---
 src/slurmmanager.jl | 52 ++++++++++++++++++++++++---------------------
 test/runtests.jl    |  6 +++++-
 test/unit.jl        | 26 +++++++++++++++++++++++
 3 files changed, 59 insertions(+), 25 deletions(-)
 create mode 100644 test/unit.jl

diff --git a/src/slurmmanager.jl b/src/slurmmanager.jl
index f0175c5..2ea97a5 100644
--- a/src/slurmmanager.jl
+++ b/src/slurmmanager.jl
@@ -12,34 +12,38 @@ mutable struct SlurmManager <: ClusterManager
   srun_proc
 
   function SlurmManager(; launch_timeout=60.0, srun_post_exit_sleep=0.01)
+    ntasks = get_slurm_ntasks_int()
+    jobid = get_slurm_jobid_int()
 
-    jobid =
-    if "SLURM_JOB_ID" in keys(ENV)
-        ENV["SLURM_JOB_ID"]
-    elseif "SLURM_JOBID" in keys(ENV)
-        ENV["SLURM_JOBID"]
-    else
-        error("""
-              SlurmManager must be constructed inside a slurm allocation environemnt.
-              SLURM_JOB_ID or SLURM_JOBID must be defined.
-              """)
-    end
-
-    ntasks =
-    if "SLURM_NTASKS" in keys(ENV)
-      ENV["SLURM_NTASKS"]
-    else
-      error("""
-            SlurmManager must be constructed inside a slurm environment with a specified number of tasks.
-            SLURM_NTASKS must be defined.
-            """)
-    end
+    new(jobid, ntasks, launch_timeout, srun_post_exit_sleep, nothing)
+  end
+end
 
-    jobid = parse(Int, jobid)
-    ntasks = parse(Int, ntasks)
+function get_slurm_ntasks_int()
+  if haskey(ENV, "SLURM_NTASKS")
+    ntasks_str = ENV["SLURM_NTASKS"]
+  else
+    msg = "SlurmManager must be constructed inside a Slurm allocation environment." *
+          "SLURM_NTASKS must be defined."
+    error(msg)
+  end
+  ntasks_int = parse(Int, ntasks_str)::Int
+  return ntasks_int
+end
 
-    new(jobid, ntasks, launch_timeout, srun_post_exit_sleep, nothing)
+function get_slurm_jobid_int()
+  if haskey(ENV, "SLURM_JOB_ID")
+      jobid_str = ENV["SLURM_JOB_ID"]
+  elseif haskey(ENV, "SLURM_JOBID")
+      jobid_str = ENV["SLURM_JOBID"]
+  else
+      msg = "SlurmManager must be constructed inside a Slurm allocation environment." *
+            "SLURM_JOB_ID or SLURM_JOBID must be defined."
+      error(msg)
   end
+
+  jobid_int = parse(Int, jobid_str)::Int
+  return jobid_int
 end
 
 @static if Base.VERSION >= v"1.9.0"
diff --git a/test/runtests.jl b/test/runtests.jl
index 8772e5b..db3f4e0 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -6,7 +6,7 @@ import Distributed
 import Test
 
 # Bring some names into scope, just for convenience:
-using Test: @testset, @test
+using Test: @testset, @test, @test_throws, @test_logs
 
 const original_JULIA_DEBUG = strip(get(ENV, "JULIA_DEBUG", ""))
 if isempty(original_JULIA_DEBUG)
@@ -16,6 +16,10 @@ else
 end
 
 @testset "SlurmClusterManager.jl" begin
+  @testset "Unit tests" begin
+    include("unit.jl")
+  end
+  
   # test that slurm is available
   @test !(Sys.which("sinfo") === nothing)
 
diff --git a/test/unit.jl b/test/unit.jl
new file mode 100644
index 0000000..9d6f6d2
--- /dev/null
+++ b/test/unit.jl
@@ -0,0 +1,26 @@
+@testset "get_slurm_ntasks_int()" begin
+  x = withenv("SLURM_NTASKS" => "12") do
+    SlurmClusterManager.get_slurm_ntasks_int()
+  end
+  @test x == 12
+
+  withenv("SLURM_NTASKS" => nothing) do
+    @test_throws ErrorException SlurmClusterManager.get_slurm_ntasks_int()
+  end
+end
+
+@testset "get_slurm_jobid_int()" begin
+  x = withenv("SLURM_JOB_ID" => "34", "SLURM_JOBID" => nothing) do
+    SlurmClusterManager.get_slurm_jobid_int()
+  end
+  @test x == 34
+
+  x = withenv("SLURM_JOB_ID" => nothing, "SLURM_JOBID" => "56") do
+    SlurmClusterManager.get_slurm_jobid_int()
+  end
+  @test x == 56
+  
+  withenv("SLURM_JOB_ID" => nothing, "SLURM_JOBID" => nothing) do
+    @test_throws ErrorException SlurmClusterManager.get_slurm_jobid_int()
+  end
+end

From 092db155c2ae1d94d50d8c75a31ce81190224629 Mon Sep 17 00:00:00 2001
From: Dilum Aluthge <dilum@aluthge.com>
Date: Tue, 29 Apr 2025 00:06:44 -0400
Subject: [PATCH 2/2] Test suite: Add more test coverage for
 `warn_if_unexpected_params()` (only on Julia versions prior to 1.6) (#45)

Depends on:
- [ ] #44
---
 test/unit.jl | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/test/unit.jl b/test/unit.jl
index 9d6f6d2..687fcb3 100644
--- a/test/unit.jl
+++ b/test/unit.jl
@@ -24,3 +24,16 @@ end
     @test_throws ErrorException SlurmClusterManager.get_slurm_jobid_int()
   end
 end
+
+@testset "warn_if_unexpected_params()" begin
+  if Base.VERSION >= v"1.6"
+    # This test is not relevant for Julia 1.6+
+  else
+    params = Dict(:env => ["foo" => "bar"])
+    SlurmClusterManager.warn_if_unexpected_params(params)
+    @test_logs(
+      (:warn, "The user provided the `env` kwarg, but SlurmClusterManager.jl's support for the `env` kwarg requires Julia 1.6 or later"),
+      SlurmClusterManager.warn_if_unexpected_params(params),
+    )
+  end
+end