1 files changed, 54 insertions, 0 deletions
diff --git a/oae2_test.go b/oae2_test.go
new file mode 100644
index 0000000..91194bf
--- /dev/null
+++ b/oae2_test.go
@@ -0,0 +1,54 @@
+package oae2
+
+import (
+	"bytes"
+	"io"
+	"strings"
+	"testing"
+)
+
+func TestRoundTrip(t *testing.T) {
+	const (
+		msg       = "Hello World!"
+		key       = "password123"
+		blockSize = 1
+	)
+	buf := new(bytes.Buffer)
+	w := NewWriter(buf, []byte(key), blockSize)
+	if _, err := io.WriteString(w, msg); err != nil {
+		t.Fatalf("NewWriter.Write failed: %s", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("NewWriter.Close failed: %s", err)
+	}
+	got, err := io.ReadAll(NewReader(buf, []byte(key), blockSize))
+	if err != nil {
+		t.Fatalf("NewReader.Read failed: %s", err)
+	}
+	if string(got) != msg {
+		t.Errorf("Message failed to round-trip, got %q, want %q", got, msg)
+	}
+}
+
+func TestReadFromWriteTo(t *testing.T) {
+	const (
+		msg       = "Hello World!"
+		key       = "asdf"
+		blockSize = 2
+	)
+	buf := new(bytes.Buffer)
+	w := NewWriter(buf, []byte(key), blockSize)
+	if _, err := io.Copy(w, struct{ io.Reader }{strings.NewReader(msg)}); err != nil {
+		t.Fatalf("NewWriter.WriteTo failed: %s", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("NewWriter.Close failed: %s", err)
+	}
+	got := new(strings.Builder)
+	if _, err := io.Copy(got, NewReader(buf, []byte(key), blockSize)); err != nil {
+		t.Fatalf("NewReader.Read failed: %s", err)
+	}
+	if got.String() != msg {
+		t.Errorf("Message failed to round-trip, got %q, want %q", got, msg)
+	}
+}