cloudflare/cloudflared
Publicmirrored from https://github.com/cloudflare/cloudflaredAvailable
datagramsession/session_test.go
59lines · modecode
| 1 | package datagramsession |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "net" |
| 6 | "testing" |
| 7 | |
| 8 | "github.com/google/uuid" |
| 9 | "github.com/stretchr/testify/require" |
| 10 | ) |
| 11 | |
| 12 | // TestCloseSession makes sure a session will stop after context is done |
| 13 | func TestSessionCtxDone(t *testing.T) { |
| 14 | testSessionReturns(t, true) |
| 15 | } |
| 16 | |
| 17 | // TestCloseSession makes sure a session will stop after close method is called |
| 18 | func TestCloseSession(t *testing.T) { |
| 19 | testSessionReturns(t, false) |
| 20 | } |
| 21 | |
| 22 | func testSessionReturns(t *testing.T, closeByContext bool) { |
| 23 | sessionID := uuid.New() |
| 24 | cfdConn, originConn := net.Pipe() |
| 25 | payload := testPayload(sessionID) |
| 26 | transport := &mockQUICTransport{ |
| 27 | reqChan: newDatagramChannel(), |
| 28 | respChan: newDatagramChannel(), |
| 29 | } |
| 30 | session := newSession(sessionID, transport, cfdConn) |
| 31 | |
| 32 | ctx, cancel := context.WithCancel(context.Background()) |
| 33 | sessionDone := make(chan struct{}) |
| 34 | go func() { |
| 35 | session.Serve(ctx) |
| 36 | close(sessionDone) |
| 37 | }() |
| 38 | |
| 39 | go func() { |
| 40 | n, err := session.writeToDst(payload) |
| 41 | require.NoError(t, err) |
| 42 | require.Equal(t, len(payload), n) |
| 43 | }() |
| 44 | |
| 45 | readBuffer := make([]byte, len(payload)+1) |
| 46 | n, err := originConn.Read(readBuffer) |
| 47 | require.NoError(t, err) |
| 48 | require.Equal(t, len(payload), n) |
| 49 | |
| 50 | if closeByContext { |
| 51 | cancel() |
| 52 | } else { |
| 53 | session.close() |
| 54 | } |
| 55 | |
| 56 | <-sessionDone |
| 57 | // call cancelled again otherwise the linter will warn about possible context leak |
| 58 | cancel() |
| 59 | } |