From 3314acc78ec7a8bebc37f3129924b5e62dd56488 Mon Sep 17 00:00:00 2001
From: Abin Simon <abinsimon10@gmail.com>
Date: Tue, 5 Dec 2023 13:37:45 +0530
Subject: [PATCH] Fix sanity tests for exchange emails export (#4756)

<!-- PR description-->

---

#### Does this PR need a docs update or release note?

- [ ] :white_check_mark: Yes, it's included
- [ ] :clock1: Yes, but in a later PR
- [x] :no_entry: No

#### Type of change

<!--- Please check the type of change your PR introduces: --->
- [ ] :sunflower: Feature
- [ ] :bug: Bugfix
- [ ] :world_map: Documentation
- [x] :robot: Supportability/Tests
- [ ] :computer: CI/Deployment
- [ ] :broom: Tech Debt/Cleanup

#### Issue(s)

<!-- Can reference multiple issues. Use one of the following "magic words" - "closes, fixes" to auto-close the Github issue. -->
* #<issue>

#### Test Plan

<!-- How will this be tested prior to merging.-->
- [ ] :muscle: Manual
- [ ] :zap: Unit test
- [ ] :green_heart: E2E
---
 src/cmd/sanity_test/common/sanitree.go  |  2 +-
 src/cmd/sanity_test/export/exchange.go  |  4 +---
 src/cmd/sanity_test/restore/exchange.go | 25 ++++++++++++++++++++++---
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/cmd/sanity_test/common/sanitree.go b/src/cmd/sanity_test/common/sanitree.go
index 290b5f037..79c23cc36 100644
--- a/src/cmd/sanity_test/common/sanitree.go
+++ b/src/cmd/sanity_test/common/sanitree.go
@@ -149,7 +149,7 @@ type NodeComparator[ET, EL, RT, RL any] func(
 // different data types.  The two trees are expected to represent
 // a common hierarchy.
 //
-// Additional comparisons besides the tre hierarchy are optionally
+// Additional comparisons besides the tree hierarchy are optionally
 // left to the caller by population of the NodeComparator func.
 func CompareDiffTrees[ET, EL, RT, RL any](
 	ctx context.Context,
diff --git a/src/cmd/sanity_test/export/exchange.go b/src/cmd/sanity_test/export/exchange.go
index 52700689e..452de5169 100644
--- a/src/cmd/sanity_test/export/exchange.go
+++ b/src/cmd/sanity_test/export/exchange.go
@@ -36,12 +36,10 @@ func CheckEmailExport(
 		expect *common.Sanitree[models.MailFolderable, any],
 		result *common.Sanitree[fs.FileInfo, fs.FileInfo],
 	) {
-		modifiedExpectedLeaves := map[string]*common.Sanileaf[models.MailFolderable, any]{}
 		modifiedResultLeaves := map[string]*common.Sanileaf[fs.FileInfo, fs.FileInfo]{}
 
-		for key, val := range expect.Leaves {
+		for _, val := range expect.Leaves {
 			val.Size = 0 // we cannot match up sizes
-			modifiedExpectedLeaves[key] = val
 		}
 
 		for key, val := range result.Leaves {
diff --git a/src/cmd/sanity_test/restore/exchange.go b/src/cmd/sanity_test/restore/exchange.go
index 5b7e3d63f..ba1b9ae7f 100644
--- a/src/cmd/sanity_test/restore/exchange.go
+++ b/src/cmd/sanity_test/restore/exchange.go
@@ -29,11 +29,12 @@ func CheckEmailRestoration(
 		"source_container_id", sourceTree.ID,
 		"source_container_name", sourceTree.Name)
 
-	common.AssertEqualTrees[models.MailFolderable, any](
+	// NOTE: We cannot compare leaves as the IDs of the restored items
+	// differ from the original ones.
+	common.CompareDiffTrees[models.MailFolderable, any](
 		ctx,
 		sourceTree,
 		restoredTree.Children[envs.SourceContainer],
-		nil,
 		nil)
 
 	common.Infof(ctx, "Success")
@@ -122,7 +123,7 @@ func recursivelyBuildTree(
 	}
 
 	for _, child := range childFolders {
-		if int(ptr.Val(child.GetTotalItemCount())) == 0 {
+		if int(ptr.Val(child.GetTotalItemCount()))+len(childFolders) == 0 {
 			common.Infof(ctx, "skipped empty folder: %s/%s", location, ptr.Val(child.GetDisplayName()))
 			continue
 		}
@@ -133,9 +134,27 @@ func recursivelyBuildTree(
 			ID:          ptr.Val(child.GetId()),
 			Name:        ptr.Val(child.GetDisplayName()),
 			CountLeaves: int(ptr.Val(child.GetTotalItemCount())),
+			Leaves:      map[string]*common.Sanileaf[models.MailFolderable, any]{},
 			Children:    map[string]*common.Sanitree[models.MailFolderable, any]{},
 		}
 
+		mails, err := ac.Mail().GetItemsInContainer(ctx, userID, c.ID)
+		if err != nil {
+			common.Fatal(ctx, "getting child containers", err)
+		}
+
+		for _, mail := range mails {
+			m := &common.Sanileaf[models.MailFolderable, any]{
+				Parent: c,
+				Self:   mail,
+				ID:     ptr.Val(mail.GetId()),
+				Name:   ptr.Val(mail.GetSubject()),
+				Size:   int64(len(ptr.Val(mail.GetBody().GetContent()))),
+			}
+
+			c.Leaves[m.ID] = m
+		}
+
 		stree.Children[c.Name] = c
 
 		recursivelyBuildTree(ctx, ac, c, userID, location+c.Name+"/")