Rewrite in Go

1 files changed, 33 insertions, 0 deletions

diff --git a/words_test.go b/words_test.go
new file mode 100644
index 0000000..9c79464
--- /dev/null
+++ b/words_test.go
@@ -0,0 +1,33 @@
+package main
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestBuildNGrams(t *testing.T) {
+	const testWordList = `abc
+abd
+def
+`
+
+	got, err := buildNGrams(strings.NewReader(testWordList))
+	if err != nil {
+		t.Fatalf("buildNGrams failed: %s", err)
+	}
+	want := map[biGram][]choice{
+		{-1, -1}:   {{'a', 2}, {'d', 3}},
+		{-1, 'a'}:  {{'b', 2}},
+		{-1, 'd'}:  {{'e', 1}},
+		{'a', 'b'}: {{'c', 1}, {'d', 2}},
+		{'b', 'c'}: {{-1, 1}},
+		{'b', 'd'}: {{-1, 1}},
+		{'d', 'e'}: {{'f', 1}},
+		{'e', 'f'}: {{-1, 1}},
+	}
+	if diff := cmp.Diff(want, got, cmp.AllowUnexported(biGram{}, choice{})); diff != "" {
+		t.Errorf("buildNGrams(%q) returned unexpected diff (-want +got):\n%s", testWordList, diff)
+	}
+}