refactor(fcm): add context parameter to various functions (#775)

- Introduce `graceful.NewManager` initialization in `main.go`
- Modify `pinger` function to accept a context parameter
- Update `PushToAndroid` function to accept a context parameter
- Update `InitFCMClient` function to accept a context parameter
- Replace `context.Background()` with passed context in various functions
- Add context import in `notification_fcm_test.go`

Signed-off-by: appleboy <appleboy.tw@gmail.com>

main.go

@@ -165,9 +165,13 @@ func main() {
 		}
 	}
 
+	g := graceful.NewManager(
+		graceful.WithLogger(logx.QueueLogger()),
+	)
+
 	if ping {
-		if err := pinger(cfg); err != nil {
-			logx.LogError.Warnf("ping server error: %v", err)
+		if err := pinger(g.ShutdownContext(), cfg); err != nil {
+			logx.LogError.Fatal(err)
 		}
 		return
 	}
@@ -195,7 +199,7 @@ func main() {
 			return
 		}
 
-		if _, err := notify.PushToAndroid(req, cfg); err != nil {
+		if _, err := notify.PushToAndroid(g.ShutdownContext(), req, cfg); err != nil {
 			return
 		}
 
@@ -342,10 +346,6 @@ func main() {
 		queue.WithLogger(logx.QueueLogger()),
 	)
 
-	g := graceful.NewManager(
-		graceful.WithLogger(logx.QueueLogger()),
-	)
-
 	g.AddShutdownJob(func() error {
 		// logx.LogAccess.Info("close the queue system, current queue usage: ", q.Usage())
 		// stop queue system and wait job completed
@@ -365,7 +365,7 @@ func main() {
 	}
 
 	if cfg.Android.Enabled {
-		if _, err = notify.InitFCMClient(cfg); err != nil {
+		if _, err = notify.InitFCMClient(g.ShutdownContext(), cfg); err != nil {
 			logx.LogError.Fatal(err)
 		}
 	}
@@ -440,7 +440,7 @@ func usage() {
 
 // handles pinging the endpoint and returns an error if the
 // agent is in an unhealthy state.
-func pinger(cfg *config.ConfYaml) error {
+func pinger(ctx context.Context, cfg *config.ConfYaml) error {
 	transport := &http.Transport{
 		Dial: (&net.Dialer{
 			Timeout: 5 * time.Second,
@@ -452,7 +452,7 @@ func pinger(cfg *config.ConfYaml) error {
 		Transport: transport,
 	}
 	req, _ := http.NewRequestWithContext(
-		context.Background(),
+		ctx,
 		http.MethodGet,
 		"http://localhost:"+cfg.Core.Port+cfg.API.HealthURI,
 		nil,

notify/notification.go

@@ -241,7 +241,7 @@ func SendNotification(
 	case core.PlatFormIos:
 		resp, err = PushToIOS(v, cfg)
 	case core.PlatFormAndroid:
-		resp, err = PushToAndroid(v, cfg)
+		resp, err = PushToAndroid(ctx, v, cfg)
 	case core.PlatFormHuawei:
 		resp, err = PushToHuawei(v, cfg)
 	}

notify/notification_fcm.go

@@ -24,7 +24,7 @@ func fileExists(filename string) bool {
 }
 
 // InitFCMClient use for initialize FCM Client.
-func InitFCMClient(cfg *config.ConfYaml) (*fcm.Client, error) {
+func InitFCMClient(ctx context.Context, cfg *config.ConfYaml) (*fcm.Client, error) {
 	var opts []fcm.Option
 
 	credential := os.Getenv("GOOGLE_APPLICATION_CREDENTIALS")
@@ -47,7 +47,7 @@ func InitFCMClient(cfg *config.ConfYaml) (*fcm.Client, error) {
 	}
 
 	return fcm.NewClient(
-		context.Background(),
+		ctx,
 		opts...,
 	)
 }
@@ -119,7 +119,7 @@ func GetAndroidNotification(req *PushNotification) []*messaging.Message {
 }
 
 // PushToAndroid provide send notification to Android server.
-func PushToAndroid(req *PushNotification, cfg *config.ConfYaml) (resp *ResponsePush, err error) {
+func PushToAndroid(ctx context.Context, req *PushNotification, cfg *config.ConfYaml) (resp *ResponsePush, err error) {
 	logx.LogAccess.Debug("Start push notification for Android")
 
 	var (
@@ -140,7 +140,7 @@ func PushToAndroid(req *PushNotification, cfg *config.ConfYaml) (resp *ResponseP
 	}
 
 	resp = &ResponsePush{}
-	client, err = InitFCMClient(cfg)
+	client, err = InitFCMClient(ctx, cfg)
 
 Retry:
 	notification := GetAndroidNotification(req)
@@ -150,7 +150,7 @@ Retry:
 		return resp, err
 	}
 
-	res, err := client.Send(context.Background(), notification...)
+	res, err := client.Send(ctx, notification...)
 	if err != nil {
 		// Send Message error
 		logx.LogError.Error("FCM server send message error: " + err.Error())

notify/notification_fcm_test.go

@@ -1,6 +1,7 @@
 package notify
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -27,7 +28,7 @@ func TestMissingKeyForInitFCMClient(t *testing.T) {
 	cfg, _ := config.LoadConf()
 	cfg.Android.Credential = ""
 	cfg.Android.KeyPath = ""
-	client, err := InitFCMClient(cfg)
+	client, err := InitFCMClient(context.Background(), cfg)
 
 	assert.Nil(t, client)
 	assert.Error(t, err)
@@ -47,7 +48,7 @@ func TestPushToAndroidWrongToken(t *testing.T) {
 	}
 
 	// Android Success count: 0, Failure count: 2
-	resp, err := PushToAndroid(req, cfg)
+	resp, err := PushToAndroid(context.Background(), req, cfg)
 	assert.Nil(t, err)
 	assert.Len(t, resp.Logs, 2)
 }
@@ -68,7 +69,7 @@ func TestPushToAndroidRightTokenForJSONLog(t *testing.T) {
 		Message:  "Welcome",
 	}
 
-	resp, err := PushToAndroid(req, cfg)
+	resp, err := PushToAndroid(context.Background(), req, cfg)
 	assert.Nil(t, err)
 	assert.Len(t, resp.Logs, 0)
 }
@@ -87,7 +88,7 @@ func TestPushToAndroidRightTokenForStringLog(t *testing.T) {
 		Message:  "Welcome",
 	}
 
-	resp, err := PushToAndroid(req, cfg)
+	resp, err := PushToAndroid(context.Background(), req, cfg)
 	assert.Nil(t, err)
 	assert.Len(t, resp.Logs, 0)
 }