From 761a11160e599e4583760075dfaf620e2b946e1b Mon Sep 17 00:00:00 2001
From: Ignacio Perez <ignacio.perez@rootstrap.com>
Date: Tue, 16 Dec 2025 16:04:39 -0300
Subject: [PATCH] Fix tests

---
 .rubocop.yml                 |  6 ++++--
 spec/mars/agent_spec.rb      | 34 ++++++++++++++++++++++------------
 spec/mars/aggregator_spec.rb | 35 +++++++----------------------------
 spec/mars/gate_spec.rb       |  5 +++--
 4 files changed, 36 insertions(+), 44 deletions(-)

diff --git a/.rubocop.yml b/.rubocop.yml
index 6303aba..ebb2677 100644
--- a/.rubocop.yml
+++ b/.rubocop.yml
@@ -14,8 +14,10 @@ Metrics/MethodLength:
   Exclude:
   - "lib/mars/workflows/parallel.rb"
 
-RSpec/ExampleLength:
-  Enabled: false
+Metrics/ParameterLists:
+  Enabled: true
+  Exclude:
+  - "lib/mars/agent.rb"
 
 Style/Documentation:
   Enabled: false
diff --git a/spec/mars/agent_spec.rb b/spec/mars/agent_spec.rb
index 89ce677..c59f922 100644
--- a/spec/mars/agent_spec.rb
+++ b/spec/mars/agent_spec.rb
@@ -2,12 +2,16 @@
 
 RSpec.describe Mars::Agent do
   describe "#run" do
+    subject(:run_agent) { agent.run("input text") }
+
     let(:agent) { described_class.new(name: "TestAgent", options: { model: "test-model" }) }
     let(:mock_chat_instance) do
       instance_double("RubyLLM::Chat").tap do |mock|
-        allow(mock).to receive_messages(with_tools: mock, with_schema: mock, ask: nil)
+        allow(mock).to receive_messages(with_tools: mock, with_schema: mock, with_instructions: mock,
+                                        ask: mock_chat_response)
       end
     end
+    let(:mock_chat_response) { instance_double("RubyLLM::Message", content: "response text") }
     let(:mock_chat_class) { class_double("RubyLLM::Chat", new: mock_chat_instance) }
 
     before do
@@ -15,25 +19,31 @@
     end
 
     it "initializes RubyLLM::Chat with provided options" do
-      agent.run("test input")
+      run_agent
 
       expect(mock_chat_class).to have_received(:new).with(model: "test-model")
     end
 
-    it "configures chat with tools if provided" do
-      tools = [proc { "tool" }]
-      agent_with_tools = described_class.new(name: "TestAgent", tools: tools)
-      agent_with_tools.run("test input")
+    context "when tools are provided" do
+      let(:tools) { [proc { "tool1" }, proc { "tool2" }] }
+      let(:agent) { described_class.new(name: "TestAgent", tools: tools) }
+
+      it "configures chat with tools" do
+        run_agent
 
-      expect(mock_chat_instance).to have_received(:with_tools).with(tools)
+        expect(mock_chat_instance).to have_received(:with_tools).with(*tools)
+      end
     end
 
-    it "configures chat with schema if provided" do
-      schema = { type: "object" }
-      agent_with_schema = described_class.new(name: "TestAgent", schema: schema)
+    context "when schema is provided" do
+      let(:schema) { { type: "object" } }
+      let(:agent) { described_class.new(name: "TestAgent", schema: schema) }
 
-      agent_with_schema.run("test input")
-      expect(mock_chat_instance).to have_received(:with_schema).with(schema)
+      it "configures chat with schema" do
+        run_agent
+
+        expect(mock_chat_instance).to have_received(:with_schema).with(schema)
+      end
     end
   end
 end
diff --git a/spec/mars/aggregator_spec.rb b/spec/mars/aggregator_spec.rb
index 1c2b6c1..cb266ed 100644
--- a/spec/mars/aggregator_spec.rb
+++ b/spec/mars/aggregator_spec.rb
@@ -2,9 +2,9 @@
 
 RSpec.describe Mars::Aggregator do
   describe "#run" do
-    let(:aggregator) { described_class.new }
-
     context "when called without a block" do
+      let(:aggregator) { described_class.new }
+
       it "joins inputs with newlines" do
         inputs = %w[first second third]
         result = aggregator.run(inputs)
@@ -15,35 +15,14 @@
         result = aggregator.run([])
         expect(result).to eq("")
       end
-
-      it "handles single input" do
-        result = aggregator.run(["single"])
-        expect(result).to eq("single")
-      end
-
-      it "handles numeric inputs" do
-        inputs = [1, 2, 3]
-        result = aggregator.run(inputs)
-        expect(result).to eq("1\n2\n3")
-      end
     end
 
-    context "when called with a block" do
-      it "executes the block and returns its value" do
-        result = aggregator.run(["ignored"]) { "block result" }
-        expect(result).to eq("block result")
-      end
-
-      it "ignores the inputs when block is given" do
-        inputs = %w[first second]
-        result = aggregator.run(inputs) { "custom aggregation" }
-        expect(result).to eq("custom aggregation")
-      end
+    context "when initialized with a block operation" do
+      let(:aggregator) { described_class.new("Aggregator", operation: lambda(&:join)) }
 
-      it "can perform custom aggregation logic" do
-        inputs = [1, 2, 3, 4, 5]
-        result = aggregator.run(inputs) { inputs.sum }
-        expect(result).to eq(15)
+      it "executes the block and returns its value" do
+        result = aggregator.run(%w[a b c])
+        expect(result).to eq("abc")
       end
     end
   end
diff --git a/spec/mars/gate_spec.rb b/spec/mars/gate_spec.rb
index aac613e..b133277 100644
--- a/spec/mars/gate_spec.rb
+++ b/spec/mars/gate_spec.rb
@@ -69,9 +69,10 @@
         expect(high_branch).to have_received(:run).with(10)
       end
 
-      it "raises an error when branch is not defined" do
+      it "returns the input when branch is not defined" do
         # For input 3, condition returns "low" which is not in branches
-        expect { gate.run(3) }.to raise_error(NoMethodError)
+        result = gate.run(3)
+        expect(result).to eq(3)
       end
     end