Run container tests in async context

Author: samuel-williams-shopify

fixtures/container_context.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require "async"
+require "async/container/group"
+
+module ContainerContext
+	def container_context(&block)
+		# Newer async-container supervises containers using Async tasks, while
+		# released versions can fail when forking under a scheduler on Ruby 3.x.
+		if Async::Container::Group.method_defined?(:supervise)
+			Sync(&block)
+		else
+			block.call
+		end
+	end
+end

test/async/service/managed/service.rb

@@ -10,8 +10,11 @@
 require "async"
 
 require "sus/fixtures/async/scheduler_context"
+require "container_context"
 
 describe Async::Service::Managed::Service do
+	include ContainerContext
+	
 	let(:configuration) do
 		Async::Service::Configuration.build do
 			service "test-container" do
@@ -365,15 +368,15 @@ def name = @name
 		let(:controller) {Async::Service::Controller.for(test_service)}
 
 		it "runs service with health checking and no restarts when async context is present" do
-			container = Async::Container.new
-			
-			begin
+			container_context do
+				container = Async::Container.new
+				
 				controller.setup(container)
 				controller.start
 				sleep(0.03)
 			ensure
 				controller.stop
-				container.stop
+				container&.stop
 			end
 		end
 	end

test/async/service/policy.rb

@@ -6,8 +6,11 @@
 require "async/container/best"
 require "async/service/policy"
 require "async/container/statistics"
+require "container_context"
 
 describe Async::Service::Policy do
+	include ContainerContext
+	
 	let(:policy) {subject.new(maximum_failures: 5, window: 10)}
 
 	with "::DEFAULT" do
@@ -161,18 +164,28 @@ def success_status.success?; true; end
 
 	with "concurrent failures" do
 		it "only stops container once when multiple children fail simultaneously" do
-			container = Async::Container.best_container_class.new(policy: policy)
-			expect(container).to receive(:stop)
-			
-			# Spawn 10 children that all fail immediately:
-			10.times do |i|
-				container.spawn(name: "worker-#{i}") do |instance|
-					instance.ready!
-					exit(1)
+			container_context do
+				container = Async::Container.best_container_class.new(policy: policy)
+				stop_count = 0
+				original_stop = container.method(:stop)
+				
+				container.define_singleton_method(:stop) do |*arguments|
+					stop_count += 1
+					original_stop.call(*arguments)
+				end
+				
+				# Spawn 10 children that all fail immediately:
+				10.times do |i|
+					container.spawn(name: "worker-#{i}") do |instance|
+						instance.ready!
+						exit(1)
+					end
 				end
+				
+				container.wait
+				
+				expect(stop_count).to be == 1
 			end
-			
-			container.wait
 		end
 	end
 end