From b25baa252411d09d887fd6341f5799f6774415d5 Mon Sep 17 00:00:00 2001
From: James Houlahan <james.houlahan@protonmail.com>
Date: Fri, 3 Jul 2020 08:22:28 +0200
Subject: [PATCH] test: set sent label properly

---
 internal/store/event_loop.go |  9 +--------
 pkg/pmapi/messages.go        | 10 ++++++++++
 test/store_setup_test.go     |  6 ++++++
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/internal/store/event_loop.go b/internal/store/event_loop.go
index f5e6d356..5f324734 100644
--- a/internal/store/event_loop.go
+++ b/internal/store/event_loop.go
@@ -523,14 +523,7 @@ func updateMessage(msgLog *logrus.Entry, message *pmapi.Message, updates *pmapi.
 		message.LabelIDs = updates.LabelIDs
 	} else {
 		for _, added := range updates.LabelIDsAdded {
-			hasLabel := false
-			for _, l := range message.LabelIDs {
-				if added == l {
-					hasLabel = true
-					break
-				}
-			}
-			if !hasLabel {
+			if !message.HasLabelID(added) {
 				msgLog.WithField("added", added).Trace("Adding label to message")
 				message.LabelIDs = append(message.LabelIDs, added)
 			}
diff --git a/pkg/pmapi/messages.go b/pkg/pmapi/messages.go
index adff6099..5c93cd28 100644
--- a/pkg/pmapi/messages.go
+++ b/pkg/pmapi/messages.go
@@ -241,6 +241,16 @@ func (m *Message) IsDraft() bool {
 	return (m.Flags & (FlagReceived | FlagSent)) == 0
 }
 
+// HasLabelID returns whether the message has the `labelID`.
+func (m *Message) HasLabelID(labelID string) bool {
+	for _, l := range m.LabelIDs {
+		if l == labelID {
+			return true
+		}
+	}
+	return false
+}
+
 func (m *Message) IsBodyEncrypted() bool {
 	trimmedBody := strings.TrimSpace(m.Body)
 	return strings.HasPrefix(trimmedBody, MessageHeader) &&
diff --git a/test/store_setup_test.go b/test/store_setup_test.go
index 04578c14..9db508e0 100644
--- a/test/store_setup_test.go
+++ b/test/store_setup_test.go
@@ -81,11 +81,17 @@ func thereAreMessagesInMailboxesForAddressOfUser(mailboxNames, bddAddressID, bdd
 		if err != nil {
 			return internalError(err, "getting labels %s for %s", mailboxNames, account.Username())
 		}
+
 		message := &pmapi.Message{
 			MIMEType:  "text/plain",
 			LabelIDs:  labelIDs,
 			AddressID: account.AddressID(),
 		}
+
+		if message.HasLabelID(pmapi.SentLabel) {
+			message.Flags |= pmapi.FlagSent
+		}
+
 		for n, cell := range messages.Rows[i].Cells {
 			switch head[n].Value {
 			case "from":