From 141fb896cebd5f48122a6277cb1397f6e637234f Mon Sep 17 00:00:00 2001
From: gaffo <mr.gaffo@gmail.com>
Date: Wed, 17 Jul 2024 22:51:48 -0700
Subject: [PATCH] Run::AddJobWithState

---
 run.go      | 10 +++++++---
 run_test.go | 15 +++++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/run.go b/run.go
index ce7fb5c..2af13e8 100644
--- a/run.go
+++ b/run.go
@@ -69,8 +69,7 @@ func (r *Run[OC, JC]) Init() {
 	}
 }
 
-// Add a job to the pool, this shouldn't be called once it's running
-func (r *Run[OC, JC]) AddJob(jc JC) {
+func (r *Run[OC, JC]) AddJobWithState(jc JC, state string) {
 	r.m.Lock()
 	defer r.m.Unlock()
 
@@ -79,7 +78,7 @@ func (r *Run[OC, JC]) AddJob(jc JC) {
 	j := Job[JC]{
 		Id:          id,
 		C:           jc,
-		State:       TRIGGER_STATE_NEW,
+		State:       state,
 		StateErrors: map[string][]string{},
 	}
 	j.UpdateLastEvent()
@@ -90,6 +89,11 @@ func (r *Run[OC, JC]) AddJob(jc JC) {
 	r.Jobs[id] = j
 }
 
+// Add a job to the pool, this shouldn't be called once it's running
+func (r *Run[OC, JC]) AddJob(jc JC) {
+	r.AddJobWithState(jc, TRIGGER_STATE_NEW)
+}
+
 func (r *Run[OC, JC]) NextJobForState(state string) (Job[JC], bool) {
 	r.m.Lock()
 	defer r.m.Unlock()
diff --git a/run_test.go b/run_test.go
index e36b637..f1df6b7 100644
--- a/run_test.go
+++ b/run_test.go
@@ -152,3 +152,18 @@ func TestRun_JobsInFlight2(t *testing.T) {
 	r.Return(j)
 	require.False(t, r.JobsInFlight())
 }
+
+func Test_AddJobWithState(t *testing.T) {
+	t.Parallel()
+	r := NewRun[MyOverallContext, MyJobContext]("job", MyOverallContext{})
+	r.AddJobWithState(MyJobContext{Count: 0}, "other_state")
+	assert.Equal(t, map[string]StatusCount{
+		"other_state": {
+			State:     "other_state",
+			Count:     1,
+			Executing: 0,
+		},
+	}, r.StatusCounts())
+	assert.Equal(t, 1, len(r.Jobs))
+	assert.Equal(t, "other_state", r.Jobs["0"].State)
+}