Adding NumStreams to query open stream count

session.go

@@ -108,6 +108,14 @@ func (s *Session) IsClosed() bool {
 	}
 }
 
+// NumStreams returns the number of currently open streams
+func (s *Session) NumStreams() int {
+	s.streamLock.Lock()
+	num := len(s.streams)
+	s.streamLock.Unlock()
+	return num
+}
+
 // Open is used to create a new stream as a net.Conn
 func (s *Session) Open() (net.Conn, error) {
 	return s.OpenStream()

session_test.go

@@ -76,6 +76,13 @@ func TestAccept(t *testing.T) {
 	defer client.Close()
 	defer server.Close()
 
+	if client.NumStreams() != 0 {
+		t.Fatalf("bad")
+	}
+	if server.NumStreams() != 0 {
+		t.Fatalf("bad")
+	}
+
 	wg := &sync.WaitGroup{}
 	wg.Add(4)
 
@@ -163,6 +170,10 @@ func TestSendData_Small(t *testing.T) {
 			t.Fatalf("err: %v", err)
 		}
 
+		if server.NumStreams() != 1 {
+			t.Fatalf("bad")
+		}
+
 		buf := make([]byte, 4)
 		for i := 0; i < 1000; i++ {
 			n, err := stream.Read(buf)
@@ -189,6 +200,10 @@ func TestSendData_Small(t *testing.T) {
 			t.Fatalf("err: %v", err)
 		}
 
+		if client.NumStreams() != 1 {
+			t.Fatalf("bad")
+		}
+
 		for i := 0; i < 1000; i++ {
 			n, err := stream.Write([]byte("test"))
 			if err != nil {
@@ -214,6 +229,13 @@ func TestSendData_Small(t *testing.T) {
 	case <-time.After(time.Second):
 		panic("timeout")
 	}
+
+	if client.NumStreams() != 0 {
+		t.Fatalf("bad")
+	}
+	if server.NumStreams() != 0 {
+		t.Fatalf("bad")
+	}
 }
 
 func TestSendData_Large(t *testing.T) {