1// Copyright (C) 2019 Storj Labs, Inc.
2// See LICENSE for copying information.
3
4package grpccompat
5
6import (
7	"context"
8	"testing"
9
10	"github.com/zeebo/assert"
11	"github.com/zeebo/errs"
12)
13
14func TestCancel_ErrorAfterCancel(t *testing.T) {
15	impl := &serviceImpl{
16		Method4Fn: func(stream ServerMethod4Stream) error {
17			<-stream.Context().Done()
18			return errs.New("marker")
19		},
20	}
21
22	testCompat(t, impl, func(t *testing.T, cli Client, ensure func(*Out, error)) {
23		ctx, cancel := context.WithCancel(context.Background())
24		defer cancel()
25
26		stream, err := cli.Method4(ctx)
27		assert.NoError(t, err)
28
29		cancel()
30		ensure(stream.Recv())
31		ensure(nil, stream.Send(in(0)))
32	})
33}
34
35func TestCancel_CancelAfterError(t *testing.T) {
36	impl := &serviceImpl{
37		Method4Fn: func(stream ServerMethod4Stream) error {
38			return errs.New("marker")
39		},
40	}
41
42	testCompat(t, impl, func(t *testing.T, cli Client, ensure func(*Out, error)) {
43		ctx, cancel := context.WithCancel(context.Background())
44		defer cancel()
45
46		stream, err := cli.Method4(ctx)
47		assert.NoError(t, err)
48
49		ensure(stream.Recv())
50		ensure(nil, stream.Send(in(0)))
51		cancel()
52		ensure(stream.Recv())
53		ensure(nil, stream.Send(in(1)))
54	})
55}
56
57func TestCancel_CancelAfterSuccess(t *testing.T) {
58	impl := &serviceImpl{
59		Method4Fn: func(stream ServerMethod4Stream) error {
60			_ = stream.Send(out(2))
61			_, _ = stream.Recv()
62			<-stream.Context().Done()
63			return nil
64		},
65	}
66
67	testCompat(t, impl, func(t *testing.T, cli Client, ensure func(*Out, error)) {
68		ctx, cancel := context.WithCancel(context.Background())
69		defer cancel()
70
71		stream, err := cli.Method4(ctx)
72		assert.NoError(t, err)
73
74		ensure(stream.Recv())
75		ensure(nil, stream.Send(in(0)))
76		cancel()
77		ensure(stream.Recv())
78		ensure(nil, stream.Send(in(1)))
79	})
80}
81