waitprocess/proc_ctx_test.go at master · siriusa51/waitprocess · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
package waitprocess

import (
	"context"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func TestRunWithCtx(t *testing.T) {
	t.Run("case-single-proc", func(t *testing.T) {
		wp := NewWaitProcess()

		stat := &teststate{}

		wp.RegisterProcess("test", RunWithCtx(func(ctx context.Context) error {
			<-ctx.Done()
			stat.add()
			return nil
		}))

		wp.Start()
		wp.Stop()
		err := wp.Wait()
		assert.Nil(t, err, "error should be nil")

		assert.Equal(t, 1, stat.getstate(), "state should be 1")
	})

	t.Run("case-multi-process", func(t *testing.T) {
		wp := NewWaitProcess()

		stat1 := &teststate{}
		stat2 := &teststate{}

		wp.RegisterProcess("test1", RunWithCtx(func(ctx context.Context) error {
			stat1.add()
			return nil
		})).RegisterProcess("test2", RunWithCtx(func(ctx context.Context) error {
			<-ctx.Done()
			stat2.add()
			return nil
		}))

		wp.Start()
		err := wp.Shutdown()
		assert.Nil(t, err, "error should be nil")

		assert.Equal(t, 1, stat1.getstate(), "state should be 1")
		assert.Equal(t, 1, stat2.getstate(), "state should be 1")
	})

	// test stop one process
	t.Run("case-multi-and-stop-one", func(t *testing.T) {
		wp := NewWaitProcess()

		stat := &teststate{}

		wp.RegisterProcess("noloop", RunWithCtx(func(ctx context.Context) error {
			time.Sleep(time.Second * 1)
			return nil
		})).RegisterProcess("loop", RunWithCtx(func(ctx context.Context) error {
			<-ctx.Done()
			stat.add()
			return nil
		}))

		err := wp.Run()
		assert.Nil(t, err, "error should be nil")

		assert.Equal(t, 1, stat.getstate(), "state should be 1")
	})

	// test error process
	t.Run("case-error-process", func(t *testing.T) {
		wp := NewWaitProcess()

		stat := &teststate{}

		wp.RegisterProcess("error", RunWithCtx(func(ctx context.Context) error {
			stat.add()
			return assert.AnError
		}))

		wp.RegisterProcess("loop", withTestprocess())

		err := wp.Run()
		assert.NotNil(t, err, "error should not be nil")
		assert.Equal(t, 1, stat.getstate(), "state should be 1")
	})
}