From 4b91ed2f70c590ea39c098580c0fbae1d992e7bb Mon Sep 17 00:00:00 2001
From: mikestefanello <mike.stefanello@gmail.com>
Date: Tue, 21 Dec 2021 21:42:53 -0500
Subject: [PATCH] Added test coverage for msg package.

---
 msg/msg.go      | 90 +++++++++++++++++++++++++++++--------------------
 msg/msg_test.go | 48 ++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 36 deletions(-)
 create mode 100644 msg/msg_test.go

diff --git a/msg/msg.go b/msg/msg.go
index 0576aef..65626c2 100644
--- a/msg/msg.go
+++ b/msg/msg.go
@@ -15,52 +15,70 @@ const (
 	TypeDanger  Type = "danger"
 )
 
-// TODO: Error handling and cleanup
+const (
+	// sessionName stores the name of the session which contains flash messages
+	sessionName = "msg"
+)
 
-func Success(c echo.Context, message string) {
-	Set(c, TypeSuccess, message)
+// Success sets a success flash message
+func Success(ctx echo.Context, message string) {
+	Set(ctx, TypeSuccess, message)
 }
 
-func Info(c echo.Context, message string) {
-	Set(c, TypeInfo, message)
+// Info sets an info flash message
+func Info(ctx echo.Context, message string) {
+	Set(ctx, TypeInfo, message)
 }
 
-func Warning(c echo.Context, message string) {
-	Set(c, TypeWarning, message)
+// Warning sets a warning flash message
+func Warning(ctx echo.Context, message string) {
+	Set(ctx, TypeWarning, message)
 }
 
-func Danger(c echo.Context, message string) {
-	Set(c, TypeDanger, message)
+// Danger sets a danger flash message
+func Danger(ctx echo.Context, message string) {
+	Set(ctx, TypeDanger, message)
 }
 
-// Set adds a new message into the cookie storage.
-func Set(c echo.Context, typ Type, message string) {
-	sess := getSession(c)
-	sess.AddFlash(message, string(typ))
-	_ = sess.Save(c.Request(), c.Response())
-}
-
-// Get gets flash messages from the cookie storage.
-func Get(c echo.Context, typ Type) []string {
-	sess := getSession(c)
-
-	fm := sess.Flashes(string(typ))
-	// If we have some messages.
-	if len(fm) > 0 {
-		_ = sess.Save(c.Request(), c.Response())
-		// Initiate a strings slice to return messages.
-		var flashes []string
-		for _, fl := range fm {
-			// Add message to the slice.
-			flashes = append(flashes, fl.(string))
-		}
-
-		return flashes
+// Set adds a new flash message of a given type into the session storage
+// Errors will logged and not returned
+func Set(ctx echo.Context, typ Type, message string) {
+	if sess, err := getSession(ctx); err == nil {
+		sess.AddFlash(message, string(typ))
+		save(ctx, sess)
 	}
-	return nil
 }
 
-func getSession(c echo.Context) *sessions.Session {
-	sess, _ := session.Get("msg", c)
-	return sess
+// Get gets flash messages of a given type from the session storage
+// Errors will logged and not returned
+func Get(ctx echo.Context, typ Type) []string {
+	var msgs []string
+
+	if sess, err := getSession(ctx); err == nil {
+		if flash := sess.Flashes(string(typ)); len(flash) > 0 {
+			save(ctx, sess)
+
+			for _, m := range flash {
+				msgs = append(msgs, m.(string))
+			}
+		}
+	}
+
+	return msgs
+}
+
+// getSession gets the flash message session
+func getSession(ctx echo.Context) (*sessions.Session, error) {
+	sess, err := session.Get(sessionName, ctx)
+	if err != nil {
+		ctx.Logger().Errorf("cannot load flash message session: %v", err)
+	}
+	return sess, err
+}
+
+// save saves the flash message session
+func save(ctx echo.Context, sess *sessions.Session) {
+	if err := sess.Save(ctx.Request(), ctx.Response()); err != nil {
+		ctx.Logger().Errorf("failed to set flash message: %v", err)
+	}
 }
diff --git a/msg/msg_test.go b/msg/msg_test.go
new file mode 100644
index 0000000..5676708
--- /dev/null
+++ b/msg/msg_test.go
@@ -0,0 +1,48 @@
+package msg
+
+import (
+	"testing"
+
+	"goweb/tests"
+
+	"github.com/go-playground/assert/v2"
+	"github.com/stretchr/testify/require"
+
+	"k8s.io/apimachinery/pkg/util/rand"
+
+	"github.com/labstack/echo/v4"
+)
+
+func TestMsg(t *testing.T) {
+	e := echo.New()
+	ctx, _ := tests.NewContext(e, "/")
+	tests.InitSession(ctx)
+
+	assertMsg := func(typ Type, message string) {
+		ret := Get(ctx, typ)
+		require.Len(t, ret, 1)
+		assert.Equal(t, message, ret[0])
+		ret = Get(ctx, typ)
+		require.Len(t, ret, 0)
+	}
+
+	text := rand.String(10)
+	Success(ctx, text)
+	assertMsg(TypeSuccess, text)
+
+	text = rand.String(10)
+	Info(ctx, text)
+	assertMsg(TypeInfo, text)
+
+	text = rand.String(10)
+	Danger(ctx, text)
+	assertMsg(TypeDanger, text)
+
+	text = rand.String(10)
+	Warning(ctx, text)
+	assertMsg(TypeWarning, text)
+
+	text = rand.String(10)
+	Set(ctx, TypeSuccess, text)
+	assertMsg(TypeSuccess, text)
+}