More tests for waiter. (#199)

ioquatix · web-flow · commit 456e2cc7042c · 2022-12-03T13:03:11.000+13:00
diff --git a/lib/async/waiter.rb b/lib/async/waiter.rb
@@ -13,6 +13,8 @@ def initialize(parent: nil, finished: Async::Condition.new)
 			@parent = parent
 		end
 		
+		# Execute a child task and add it to the waiter.
+		# @asynchronous Executes the given block concurrently.
 		def async(parent: (@parent or Task.current), &block)
 			parent.async do |task|
 				yield(task)
@@ -22,16 +24,28 @@ def async(parent: (@parent or Task.current), &block)
 			end
 		end
 		
+		# Wait for the first `count` tasks to complete.
+		# @parameter count [Integer | Nil] The number of tasks to wait for.
+		# @returns [Array(Async::Task)] If an integer is given, the tasks which have completed.
+		# @returns [Async::Task] Otherwise, the first task to complete.
 		def first(count = nil)
-			while @done.size < count
+			minimum = count || 1
+			
+			while @done.size < minimum
 				@finished.wait
 			end
 			
 			return @done.shift(*count)
 		end
 		
+		# Wait for the first `count` tasks to complete.
+		# @parameter count [Integer | Nil] The number of tasks to wait for.
 		def wait(count = nil)
-			first(count).map(&:wait)
+			if count
+				first(count).map(&:wait)
+			else
+				first.wait
+			end
 		end
 	end
 end
diff --git a/test/async/waiter.rb b/test/async/waiter.rb
@@ -7,6 +7,14 @@
 	
 	let(:waiter) {subject.new}
 	
+	it "can wait for the first task to complete" do
+		waiter.async do
+			:result
+		end
+		
+		expect(waiter.wait).to be == :result
+	end
+	
 	it "can wait for a subset of tasks" do
 		3.times do
 			waiter.async do
