7
7
8
8
module Sus
9
9
class Receive
10
- CALL_ORIGINAL = Object . new
11
-
12
- def initialize ( base , method )
10
+ def initialize ( base , method , &block )
13
11
@base = base
14
12
@method = method
15
13
16
14
@times = Times . new
17
15
@arguments = nil
18
16
@options = nil
19
17
@block = nil
20
- @returning = CALL_ORIGINAL
18
+
19
+ @returning = block
21
20
end
22
21
23
22
def print ( output )
@@ -60,12 +59,27 @@ def with_call_count(predicate)
60
59
return self
61
60
end
62
61
63
- def and_return ( *returning )
64
- if returning . size == 1
65
- @returning = returning . first
62
+ def and_return ( *returning , &block )
63
+ if block_given?
64
+ if returning . any?
65
+ raise ArgumentError , "Cannot specify both a block and returning values."
66
+ end
67
+
68
+ @returning = block
69
+ elsif returning . size == 1
70
+ @returning = proc { returning . first }
66
71
else
67
- @returning = returning
72
+ @returning = proc { returning }
73
+ end
74
+
75
+ return self
76
+ end
77
+
78
+ def and_raise ( ...)
79
+ @returning = proc do
80
+ raise ( ...)
68
81
end
82
+
69
83
return self
70
84
end
71
85
@@ -97,7 +111,7 @@ def call(assertions, subject)
97
111
98
112
validate ( mock , assertions , arguments , options , block )
99
113
100
- next @returning
114
+ next @returning . call ( * arguments , ** options , & block )
101
115
end
102
116
end
103
117
@@ -110,7 +124,7 @@ def call(assertions, subject)
110
124
end
111
125
112
126
def call_original?
113
- @returning == CALL_ORIGINAL
127
+ @returning . nil?
114
128
end
115
129
116
130
class WithArguments
@@ -128,7 +142,7 @@ def call(assertions, subject)
128
142
end
129
143
end
130
144
end
131
-
145
+
132
146
class WithOptions
133
147
def initialize ( predicate )
134
148
@predicate = predicate
@@ -153,15 +167,15 @@ def initialize(predicate)
153
167
def print ( output )
154
168
output . write ( "with block" , @predicate )
155
169
end
156
-
170
+
157
171
def call ( assertions , subject )
158
172
assertions . nested ( self ) do |assertions |
159
173
160
174
Expect . new ( assertions , subject ) . not . to ( Be == nil )
161
175
end
162
176
end
163
177
end
164
-
178
+
165
179
class Times
166
180
ONCE = Be . new ( :== , 1 )
167
181
@@ -172,7 +186,7 @@ def initialize(condition = ONCE)
172
186
def print ( output )
173
187
output . write ( "with call count " , @condition )
174
188
end
175
-
189
+
176
190
def call ( assertions , subject )
177
191
assertions . nested ( self ) do |assertions |
178
192
Expect . new ( assertions , subject ) . to ( @condition )
@@ -182,8 +196,8 @@ def call(assertions, subject)
182
196
end
183
197
184
198
class Base
185
- def receive ( method )
186
- Receive . new ( self , method )
199
+ def receive ( method , & block )
200
+ Receive . new ( self , method , & block )
187
201
end
188
202
end
189
203
end
0 commit comments