From 3b7a56cb43a26b283c806bf80a62149ec7723c9a Mon Sep 17 00:00:00 2001
From: jbhusson <jonah@elixxir.io>
Date: Fri, 17 Sep 2021 13:27:24 -0400
Subject: [PATCH] Code cleanup part 1

---
 notifications/apns/apns.go                    | 19 ++++++++++++
 {firebase => notifications/firebase}/fcm.go   |  5 +--
 .../firebase}/fcm_test.go                     |  2 +-
 notifications/notifications.go                | 31 +++++++++----------
 notifications/notifications_test.go           | 23 +++++---------
 5 files changed, 45 insertions(+), 35 deletions(-)
 create mode 100644 notifications/apns/apns.go
 rename {firebase => notifications/firebase}/fcm.go (96%)
 rename {firebase => notifications/firebase}/fcm_test.go (98%)

diff --git a/notifications/apns/apns.go b/notifications/apns/apns.go
new file mode 100644
index 0000000..2e3f8e1
--- /dev/null
+++ b/notifications/apns/apns.go
@@ -0,0 +1,19 @@
+package apns
+
+import "github.com/sideshow/apns2"
+
+type ApnsComm struct {
+	*apns2.Client
+	topic string
+}
+
+func NewApnsComm(cl *apns2.Client, topic string) *ApnsComm {
+	return &ApnsComm{
+		Client: cl,
+		topic:  topic,
+	}
+}
+
+func (c *ApnsComm) GetTopic() string {
+	return c.topic
+}
diff --git a/firebase/fcm.go b/notifications/firebase/fcm.go
similarity index 96%
rename from firebase/fcm.go
rename to notifications/firebase/fcm.go
index cbe26d4..94bbfb4 100644
--- a/firebase/fcm.go
+++ b/notifications/firebase/fcm.go
@@ -21,13 +21,13 @@ import (
 )
 
 // function types for use in notificationsbot struct
-type SetupFunc func(string) (*messaging.Client, context.Context, error)
 type SendFunc func(FBSender, string, *mixmessages.NotificationData) (string, error)
 
 // FirebaseComm is a struct which holds the functions to setup the messaging app and sending notifications
 // Using a struct in this manner allows us to properly unit test the NotifyUser function
 type FirebaseComm struct {
 	SendNotification SendFunc
+	*messaging.Client
 }
 
 // FBSender is an interface which matches the send function in the messaging app, allowing us to unit test sendNotification
@@ -36,9 +36,10 @@ type FBSender interface {
 }
 
 // NewFirebaseComm create a *FirebaseComm object with the proper setup and send functions
-func NewFirebaseComm() *FirebaseComm {
+func NewFirebaseComm(cl *messaging.Client) *FirebaseComm {
 	return &FirebaseComm{
 		SendNotification: sendNotification,
+		Client:           cl,
 	}
 }
 
diff --git a/firebase/fcm_test.go b/notifications/firebase/fcm_test.go
similarity index 98%
rename from firebase/fcm_test.go
rename to notifications/firebase/fcm_test.go
index 70741ae..08650bb 100644
--- a/firebase/fcm_test.go
+++ b/notifications/firebase/fcm_test.go
@@ -37,7 +37,7 @@ func TestSendNotification(t *testing.T) {
 
 // Unit test the NewFirebaseComm method
 func TestNewFirebaseComm(t *testing.T) {
-	comm := NewFirebaseComm()
+	comm := NewFirebaseComm(nil)
 	if comm.SendNotification == nil {
 		t.Error("Failed to set functions in comm")
 	}
diff --git a/notifications/notifications.go b/notifications/notifications.go
index aa8e941..e09a16e 100644
--- a/notifications/notifications.go
+++ b/notifications/notifications.go
@@ -10,7 +10,8 @@ package notifications
 
 import (
 	"encoding/base64"
-	"firebase.google.com/go/messaging"
+	"gitlab.com/elixxir/notifications-bot/notifications/apns"
+
 	// "github.com/jonahh-yeti/apns"
 	"github.com/pkg/errors"
 	"github.com/sideshow/apns2"
@@ -22,7 +23,7 @@ import (
 	"gitlab.com/elixxir/comms/notificationBot"
 	"gitlab.com/elixxir/crypto/hash"
 	"gitlab.com/elixxir/crypto/registration"
-	"gitlab.com/elixxir/notifications-bot/firebase"
+	"gitlab.com/elixxir/notifications-bot/notifications/firebase"
 	"gitlab.com/elixxir/notifications-bot/storage"
 	"gitlab.com/xx_network/comms/connect"
 	"gitlab.com/xx_network/crypto/signature/rsa"
@@ -38,7 +39,7 @@ import (
 )
 
 // Function type definitions for the main operations (poll and notify)
-type NotifyFunc func(*pb.NotificationData, ApnsSender, *messaging.Client, *firebase.FirebaseComm, *storage.Storage, string) error
+type NotifyFunc func(*pb.NotificationData, *apns.ApnsComm, *firebase.FirebaseComm, *storage.Storage) error
 type ApnsSender interface {
 	//Send(token string, p apns.Payload, opts ...apns.SendOption) (*apns.Response, error)
 	Push(n *apns2.Notification) (*apns2.Response, error)
@@ -66,10 +67,9 @@ type Impl struct {
 	Storage     *storage.Storage
 	inst        *network.Instance
 	notifyFunc  NotifyFunc
-	fcm         *messaging.Client
-	apnsClient  *apns2.Client
+	fcm         *firebase.FirebaseComm
+	apnsClient  *apns.ApnsComm
 	receivedNdf *uint32
-	params      *Params
 
 	ndfStopper Stopper
 }
@@ -98,20 +98,20 @@ func StartNotifications(params Params, noTLS, noFirebase bool) (*Impl, error) {
 	}
 
 	// Set up firebase messaging client
-	var app *messaging.Client
+	var fbComm *firebase.FirebaseComm
 	if !noFirebase {
-		app, err = firebase.SetupMessagingApp(params.FBCreds)
+		app, err := firebase.SetupMessagingApp(params.FBCreds)
 		if err != nil {
 			return nil, errors.Wrap(err, "Failed to setup firebase messaging app")
 		}
+		fbComm = firebase.NewFirebaseComm(app)
 	}
 	receivedNdf := uint32(0)
 
 	impl := &Impl{
 		notifyFunc:  notifyUser,
-		fcm:         app,
+		fcm:         fbComm,
 		receivedNdf: &receivedNdf,
-		params:      &params,
 	}
 
 	if params.APNS.KeyPath == "" {
@@ -140,7 +140,7 @@ func StartNotifications(params Params, noTLS, noFirebase bool) (*Impl, error) {
 			apnsClient.Production()
 		}
 
-		impl.apnsClient = apnsClient
+		impl.apnsClient = apns.NewApnsComm(apnsClient, params.APNS.BundleID)
 	}
 
 	// Start notification comms server
@@ -178,7 +178,7 @@ func NewImplementation(instance *Impl) *notificationBot.Implementation {
 
 // NotifyUser accepts a UID and service key file path.
 // It handles the logic involved in retrieving a user's token and sending the notification
-func notifyUser(data *pb.NotificationData, apnsClient ApnsSender, fcm *messaging.Client, fc *firebase.FirebaseComm, db *storage.Storage, topic string) error {
+func notifyUser(data *pb.NotificationData, apnsClient *apns.ApnsComm, fc *firebase.FirebaseComm, db *storage.Storage) error {
 	elist, err := db.GetEphemeral(data.EphemeralID)
 	if err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
@@ -209,7 +209,7 @@ func notifyUser(data *pb.NotificationData, apnsClient ApnsSender, fcm *messaging
 				Priority:    apns2.PriorityHigh,
 				Payload:     notifPayload,
 				PushType:    apns2.PushTypeAlert,
-				Topic:       topic,
+				Topic:       apnsClient.GetTopic(),
 			}
 			resp, err := apnsClient.Push(notif)
 			//resp, err := apnsClient.Send(u.Token, apns.Payload{
@@ -239,7 +239,7 @@ func notifyUser(data *pb.NotificationData, apnsClient ApnsSender, fcm *messaging
 				jww.INFO.Printf("Notified ephemeral ID %+v [%+v] and received response %+v", data.EphemeralID, u.Token, resp)
 			}
 		} else {
-			resp, err := fc.SendNotification(fcm, u.Token, data)
+			resp, err := fc.SendNotification(fc.Client, u.Token, data)
 			if err != nil {
 				// Catch two firebase errors that we don't want to crash on
 				// 400 and 404 indicate that the token stored is incorrect
@@ -356,9 +356,8 @@ func (nb *Impl) ReceiveNotificationBatch(notifBatch *pb.NotificationBatch, auth
 
 	jww.INFO.Printf("Received notification batch for round %+v", notifBatch.RoundID)
 
-	fbComm := firebase.NewFirebaseComm()
 	for _, notifData := range notifBatch.GetNotifications() {
-		err := nb.notifyFunc(notifData, nb.apnsClient, nb.fcm, fbComm, nb.Storage, nb.params.APNS.BundleID)
+		err := nb.notifyFunc(notifData, nb.apnsClient, nb.fcm, nb.Storage)
 		if err != nil {
 			return err
 		}
diff --git a/notifications/notifications_test.go b/notifications/notifications_test.go
index d54017a..b2d9de8 100644
--- a/notifications/notifications_test.go
+++ b/notifications/notifications_test.go
@@ -6,15 +6,14 @@
 package notifications
 
 import (
-	"firebase.google.com/go/messaging"
 	"fmt"
-	"github.com/jonahh-yeti/apns"
 	"github.com/pkg/errors"
 	"github.com/sideshow/apns2"
 	pb "gitlab.com/elixxir/comms/mixmessages"
 	"gitlab.com/elixxir/crypto/hash"
 	"gitlab.com/elixxir/crypto/registration"
-	"gitlab.com/elixxir/notifications-bot/firebase"
+	"gitlab.com/elixxir/notifications-bot/notifications/apns"
+	"gitlab.com/elixxir/notifications-bot/notifications/firebase"
 	"gitlab.com/elixxir/notifications-bot/storage"
 	"gitlab.com/xx_network/comms/connect"
 	"gitlab.com/xx_network/crypto/csprng"
@@ -30,16 +29,6 @@ import (
 
 var port = 4200
 
-type MockApns struct{}
-
-func (m *MockApns) Send(token string, p apns.Payload, opts ...apns.SendOption) (*apns.Response, error) {
-	return nil, nil
-}
-
-func (m *MockApns) Push(n *apns2.Notification) (*apns2.Response, error) {
-	return nil, nil
-}
-
 // Test notificationbot's notifyuser function
 // this mocks the setup and send functions, and only tests the core logic of this function
 func TestNotifyUser(t *testing.T) {
@@ -73,11 +62,13 @@ func TestNotifyUser(t *testing.T) {
 	if err != nil {
 		t.Errorf("Failed to add latest ephemeral: %+v", err)
 	}
+
+	ac := apns.NewApnsComm(apns2.NewTokenClient(nil), "")
 	err = notifyUser(&pb.NotificationData{
 		EphemeralID: eph.EphemeralId,
 		IdentityFP:  nil,
 		MessageHash: nil,
-	}, &MockApns{}, nil, fcBadSend, s, "")
+	}, ac, fcBadSend, s)
 	if err == nil {
 		t.Errorf("Should have returned an error")
 	}
@@ -86,7 +77,7 @@ func TestNotifyUser(t *testing.T) {
 		EphemeralID: eph.EphemeralId,
 		IdentityFP:  nil,
 		MessageHash: nil,
-	}, &MockApns{}, nil, fc, s, "")
+	}, ac, fc, s)
 	if err != nil {
 		t.Errorf("Failed to notify user properly")
 	}
@@ -297,7 +288,7 @@ func TestImpl_UnregisterForNotifications(t *testing.T) {
 func TestImpl_ReceiveNotificationBatch(t *testing.T) {
 	impl := getNewImpl()
 	dataChan := make(chan *pb.NotificationData)
-	impl.notifyFunc = func(data *pb.NotificationData, apns ApnsSender, f *messaging.Client, fc *firebase.FirebaseComm, s *storage.Storage, topic string) error {
+	impl.notifyFunc = func(data *pb.NotificationData, apns *apns.ApnsComm, fc *firebase.FirebaseComm, s *storage.Storage) error {
 		go func() { dataChan <- data }()
 		return nil
 	}
-- 
GitLab