Update Reader whence and tests

2 files changed, 67 insertions, 74 deletions

diff --git a/reader.go b/reader.go
index 426afc0..57a3172 100644
--- a/reader.go
+++ b/reader.go
@@ -13,9 +13,11 @@ var (
 	ErrInvalidOffset = errors.New("invalid offset")
 )
 
+type SeekFrom int
+
 const (
 	// SeekStart seeks relative to the origin of the file
-	SeekStart = 0
+	SeekStart SeekFrom = iota
 	// SeekCurrent seeks relative to the current offset
 	SeekCurrent = 1
 	// SeekEnd seeks relative to the end
@@ -60,14 +62,14 @@ func (r *Reader) ReadBit() bool {
 	return out
 }
 
-// Seek sets the internal pointer to position n.
-func (r *Reader) Seek(offset int64, whence int) (int64, error) {
+// Seek sets the internal pointer to position n. It returns the resulting offset.
+func (r *Reader) Seek(offset int64, whence SeekFrom) (int64, error) {
 	switch whence {
-	case 0:
+	case SeekStart:
 		r.i = offset
-	case 1:
+	case SeekCurrent:
 		r.i += offset
-	case 2:
+	case SeekEnd:
 		r.i = int64(r.ba.size) - offset
 	default:
 		return r.i, ErrInvalidWhence
diff --git a/reader_test.go b/reader_test.go
index db087a6..6c209c1 100644
--- a/reader_test.go
+++ b/reader_test.go
@@ -5,76 +5,67 @@ import (
 )
 
 func TestReadBits(t *testing.T) {
-	ba := NewFromBytes([]byte{0xf0, 0x01}, 16)
-	r := NewReader(ba)
-
-	if r.Pos() != 0 {
-		t.Errorf("got pos %d, want %d", r.Pos(), 0)
-	}
-
-	var test uint
-	if err := r.ReadBits(&test, 4); err != nil {
-		t.Errorf("failed with %q", err)
-	}
-	if test != 15 {
-		t.Errorf("got 15, want %d", test)
-	}
-	if r.Pos() != 4 {
-		t.Errorf("got pos %d, want %d", r.Pos(), 4)
-	}
-	r.Seek(-4, SeekCurrent)
-	if r.Pos() != 0 {
-		t.Errorf("got pos %d, want %d", r.Pos(), 0)
-	}
-
-	r.Seek(4, SeekStart)
-	if r.Pos() != 4 {
-		t.Errorf("got pos %d, want %d", r.Pos(), 4)
-	}
-	if err := r.ReadBits(&test, 4); err != nil {
-		t.Errorf("failed with %q", err)
-	}
-	if test != 0 {
-		t.Errorf("got 0, want %d", test)
-	}
-	if r.Pos() != 8 {
-		t.Errorf("got pos %d, want %d", r.Pos(), 8)
+	tests := map[string]struct {
+		ba *BitArray
+		// Seeking
+		offset     int64
+		dir        SeekFrom
+		seekOffset int64
+		// Reading
+		count      int
+		readOffset int64
+		want       uint
+	}{
+		"fromStart": {
+			ba:         NewFromBytes([]byte{0xf0, 0x01}, 16),
+			offset:     0,
+			dir:        SeekStart,
+			seekOffset: 0,
+			count:      4,
+			readOffset: 4,
+			want:       15,
+		},
+		"fromEnd": {
+			ba:         NewFromBytes([]byte{0xf0, 0x01}, 16),
+			offset:     4,
+			dir:        SeekEnd,
+			seekOffset: 12,
+			count:      4,
+			readOffset: 16,
+			want:       1,
+		},
+		"fromCurrent": {
+			ba:         NewFromBytes([]byte{0xf0, 0x01}, 16),
+			offset:     12,
+			dir:        SeekCurrent,
+			seekOffset: 12,
+			count:      4,
+			readOffset: 16,
+			want:       1,
+		},
 	}
 
-	if err := r.ReadBits(&test, 8); err != nil {
-		t.Errorf("failed with %q", err)
-	}
-	if test != 1 {
-		t.Errorf("got 1, want %d", test)
-	}
-	if r.Pos() != 16 {
-		t.Errorf("got pos %d, want %d", r.Pos(), 16)
-	}
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			r := NewReader(tt.ba)
+			actualPos, err := r.Seek(tt.offset, tt.dir)
+			if err != nil {
+				t.Fatalf("failed to seek: %s", err)
+			}
+			if actualPos != tt.seekOffset {
+				t.Errorf("got %d, want %d", actualPos, tt.seekOffset)
+			}
+			var actual uint
+			if err := r.ReadBits(&actual, tt.count); err != nil {
+				t.Fatalf("failed to read: %s", err)
+			}
+			if actual != tt.want {
+				t.Errorf("got %d, want %d", actual, tt.want)
+			}
+			if r.Pos() != tt.readOffset {
+				t.Errorf("got %d, want %d", r.Pos(), tt.readOffset)
+			}
 
-	// Read past end
-	if err := r.ReadBits(&test, 1); err == nil {
-		t.Errorf("expected error")
-	}
-
-	r.Seek(1, SeekEnd)
-	if r.Pos() != 15 {
-		t.Errorf("got pos %d, want %d", r.Pos(), 15)
-	}
-	if !r.ReadBit() {
-		t.Errorf("got 1, want 0")
-	}
-	if r.Pos() != 16 {
-		t.Errorf("got pos %d, want %d", r.Pos(), 16)
-	}
-
-	// Some errors
-	if _, err := r.Seek(1, 4); err == nil {
-		t.Errorf("expected err got none")
-	}
-	if _, err := r.Seek(17, SeekStart); err == nil {
-		t.Errorf("got none, want err")
-	}
-	if _, err := r.Seek(-1, SeekStart); err == nil {
-		t.Errorf("got none, want err")
+		})
 	}
 }