1 files changed, 14 insertions, 3 deletions
diff --git a/xmpp/element_buffer_test.go b/xmpp/element_buffer_test.go
index 3d03f0b..4b5700c 100644
--- a/xmpp/element_buffer_test.go
+++ b/xmpp/element_buffer_test.go
@@ -12,10 +12,12 @@ type bufTest struct {
 }
 
 func TestElementBuffer(t *testing.T) {
-	var tests = []bufTest{
+	tests := []bufTest{
 		bufTest{`<stream></stream>`, []int{1, 0}},
+		bufTest{`<stream/>`, []int{1, 0}},
 		bufTest{`<a><b>testing</b></a>`, []int{1, 2, 2, 1, 0}},
 		bufTest{`<a><!-- comment --><b>testing</b></a>`, []int{1, 1, 2, 2, 1, 0}},
+		bufTest{`<!-- comment --><a><b>testing</b></a>`, []int{0, 1, 2, 2, 1, 0}},
 	}
 
 	for _, v := range tests {
@@ -23,18 +25,27 @@ func TestElementBuffer(t *testing.T) {
 		d := xml.NewDecoder(r)
 		b := newElementBuffer()
 
-		for i, _ := range v.levels {
+		i := 0
+
+		for {
 			token, err := d.Token()
 			if err != nil {
+				if i != len(v.levels) {
+					t.Fatalf("Stopped parsing at unexpected index due to error `%v`!\n", err)
+				}
 				break
 			}
+
 			err = b.add(token)
 			if err != nil {
-				t.Fatalf("xmpp.elementBuffer.add() returned error: `%v`!\n", err)
+				t.Fatalf("add(token) failed with error `%v`!\n", err)
 			}
+
 			if b.level != v.levels[i] {
 				t.Fatalf("Indent level of xmpp.elementBuffer %d does not match value given by test data %d!\n", b.level, v.levels[i])
 			}
+
+			i += 1
 		}
 	}
 }