Include additional test case for three people

spec/services/tables/discomfort_calculator_spec.rb

@@ -2,7 +2,7 @@ require 'rails_helper'
 module Tables
   RSpec.describe DiscomfortCalculator do
     let(:calculator) { described_class.new(table) }
- 
+
     describe '#cohesion_penalty' do
       context 'when the table contains just two guests' do
         let(:table) do
@@ -46,6 +46,26 @@ module Tables
           it { expect(calculator.send(:cohesion_penalty)).to eq(Rational(3, 4)) }
         end
       end
+
+      context 'when the table contains three guests' do
+        let(:table) do
+          [
+            create(:guest, affinity_group_list: ['family']),
+            create(:guest, affinity_group_list: ['friends']),
+            create(:guest, affinity_group_list: ['work'])
+          ]
+        end
+
+        before do
+          allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('family', 'friends').and_return(nil)
+          allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('friends', 'work').and_return(1)
+          allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('family', 'work').and_return(2)
+        end
+
+        it 'returns the sum of the penalties for each pair of guests' do
+          expect(calculator.send(:cohesion_penalty)).to eq(1 + Rational(1, 2) + Rational(2, 3))
+        end
+      end
     end
   end
 end