talos

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package conditions_test
6

7
import (
8
	"context"
9
	"testing"
10
	"time"
11

12
	"github.com/stretchr/testify/suite"
13

14
	"github.com/siderolabs/talos/pkg/conditions"
15
)
16

17
type AllSuite struct {
18
	suite.Suite
19
}
20

21
type MockCondition struct {
22
	description string
23
	errCh       chan error
24
}
25

26
func (mc *MockCondition) String() string {
27
	return mc.description
28
}
29

30
func (mc *MockCondition) Wait(ctx context.Context) error {
31
	select {
32
	case <-ctx.Done():
33
		return ctx.Err()
34
	case err := <-mc.errCh:
35
		return err
36
	}
37
}
38

39
func (suite *AllSuite) TestString() {
40
	suite.Require().Equal("A, B", conditions.WaitForAll(
41
		&MockCondition{description: "A"},
42
		&MockCondition{description: "B"},
43
	).String())
44

45
	suite.Require().Equal("A", conditions.WaitForAll(
46
		&MockCondition{description: "A"},
47
	).String())
48

49
	conds := []conditions.Condition{
50
		&MockCondition{description: "A", errCh: make(chan error)},
51
		&MockCondition{description: "B", errCh: make(chan error)},
52
	}
53

54
	waiter := conditions.WaitForAll(conds...)
55

56
	done := make(chan error)
57

58
	go func() {
59
		done <- waiter.Wait(context.Background())
60
	}()
61

62
	suite.Require().Equal("A, B", waiter.String())
63
	conds[0].(*MockCondition).errCh <- nil
64
	time.Sleep(50 * time.Millisecond)
65

66
	// done waiting for 'A', so description should now be shorter
67
	suite.Require().Equal("B", waiter.String())
68

69
	conds[1].(*MockCondition).errCh <- nil
70
	<-done
71
}
72

73
func (suite *AllSuite) TestFlatten() {
74
	conds1 := []conditions.Condition{
75
		&MockCondition{description: "A", errCh: make(chan error)},
76
		&MockCondition{description: "B", errCh: make(chan error)},
77
	}
78
	conds2 := []conditions.Condition{
79
		&MockCondition{description: "C", errCh: make(chan error)},
80
		&MockCondition{description: "D", errCh: make(chan error)},
81
	}
82

83
	waiter := conditions.WaitForAll(conditions.WaitForAll(conds1...), conditions.WaitForAll(conds2...))
84
	suite.Require().Equal("A, B, C, D", waiter.String())
85

86
	ctx, ctxCancel := context.WithCancel(context.Background())
87
	defer ctxCancel()
88

89
	done := make(chan error)
90

91
	go func() {
92
		done <- waiter.Wait(ctx)
93
	}()
94

95
	conds1[0].(*MockCondition).errCh <- nil
96
	conds2[1].(*MockCondition).errCh <- nil
97

98
	time.Sleep(50 * time.Millisecond)
99

100
	suite.Require().Equal("B, C", waiter.String())
101

102
	ctxCancel()
103

104
	suite.Require().Equal(context.Canceled, <-done)
105
}
106

107
func TestAllSuite(t *testing.T) {
108
	suite.Run(t, new(AllSuite))
109
}
110