From bbae93813b57a04b454b80fb9ae3d39866ad5eb5 Mon Sep 17 00:00:00 2001
From: Ken Hibino <ken.hibino7@gmail.com>
Date: Sat, 22 Feb 2020 14:30:24 -0800
Subject: [PATCH] Add test for cancelation pubsub

---
 internal/asynqtest/asynqtest.go |  7 ++++++
 internal/rdb/rdb_test.go        | 42 +++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/internal/asynqtest/asynqtest.go b/internal/asynqtest/asynqtest.go
index 531605d..7b8a0ae 100644
--- a/internal/asynqtest/asynqtest.go
+++ b/internal/asynqtest/asynqtest.go
@@ -53,6 +53,13 @@ var SortProcessInfoOpt = cmp.Transformer("SortProcessInfo", func(in []*base.Proc
 	return out
 })
 
+// SortStringSliceOpt is a cmp.Option to sort string slice.
+var SortStringSliceOpt = cmp.Transformer("SortStringSlice", func(in []string) []string {
+	out := append([]string(nil), in...)
+	sort.Strings(out)
+	return out
+})
+
 // IgnoreIDOpt is an cmp.Option to ignore ID field in task messages when comparing.
 var IgnoreIDOpt = cmpopts.IgnoreFields(base.TaskMessage{}, "ID")
 
diff --git a/internal/rdb/rdb_test.go b/internal/rdb/rdb_test.go
index e0b4669..b44cc0f 100644
--- a/internal/rdb/rdb_test.go
+++ b/internal/rdb/rdb_test.go
@@ -7,6 +7,7 @@ package rdb
 import (
 	"encoding/json"
 	"fmt"
+	"sync"
 	"testing"
 	"time"
 
@@ -975,3 +976,44 @@ func TestClearProcessState(t *testing.T) {
 		t.Errorf("%q contained %v, want %v", base.AllWorkers, gotWorkerKeys, wantWorkerKeys)
 	}
 }
+
+func TestCancelationPubSub(t *testing.T) {
+	r := setup(t)
+
+	pubsub, err := r.CancelationPubSub()
+	if err != nil {
+		t.Fatalf("(*RDB).CancelationPubSub() returned an error: %v", err)
+	}
+
+	cancelCh := pubsub.Channel()
+
+	var (
+		mu       sync.Mutex
+		received []string
+	)
+
+	go func() {
+		for msg := range cancelCh {
+			mu.Lock()
+			received = append(received, msg.Payload)
+			mu.Unlock()
+		}
+	}()
+
+	publish := []string{"one", "two", "three"}
+
+	for _, msg := range publish {
+		r.PublishCancelation(msg)
+	}
+
+	// allow for message to reach subscribers.
+	time.Sleep(time.Second)
+
+	pubsub.Close()
+
+	mu.Lock()
+	if diff := cmp.Diff(publish, received, h.SortStringSliceOpt); diff != "" {
+		t.Errorf("subscriber received %v, want %v; (-want,+got)\n%s", received, publish, diff)
+	}
+	mu.Unlock()
+}