From 35fd6e7c8ab044a9f330b1b86e90f1c7256426ea Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 25 Apr 2024 12:24:23 +0200
Subject: [PATCH] add a test case

---
 go.mod               |  2 +-
 webtransport_test.go | 57 +++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index c72833b..2a85961 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/quic-go/quic-go v0.43.0
 	github.com/stretchr/testify v1.8.0
 	go.uber.org/mock v0.4.0
+	golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63
 )
 
 require (
@@ -19,7 +20,6 @@ require (
 	github.com/quic-go/qpack v0.4.0 // indirect
 	github.com/rogpeppe/go-internal v1.10.0 // indirect
 	golang.org/x/crypto v0.14.0 // indirect
-	golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63 // indirect
 	golang.org/x/mod v0.12.0 // indirect
 	golang.org/x/net v0.17.0 // indirect
 	golang.org/x/sys v0.15.0 // indirect
diff --git a/webtransport_test.go b/webtransport_test.go
index d819191..3d84ca8 100644
--- a/webtransport_test.go
+++ b/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,57 @@ func TestWriteCloseRace(t *testing.T) {
 	<-ready
 	close(ch)
 }
+
+func TestDatagrams(t *testing.T) {
+	const num = 100
+	var mx sync.Mutex
+	m := make(map[string]bool, num)
+
+	var counter int
+	done := make(chan struct{})
+	serverErrChan := make(chan error, 1)
+	sess, closeServer := establishSession(t, func(sess *webtransport.Session) {
+		defer close(done)
+		for {
+			b, err := sess.ReceiveDatagram(context.Background())
+			if err != nil {
+				return
+			}
+			mx.Lock()
+			if _, ok := m[string(b)]; !ok {
+				serverErrChan <- errors.New("received unexpected datagram")
+				return
+			}
+			m[string(b)] = true
+			mx.Unlock()
+			counter++
+		}
+	})
+	defer closeServer()
+
+	errChan := make(chan error, 1)
+
+	for i := 0; i < num; i++ {
+		b := make([]byte, 800)
+		rand.Read(b)
+		mx.Lock()
+		m[string(b)] = false
+		mx.Unlock()
+		if err := sess.SendDatagram(b); err != nil {
+			break
+		}
+	}
+	time.Sleep(scaleDuration(10 * time.Millisecond))
+	sess.CloseWithError(0, "")
+	select {
+	case err := <-serverErrChan:
+		t.Fatal(err)
+	case err := <-errChan:
+		t.Fatal(err)
+	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")
+	}
+}