From 30921b708f1ef7910842cc63fc27fe383c7e5d81 Mon Sep 17 00:00:00 2001
From: "Son L. Phan" <56083989+drwpls@users.noreply.github.com>
Date: Tue, 12 Dec 2023 22:53:46 +0700
Subject: [PATCH] Ability to register for events associated with jobs starting
 immediately (#632)

* Ability to register for events associated with jobs starting immediately

This commit add some functions to support register events for job which
is on creating and run immediately.

Because of there is RegisterEventListeners in Scheduler receiver
interface (which will register event for all exsiting jobs), I don't
change the function.

* add TestScheduler Register Event Tests

* lint and fix typo

---------

Co-authored-by: Your Name <you@example.com>
---
 scheduler.go      | 40 +++++++++++++++++++++++++++++++++
 scheduler_test.go | 57 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/scheduler.go b/scheduler.go
index e681fbe6..69fc86c6 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -1527,6 +1527,46 @@ func (s *Scheduler) RegisterEventListeners(eventListeners ...EventListener) {
 	}
 }
 
+// BeforeJobRuns registers an event listener that is called before a job runs.
+func (s *Scheduler) BeforeJobRuns(eventListenerFunc func(jobName string)) *Scheduler {
+	job := s.getCurrentJob()
+	job.mu.Lock()
+	defer job.mu.Unlock()
+	job.eventListeners.beforeJobRuns = eventListenerFunc
+
+	return s
+}
+
+// AfterJobRuns registers an event listener that is called after a job runs.
+func (s *Scheduler) AfterJobRuns(eventListenerFunc func(jobName string)) *Scheduler {
+	job := s.getCurrentJob()
+	job.mu.Lock()
+	defer job.mu.Unlock()
+	job.eventListeners.afterJobRuns = eventListenerFunc
+
+	return s
+}
+
+// WhenJobStarts registers an event listener that is called when a job starts.
+func (s *Scheduler) WhenJobReturnsError(eventListenerFunc func(jobName string, err error)) *Scheduler {
+	job := s.getCurrentJob()
+	job.mu.Lock()
+	defer job.mu.Unlock()
+	job.eventListeners.onError = eventListenerFunc
+
+	return s
+}
+
+// WhenJobStarts registers an event listener that is called when a job starts.
+func (s *Scheduler) WhenJobReturnsNoError(eventListenerFunc func(jobName string)) *Scheduler {
+	job := s.getCurrentJob()
+	job.mu.Lock()
+	defer job.mu.Unlock()
+	job.eventListeners.noError = eventListenerFunc
+
+	return s
+}
+
 func (s *Scheduler) PauseJobExecution(shouldPause bool) {
 	s.executor.skipExecution.Store(shouldPause)
 }
diff --git a/scheduler_test.go b/scheduler_test.go
index 14485a56..9933cd63 100644
--- a/scheduler_test.go
+++ b/scheduler_test.go
@@ -2659,6 +2659,63 @@ func TestScheduler_ChainOrder(t *testing.T) {
 	require.Len(t, s.jobs, 2)
 }
 
+func TestScheduler_Register_Event(t *testing.T) {
+	userDefinedError := errors.New("user defined error")
+	testCases := []struct {
+		description   string
+		jobFunc       func() error
+		expected      []uint8
+		expectedError error
+	}{
+		{"event order: no error", func() error { return nil }, []uint8{1, 2, 3, 4}, nil},
+		{"event order: on error", func() error { return userDefinedError }, []uint8{1, 2, 3, 4}, userDefinedError},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			s := NewScheduler(time.UTC)
+			order := make(chan uint8, 10)
+			done := make(chan struct{})
+			var expectedErr error
+			s = s.BeforeJobRuns(func(jobName string) {
+				order <- 1
+			})
+			s = s.WhenJobReturnsError(func(jobName string, err error) {
+				order <- 3
+				expectedErr = err
+			})
+			s = s.WhenJobReturnsNoError(func(jobName string) {
+				order <- 3
+			})
+			s = s.AfterJobRuns(func(jobName string) {
+				order <- 4
+				done <- struct{}{}
+			})
+
+			_, err := s.Day().Every(1).StartImmediately().Do(func() error {
+				order <- 2
+				return tc.jobFunc()
+			})
+
+			require.NoError(t, err)
+			s.StartAsync()
+			select {
+			case <-done:
+				assert.Equal(t, tc.expectedError, expectedErr)
+				assert.Equal(t, len(tc.expected), len(order))
+				if len(tc.expected) == len(order) {
+					for i := 0; i < len(tc.expected); i++ {
+						assert.Equal(t, tc.expected[i], <-order)
+					}
+				}
+				s.Clear()
+			case <-time.After(1 * time.Second):
+				t.Fatal("timeout")
+			}
+		})
+	}
+}
+
 var _ Locker = (*locker)(nil)
 
 type locker struct {