diff --git a/app/services/tables/discomfort_calculator.rb b/app/services/tables/discomfort_calculator.rb
index 4a1f5c8..e1b413d 100644
--- a/app/services/tables/discomfort_calculator.rb
+++ b/app/services/tables/discomfort_calculator.rb
@@ -14,6 +14,7 @@ module Tables
     def cohesion_penalty
       table.map { |guest| guest.affinity_group_list.first }.combination(2).sum do |a, b|
         distance = AffinityGroupsHierarchy.instance.distance(a, b)
+
         next 1 if distance.nil?
         next 0 if distance.zero?
 
diff --git a/spec/services/tables/discomfort_calculator_spec.rb b/spec/services/tables/discomfort_calculator_spec.rb
index 9f391f7..939ebb4 100644
--- a/spec/services/tables/discomfort_calculator_spec.rb
+++ b/spec/services/tables/discomfort_calculator_spec.rb
@@ -4,44 +4,67 @@ module Tables
     let(:calculator) { described_class.new(table) }
 
     describe '#cohesion_penalty' do
+      before do
+        # Overridden in each test except trivial cases
+        allow(AffinityGroupsHierarchy.instance).to receive(:distance).and_call_original
+
+        %w[family friends work school].each do |group|
+          allow(AffinityGroupsHierarchy.instance).to receive(:distance).with(group, group).and_return(0)
+        end
+
+        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)
+        allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('family', 'school').and_return(3)
+        allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('friends', 'school').and_return(4)
+        allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('work', 'school').and_return(5)
+      end
       context 'when the table contains just two guests' do
-        let(:table) do
-          [
-            create(:guest, affinity_group_list: ['family']),
-            create(:guest, affinity_group_list: ['friends'])
-          ]
-        end
-
-        before do
-          allow(AffinityGroupsHierarchy.instance).to receive(:distance).and_return(distance)
-        end
-
         context 'when they belong to the same group' do
-          let(:distance) { 0 }
+          let(:table) { create_list(:guest, 2, affinity_group_list: ['family']) }
 
           it { expect(calculator.send(:cohesion_penalty)).to eq(0) }
         end
 
         context 'when they belong to completely unrelated groups' do
-          let(:distance) { nil }
-
+          let(:table) do
+            [
+              create(:guest, affinity_group_list: ['family']),
+              create(:guest, affinity_group_list: ['friends'])
+            ]
+          end
           it { expect(calculator.send(:cohesion_penalty)).to eq(1) }
         end
 
         context 'when they belong to groups at a distance of 1' do
-          let(:distance) { 1 }
+          let(:table) do
+            [
+              create(:guest, affinity_group_list: ['friends']),
+              create(:guest, affinity_group_list: ['work'])
+            ]
+          end
 
           it { expect(calculator.send(:cohesion_penalty)).to eq(0.5) }
         end
 
         context 'when they belong to groups at a distance of 2' do
-          let(:distance) { 2 }
+          let(:table) do
+            [
+              create(:guest, affinity_group_list: ['family']),
+              create(:guest, affinity_group_list: ['work'])
+            ]
+          end
 
           it { expect(calculator.send(:cohesion_penalty)).to eq(Rational(2, 3)) }
         end
 
         context 'when they belong to groups at a distance of 3' do
-          let(:distance) { 3 }
+          let(:table) do
+            [
+              create(:guest, affinity_group_list: ['family']),
+              create(:guest, affinity_group_list: ['school'])
+            ]
+          end
 
           it { expect(calculator.send(:cohesion_penalty)).to eq(Rational(3, 4)) }
         end
@@ -56,18 +79,12 @@ module Tables
           ]
         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
 
-      context 'when the table contains four guests' do
+      context 'when the table contains four guests of different groups' do
         let(:table) do
           [
             create(:guest, affinity_group_list: ['family']),
@@ -77,20 +94,26 @@ module Tables
           ]
         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)
-          allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('family', 'school').and_return(3)
-          allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('friends', 'school').and_return(4)
-          allow(AffinityGroupsHierarchy.instance).to receive(:distance).with('work', 'school').and_return(5)
-        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) + Rational(3, 4) + Rational(4, 5) + Rational(5, 6))
         end
       end
+
+      context 'when the table contains four guests of two split groups' do
+        let(:table) do
+          [
+            create(:guest, affinity_group_list: ['family']),
+            create(:guest, affinity_group_list: ['family']),
+            create(:guest, affinity_group_list: ['friends']),
+            create(:guest, affinity_group_list: ['friends'])
+          ]
+        end
+
+        it 'returns the sum of the penalties for each pair of guests' do
+          expect(calculator.send(:cohesion_penalty)).to eq(4)
+        end
+      end
     end
   end
 end