From 555a81f05ca941e28538d27bb8b2f1d65efcf361 Mon Sep 17 00:00:00 2001
From: Alex Shtin <alex@shtin.com>
Date: Wed, 31 May 2023 17:57:54 -0700
Subject: [PATCH] Add abiliy to add comments to exepcted history in tests

---
 tests/integrationbase.go      | 51 +++++++++++++++++++++++++++++++++--
 tests/update_workflow_test.go |  5 ++--
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/tests/integrationbase.go b/tests/integrationbase.go
index 680dad651c85..8cee71f5056d 100644
--- a/tests/integrationbase.go
+++ b/tests/integrationbase.go
@@ -350,12 +350,59 @@ func (s *IntegrationBase) formatHistoryCompact(history *historypb.History) strin
 
 func (s *IntegrationBase) EqualHistory(expectedHistory string, actualHistory *historypb.History) {
 	s.T().Helper()
-	expectedHistoryTrimmed := strings.Trim(expectedHistory, "\n")
+	expectedHistorySanitized := s.sanitizeHistory(expectedHistory)
 	actualHistoryStr := s.formatHistoryCompact(actualHistory)
-	s.Equal(expectedHistoryTrimmed, actualHistoryStr)
+	s.Equal(expectedHistorySanitized, actualHistoryStr)
 }
 
 func (s *IntegrationBase) EqualHistoryEvents(expectedHistory string, actualHistoryEvents []*historypb.HistoryEvent) {
 	s.T().Helper()
 	s.EqualHistory(expectedHistory, &historypb.History{Events: actualHistoryEvents})
 }
+
+// sanitizeHistory removes all the new lines and comments from expectedHistory string.
+func (s *IntegrationBase) sanitizeHistory(expectedHistory string) string {
+	type state int
+	const (
+		stateStart state = iota
+		stateNormal
+		stateComment
+		stateNewLine
+	)
+
+	var ret strings.Builder
+	st := stateStart
+	spaceCounter := 0
+	for _, r := range expectedHistory {
+		switch r {
+		case '\n':
+			if st == stateStart {
+				continue
+			}
+			spaceCounter = 0
+			st = stateNewLine
+		case '/':
+			spaceCounter = 0
+			st = stateComment
+		case ' ':
+			if st == stateComment {
+				continue
+			}
+			spaceCounter++
+		default:
+			if st == stateComment {
+				continue
+			}
+			if st == stateNewLine {
+				ret.WriteRune('\n')
+			}
+			if spaceCounter > 0 {
+				ret.WriteString(strings.Repeat(" ", spaceCounter))
+				spaceCounter = 0
+			}
+			ret.WriteRune(r)
+			st = stateNormal
+		}
+	}
+	return ret.String()
+}
diff --git a/tests/update_workflow_test.go b/tests/update_workflow_test.go
index 1b110c63ab38..ab3aedbd5802 100644
--- a/tests/update_workflow_test.go
+++ b/tests/update_workflow_test.go
@@ -314,7 +314,7 @@ func (s *integrationSuite) TestUpdateWorkflow_NewSpeculativeWorkflowTask_AcceptC
   3 WorkflowTaskStarted
   4 WorkflowTaskCompleted
   5 ActivityTaskScheduled
-  6 WorkflowTaskScheduled
+  6 WorkflowTaskScheduled // Was speculative WT but was converted to normal WT.
   7 WorkflowTaskStarted
   8 WorkflowTaskCompleted
   9 WorkflowExecutionUpdateAccepted
@@ -322,7 +322,8 @@ func (s *integrationSuite) TestUpdateWorkflow_NewSpeculativeWorkflowTask_AcceptC
  11 WorkflowTaskScheduled
  12 WorkflowTaskStarted
  13 WorkflowTaskCompleted
- 14 WorkflowExecutionCompleted`, events)
+ 14 WorkflowExecutionCompleted
+`, events)
 		})
 	}
 }