From 239d6452e0d24089470b7889cc1fda9fdfea795f Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Mon, 22 Jul 2024 13:52:42 +1200
Subject: [PATCH] Make it possible to include/prepend shared modules.

---
 lib/sus/context.rb          | 46 +++++++++++++++++++++++++------------
 lib/sus/include_context.rb  |  4 ++--
 lib/sus/shared.rb           | 10 +++++---
 test/sus/include_context.rb | 28 ++++++++++++++--------
 4 files changed, 58 insertions(+), 30 deletions(-)

diff --git a/lib/sus/context.rb b/lib/sus/context.rb
index 303d2dd..c0f4bec 100644
--- a/lib/sus/context.rb
+++ b/lib/sus/context.rb
@@ -76,38 +76,54 @@ def each(&block)
 			end
 		end
 		
-		def before(&block)
+		# Include an around method to the context class, that invokes the given block before running the test.
+		#
+		# Before hooks are called in the reverse order they are defined, in other words the last defined before hook is called first.
+		#
+		# @parameter hook [Proc] The block to execute before each test.
+		def before(&hook)
 			wrapper = Module.new
 			
-			wrapper.define_method(:before) do
-				super()
-				instance_exec(&block)
+			wrapper.define_method(:around) do |&block|
+				instance_exec(&hook)
+				super(&block)
 			end
 			
 			self.include(wrapper)
 		end
 		
-		def after(&block)
+		# Include an around method to the context class, that invokes the given block after running the test.
+		#
+		# After hooks are called in the order they are defined, in other words the last defined after hook is called last.
+		#
+		# @parameter hook [Proc] The block to execute after each test. An `error` argument is passed if the test failed with an exception.
+		def after(&hook)
 			wrapper = Module.new
 			
-			wrapper.define_method(:after) do
-				instance_exec(&block)
-				super()
+			wrapper.define_method(:around) do |&block|
+				error = nil
+				
+				super(&block)
+			rescue => error
+				raise
+			ensure
+				instance_exec(error, &hook)
 			end
 			
 			self.include(wrapper)
 		end
 		
+		# Add an around hook to the context class.
+		#
+		# Around hooks are called in the reverse order they are defined.
+		#
+		# The top level `around` implementation invokes before and after hooks.
+		#
+		# @paremeter block [Proc] The block to execute around each test.
 		def around(&block)
 			wrapper = Module.new
 			
-			wrapper.define_method(:around) do
-				call_super = proc do
-					super()
-				end
-				
-				instance_exec(call_super, &block)
-			end
+			wrapper.define_method(:around, &block)
 			
 			self.include(wrapper)
 		end
diff --git a/lib/sus/include_context.rb b/lib/sus/include_context.rb
index f3cafda..0ce6452 100644
--- a/lib/sus/include_context.rb
+++ b/lib/sus/include_context.rb
@@ -7,8 +7,8 @@
 
 module Sus
 	module Context
-		def include_context(shared, ...)
-			shared.included(self, ...)
+		def include_context(shared, *arguments, **options)
+			self.class_exec(*arguments, **options, &shared.block)
 		end
 	end
 end
diff --git a/lib/sus/shared.rb b/lib/sus/shared.rb
index 71bb0d2..e340d80 100644
--- a/lib/sus/shared.rb
+++ b/lib/sus/shared.rb
@@ -11,7 +11,7 @@ module Shared
 		attr_accessor :block
 		
 		def self.build(name, block)
-			base = Class.new
+			base = Module.new
 			base.extend(Shared)
 			base.name = name
 			base.block = block
@@ -19,8 +19,12 @@ def self.build(name, block)
 			return base
 		end
 		
-		def included(base, *arguments, **options)
-			base.class_exec(*arguments, **options, &self.block)
+		def included(base)
+			base.class_exec(&self.block)
+		end
+		
+		def prepended(base)
+			base.class_exec(&self.block)
 		end
 	end
 	
diff --git a/test/sus/include_context.rb b/test/sus/include_context.rb
index c9a7e6a..5199824 100644
--- a/test/sus/include_context.rb
+++ b/test/sus/include_context.rb
@@ -3,38 +3,46 @@
 # Released under the MIT License.
 # Copyright, 2023, by Samuel Williams.
 
-AThing = Sus::Shared("a thing") do |key, value: 42|
+AContextWithArguments = Sus::Shared("a context with arguments") do |key, value: 42|
 	let(:a_thing) {{key => value}}
-	
+end
+
+AContextWithHooks = Sus::Shared("a context with hooks") do
 	before do
-		$stderr.puts "before: #{self}"
 		events << :shared_before
 	end
 	
 	after do
-		$stderr.puts "after: #{self}"
 		events << :shared_after
 	end
 	
-	around do |block|
-		$stderr.puts "around: #{self}"
-		block.call
+	around do |&block|
+		events << :shared_around_before
+		super(&block)
 	end
 end
 
 describe Sus::Context do
 	with '.include_context' do
 		with "a shared context with an option" do
+			include_context AContextWithArguments, :key, value: 42
+			
+			it "can include a shared context with arguments" do
+				expect(a_thing).to be == {:key => 42}
+			end
+		end
+		
+		with "a shared context with arguments" do
 			let(:events) {Array.new}
-			include_context AThing, :key, value: 42
+			
+			include AContextWithHooks
 			
 			before do
 				events << :example_before
 			end
 			
 			it "can include a shared context" do
-				expect(a_thing).to be == {:key => 42}
-				expect(events).to be == [:shared_before, :example_before]
+				expect(events).to be == [:example_before, :shared_around_before, :shared_before]
 			end
 		end
 	end