add a test case

webtransport_test.go

@@ -2,7 +2,6 @@ package webtransport_test
 
 import (
 	"context"
-	"crypto/rand"
 	"crypto/tls"
 	"errors"
 	"fmt"
@@ -15,6 +14,8 @@ import (
 	"testing"
 	"time"
 
+	"golang.org/x/exp/rand"
+
 	"github.com/quic-go/webtransport-go"
 
 	"github.com/quic-go/quic-go"
@@ -595,3 +596,65 @@ func TestWriteCloseRace(t *testing.T) {
 	<-ready
 	close(ch)
 }
+
+func TestDatagrams(t *testing.T) {
+	errChan := make(chan error, 1)
+	sess, closeServer := establishSession(t, func(sess *webtransport.Session) {
+		for {
+			d, err := sess.ReceiveDatagram(context.Background())
+			if err != nil {
+				errChan <- err
+				return
+			}
+			if err := sess.SendDatagram(d); err != nil {
+				errChan <- err
+				return
+			}
+		}
+	})
+	defer closeServer()
+
+	const num = 10
+	m := make(map[string]bool, num)
+	var mx sync.Mutex
+	done := make(chan struct{})
+	var counter int
+	go func() {
+		defer close(done)
+		for {
+			b, err := sess.ReceiveDatagram(context.Background())
+			fmt.Println("received datagram")
+			if err != nil {
+				fmt.Println(err)
+				return
+			}
+			mx.Lock()
+			if _, ok := m[string(b)]; !ok {
+				t.Fatal("received unexpected datagram")
+			}
+			m[string(b)] = true
+			mx.Unlock()
+			counter++
+		}
+	}()
+	for i := 0; i < num; i++ {
+		b := make([]byte, 800)
+		rand.Read(b)
+		mx.Lock()
+		m[string(b)] = false
+		mx.Unlock()
+		fmt.Println("sending datagram", i)
+		if err := sess.SendDatagram(b); err != nil {
+			break
+		}
+	}
+	time.Sleep(time.Second)
+	sess.CloseWithError(0, "")
+	select {
+	case <-done:
+		t.Logf("sent: %d, received: %d", num, counter)
+		require.Greater(t, counter, num*4/5)
+	case <-time.After(5 * time.Second):
+		t.Fatal("timeout")
+	}
+}