From 6d62034e6a694844ba664e01360ec95ea44bebc8 Mon Sep 17 00:00:00 2001
From: Arjun Mahishi <arjun.mahishi@gmail.com>
Date: Sun, 4 Oct 2020 12:50:13 +0530
Subject: [PATCH] Add a new Job method called 'LimitRunsTo' to control how many
 times a job should run

---
 example_test.go |  7 +++++++
 job.go          | 24 ++++++++++++++++++++++
 job_test.go     | 54 +++++++++++++++++++++++++++++++++++++++++++++++++
 scheduler.go    |  2 +-
 4 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/example_test.go b/example_test.go
index 120177d7..5d2e7c9e 100644
--- a/example_test.go
+++ b/example_test.go
@@ -88,3 +88,10 @@ func ExampleScheduler_Clear() {
 	// 3
 	// 0
 }
+
+func ExampleJob_LimitRunsTo() {
+	s := gocron.NewScheduler(time.UTC)
+	job, _ := s.Every(1).Second().Do(task)
+	job.LimitRunsTo(2)
+	s.StartAsync()
+}
diff --git a/job.go b/job.go
index e4734e8a..65f5bddb 100644
--- a/job.go
+++ b/job.go
@@ -21,6 +21,13 @@ type Job struct {
 	fparams           map[string][]interface{} // Map for function and  params of function
 	lock              bool                     // lock the Job from running at same time form multiple instances
 	tags              []string                 // allow the user to tag Jobs with certain labels
+	runConfig         runConfig                // configuration for how many times to run the job
+	runCount          int                      // number of time the job ran
+}
+
+type runConfig struct {
+	finiteRuns bool
+	maxRuns    int
 }
 
 // NewJob creates a new Job with the provided interval
@@ -38,6 +45,7 @@ func NewJob(interval uint64) *Job {
 // Run the Job and immediately reschedule it
 func (j *Job) run() {
 	callJobFuncWithParams(j.funcs[j.jobFunc], j.fparams[j.jobFunc])
+	j.runCount++
 }
 
 func (j Job) neverRan() bool {
@@ -93,3 +101,19 @@ func (j *Job) Weekday() (time.Weekday, error) {
 	}
 	return *j.scheduledWeekday, nil
 }
+
+// LimitRunsTo limits the number of executions of this
+// job to n. However, the job will still remain in the
+// scheduler
+func (j *Job) LimitRunsTo(n int) {
+	j.runConfig = runConfig{
+		finiteRuns: true,
+		maxRuns:    n,
+	}
+}
+
+// shouldRun eveluates if this job should run again
+// based on the runConfig
+func (j *Job) shouldRun() bool {
+	return !j.runConfig.finiteRuns || j.runCount < j.runConfig.maxRuns
+}
diff --git a/job_test.go b/job_test.go
index da2faed4..470308c0 100644
--- a/job_test.go
+++ b/job_test.go
@@ -53,3 +53,57 @@ func TestGetWeekday(t *testing.T) {
 		})
 	}
 }
+
+func TestJob_shouldRunAgain(t *testing.T) {
+	tests := []struct {
+		name      string
+		runConfig runConfig
+		runCount  int
+		want      bool
+	}{
+		{
+			name:      "should run again (infinite)",
+			runConfig: runConfig{finiteRuns: false},
+			want:      true,
+		},
+		{
+			name:      "should run again (finite)",
+			runConfig: runConfig{finiteRuns: true, maxRuns: 2},
+			runCount:  1,
+			want:      true,
+		},
+		{
+			name:      "shouldn't run again #1",
+			runConfig: runConfig{finiteRuns: true, maxRuns: 2},
+			runCount:  2,
+			want:      false,
+		},
+		{
+			name:      "shouldn't run again #2",
+			runConfig: runConfig{finiteRuns: true, maxRuns: 2},
+			runCount:  4,
+			want:      false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			j := &Job{
+				runConfig: tt.runConfig,
+				runCount:  tt.runCount,
+			}
+			if got := j.shouldRun(); got != tt.want {
+				t.Errorf("Job.shouldRunAgain() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestJob_LimitRunsTo(t *testing.T) {
+	j, _ := NewScheduler(time.Local).Every(1).Second().Do(func() {})
+	j.LimitRunsTo(2)
+	assert.Equal(t, j.shouldRun(), true, "Expecting it to run again")
+	j.run()
+	assert.Equal(t, j.shouldRun(), true, "Expecting it to run again")
+	j.run()
+	assert.Equal(t, j.shouldRun(), false, "Not expecting it to run again")
+}
diff --git a/scheduler.go b/scheduler.go
index e4fd7696..f1e80b2f 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -435,7 +435,7 @@ func (s *Scheduler) StartImmediately() *Scheduler {
 
 // shouldRun returns true if the Job should be run now
 func (s *Scheduler) shouldRun(j *Job) bool {
-	return s.time.Now(s.loc).Unix() >= j.nextRun.Unix()
+	return j.shouldRun() && s.time.Now(s.loc).Unix() >= j.nextRun.Unix()
 }
 
 // setUnit sets the unit type