From 584cd8aecad35d8862e3172e006c7d651274effe Mon Sep 17 00:00:00 2001
From: John Roesler <johnrroesler@gmail.com>
Date: Wed, 8 Nov 2023 13:53:56 -0600
Subject: [PATCH] elector test & makefile test coverage (#605)

---
 Makefile          |   4 ++
 scheduler_test.go | 104 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+)

diff --git a/Makefile b/Makefile
index 1e16aef..1dcd6fc 100644
--- a/Makefile
+++ b/Makefile
@@ -11,3 +11,7 @@ lint:
 
 test:
 	@go test -race -v $(GO_FLAGS) -count=1 $(GO_PKGS)
+
+test_coverage:
+	@go test -race -v $(GO_FLAGS) -count=1 -coverprofile=coverage.out -covermode=atomic $(GO_PKGS)
+	@go tool cover -html coverage.out
diff --git a/scheduler_test.go b/scheduler_test.go
index 682fd2b..509461e 100644
--- a/scheduler_test.go
+++ b/scheduler_test.go
@@ -2,7 +2,9 @@ package gocron
 
 import (
 	"context"
+	"fmt"
 	"log/slog"
+	"sync"
 	"testing"
 	"time"
 
@@ -876,3 +878,105 @@ func TestScheduler_LimitMode(t *testing.T) {
 		})
 	}
 }
+
+var _ Elector = (*testElector)(nil)
+
+type testElector struct {
+	mu            sync.Mutex
+	leaderElected bool
+}
+
+func (t *testElector) IsLeader(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return fmt.Errorf("done")
+	default:
+	}
+
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if t.leaderElected {
+		return fmt.Errorf("already elected leader")
+	}
+	t.leaderElected = true
+	return nil
+}
+
+func TestScheduler_WithDistributedElector(t *testing.T) {
+	goleak.VerifyNone(t)
+	tests := []struct {
+		name  string
+		count int
+	}{
+		{
+			"3 schedulers",
+			3,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			elector := &testElector{}
+			jobsRan := make(chan struct{}, 20)
+			ctx, cancel := context.WithCancel(context.Background())
+			schedulersDone := make(chan struct{}, tt.count)
+
+			for i := tt.count; i > 0; i-- {
+				go func() {
+					s, err := newTestScheduler(
+						WithDistributedElector(elector),
+					)
+					require.NoError(t, err)
+
+					_, err = s.NewJob(
+						DurationJob(
+							time.Second,
+						),
+						NewTask(
+							func() {
+								jobsRan <- struct{}{}
+							},
+						),
+						WithStartAt(
+							WithStartImmediately(),
+						),
+					)
+					require.NoError(t, err)
+
+					s.Start()
+
+					<-ctx.Done()
+					err = s.Shutdown()
+					require.NoError(t, err)
+					schedulersDone <- struct{}{}
+				}()
+			}
+
+			var runCount int
+			select {
+			case <-jobsRan:
+				cancel()
+				runCount++
+			case <-time.After(time.Second):
+				cancel()
+				t.Error("timed out waiting for job to run")
+			}
+
+			var doneCount int
+			timeout := time.Now().Add(3 * time.Second)
+			for doneCount < tt.count && time.Now().After(timeout) {
+				select {
+				case <-schedulersDone:
+					doneCount++
+				default:
+				}
+			}
+			close(jobsRan)
+			for range jobsRan {
+				runCount++
+			}
+
+			assert.Equal(t, 1, runCount)
+		})
+	}
+}