1
package internal
2

3
import (
4
	"context"
5
	"time"
6

7
	"github.com/onsi/gomega/types"
8
)
9

10
type Gomega struct {
11
	Fail           types.GomegaFailHandler
12
	THelper        func()
13
	DurationBundle DurationBundle
14
}
15

16
func NewGomega(bundle DurationBundle) *Gomega {
17
	return &Gomega{
18
		Fail:           nil,
19
		THelper:        nil,
20
		DurationBundle: bundle,
21
	}
22
}
23

24
func (g *Gomega) IsConfigured() bool {
25
	return g.Fail != nil && g.THelper != nil
26
}
27

28
func (g *Gomega) ConfigureWithFailHandler(fail types.GomegaFailHandler) *Gomega {
29
	g.Fail = fail
30
	g.THelper = func() {}
31
	return g
32
}
33

34
func (g *Gomega) ConfigureWithT(t types.GomegaTestingT) *Gomega {
35
	g.Fail = func(message string, _ ...int) {
36
		t.Helper()
37
		t.Fatalf("\n%s", message)
38
	}
39
	g.THelper = t.Helper
40
	return g
41
}
42

43
func (g *Gomega) Ω(actual any, extra ...any) types.Assertion {
44
	return g.ExpectWithOffset(0, actual, extra...)
45
}
46

47
func (g *Gomega) Expect(actual any, extra ...any) types.Assertion {
48
	return g.ExpectWithOffset(0, actual, extra...)
49
}
50

51
func (g *Gomega) ExpectWithOffset(offset int, actual any, extra ...any) types.Assertion {
52
	return NewAssertion(actual, g, offset, extra...)
53
}
54

55
func (g *Gomega) Eventually(actualOrCtx any, args ...any) types.AsyncAssertion {
56
	return g.makeAsyncAssertion(AsyncAssertionTypeEventually, 0, actualOrCtx, args...)
57
}
58

59
func (g *Gomega) EventuallyWithOffset(offset int, actualOrCtx any, args ...any) types.AsyncAssertion {
60
	return g.makeAsyncAssertion(AsyncAssertionTypeEventually, offset, actualOrCtx, args...)
61
}
62

63
func (g *Gomega) Consistently(actualOrCtx any, args ...any) types.AsyncAssertion {
64
	return g.makeAsyncAssertion(AsyncAssertionTypeConsistently, 0, actualOrCtx, args...)
65
}
66

67
func (g *Gomega) ConsistentlyWithOffset(offset int, actualOrCtx any, args ...any) types.AsyncAssertion {
68
	return g.makeAsyncAssertion(AsyncAssertionTypeConsistently, offset, actualOrCtx, args...)
69
}
70

71
func (g *Gomega) makeAsyncAssertion(asyncAssertionType AsyncAssertionType, offset int, actualOrCtx any, args ...any) types.AsyncAssertion {
72
	baseOffset := 3
73
	timeoutInterval := -time.Duration(1)
74
	pollingInterval := -time.Duration(1)
75
	intervals := []any{}
76
	var ctx context.Context
77

78
	actual := actualOrCtx
79
	startingIndex := 0
80
	if _, isCtx := actualOrCtx.(context.Context); isCtx && len(args) > 0 {
81
		// the first argument is a context, we should accept it as the context _only if_ it is **not** the only argument **and** the second argument is not a parseable duration
82
		// this is due to an unfortunate ambiguity in early version of Gomega in which multi-type durations are allowed after the actual
83
		if _, err := toDuration(args[0]); err != nil {
84
			ctx = actualOrCtx.(context.Context)
85
			actual = args[0]
86
			startingIndex = 1
87
		}
88
	}
89

90
	for _, arg := range args[startingIndex:] {
91
		switch v := arg.(type) {
92
		case context.Context:
93
			ctx = v
94
		default:
95
			intervals = append(intervals, arg)
96
		}
97
	}
98
	var err error
99
	if len(intervals) > 0 {
100
		timeoutInterval, err = toDuration(intervals[0])
101
		if err != nil {
102
			g.Fail(err.Error(), offset+baseOffset)
103
		}
104
	}
105
	if len(intervals) > 1 {
106
		pollingInterval, err = toDuration(intervals[1])
107
		if err != nil {
108
			g.Fail(err.Error(), offset+baseOffset)
109
		}
110
	}
111

112
	return NewAsyncAssertion(asyncAssertionType, actual, g, timeoutInterval, pollingInterval, 1, ctx, offset)
113
}
114

115
func (g *Gomega) SetDefaultEventuallyTimeout(t time.Duration) {
116
	g.DurationBundle.EventuallyTimeout = t
117
}
118

119
func (g *Gomega) SetDefaultEventuallyPollingInterval(t time.Duration) {
120
	g.DurationBundle.EventuallyPollingInterval = t
121
}
122

123
func (g *Gomega) SetDefaultConsistentlyDuration(t time.Duration) {
124
	g.DurationBundle.ConsistentlyDuration = t
125
}
126

127
func (g *Gomega) SetDefaultConsistentlyPollingInterval(t time.Duration) {
128
	g.DurationBundle.ConsistentlyPollingInterval = t
129
}
130

131
func (g *Gomega) EnforceDefaultTimeoutsWhenUsingContexts() {
132
	g.DurationBundle.EnforceDefaultTimeoutsWhenUsingContexts = true
133
}
134

135
func (g *Gomega) DisableDefaultTimeoutsWhenUsingContext() {
136
	g.DurationBundle.EnforceDefaultTimeoutsWhenUsingContexts = false
137
}
138

139