From 6c93f1f1ecbaa11ac040ae04acf3d606a2b9d903 Mon Sep 17 00:00:00 2001
From: Michal Horejsek <horejsek@protonmail.com>
Date: Mon, 17 Aug 2020 09:10:03 +0200
Subject: [PATCH] Fix integration tests - compiting message flags

---
 pkg/message/flags.go     |  1 +
 pkg/pmapi/messages.go    | 19 +++++++++++++++++++
 test/liveapi/messages.go |  4 ++++
 3 files changed, 24 insertions(+)

diff --git a/pkg/message/flags.go b/pkg/message/flags.go
index 798174ed..c0b0ad38 100644
--- a/pkg/message/flags.go
+++ b/pkg/message/flags.go
@@ -60,6 +60,7 @@ func GetFlags(m *pmapi.Message) (flags []string) {
 }
 
 func ParseFlags(m *pmapi.Message, flags []string) {
+	// Consider to use ComputeMessageFlagsByLabels to keep logic in one place.
 	if (m.Flags & pmapi.FlagSent) == 0 {
 		m.Flags |= pmapi.FlagReceived
 	}
diff --git a/pkg/pmapi/messages.go b/pkg/pmapi/messages.go
index 5c93cd28..73a9539c 100644
--- a/pkg/pmapi/messages.go
+++ b/pkg/pmapi/messages.go
@@ -839,3 +839,22 @@ func (c *client) EmptyFolder(labelID, addressID string) (err error) {
 	err = res.Err()
 	return
 }
+
+// ComputeMessageFlagsByLabels returns flags based on labels.
+func ComputeMessageFlagsByLabels(labels []string) (flag int64) {
+	for _, labelID := range labels {
+		switch labelID {
+		case SentLabel:
+			flag = (flag | FlagSent)
+		case ArchiveLabel, InboxLabel:
+			flag = (flag | FlagReceived)
+		}
+	}
+
+	// NOTE: if the labels are custom only
+	if flag == 0 {
+		flag = FlagReceived
+	}
+
+	return flag
+}
diff --git a/test/liveapi/messages.go b/test/liveapi/messages.go
index f7232095..f064cffd 100644
--- a/test/liveapi/messages.go
+++ b/test/liveapi/messages.go
@@ -36,6 +36,10 @@ func (ctl *Controller) AddUserMessage(username string, message *pmapi.Message) e
 		return fmt.Errorf("user %s does not exist", username)
 	}
 
+	if message.Flags == 0 {
+		message.Flags = pmapi.ComputeMessageFlagsByLabels(message.LabelIDs)
+	}
+
 	body, err := buildMessage(client, message)
 	if err != nil {
 		return errors.Wrap(err, "failed to build message")