Add unit test for Generate. Trim leading and trailing whitespace from…

… generated sentences.
colinfike · Feb 8, 2020 · d3d2dff · d3d2dff
1 parent 111f11b
commit d3d2dff
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/mimic.go b/mimic.go
@@ -66,7 +66,7 @@ func (m *MarkovChain) Generate() string {
 		sentence = sentence + delim + suffix
 		prefix = nextPrefix(prefix, suffix)
 	}
-	return sentence
+	return strings.TrimSpace(sentence)
 }
 
 func nextPrefix(prefix, suffix string) string {

diff --git a/mimic_test.go b/mimic_test.go
@@ -86,6 +86,14 @@ func TestTrainThreeValid(t *testing.T) {
 	assert.Equal(t, markov.chain, expectedMap)
 }
 
+func TestGenerate(t *testing.T) {
+	markov := NewMarkovChain(2)
+	testSet := []string{"Hello there jon"}
+	markov.Train(testSet)
+
+	assert.Equal(t, "hello there jon", markov.Generate())
+}
+
 var nextSuffixTests = []struct {
 	prefix         string
 	suffix         string