procedure.find_column raise NotFound to fit AR interface

app/models/concerns/columns_concern.rb

@@ -9,8 +9,12 @@ module ColumnsConcern
     # instead, we are using h_id == { procedure_id:, column_id: }
     # another way to find a column is to look for its label
     def find_column(h_id: nil, label: nil)
-      return columns.find { _1.h_id == h_id } if h_id.present?
-      return columns.find { _1.label == label } if label.present?
+      column = columns.find { _1.h_id == h_id } if h_id.present?
+      column = columns.find { _1.label == label } if label.present?
+
+      raise ActiveRecord::RecordNotFound if column.nil?
+
+      column
     end
 
     def columns

spec/models/concerns/columns_concern_spec.rb

@@ -1,6 +1,22 @@
 # frozen_string_literal: true
 
 describe ColumnsConcern do
+  describe '#find_column' do
+    let(:procedure) { build(:procedure) }
+    let(:notifications_column) { procedure.notifications_column }
+
+    it do
+      label = notifications_column.label
+      expect(procedure.find_column(label:)).to eq(notifications_column)
+
+      h_id = notifications_column.h_id
+      expect(procedure.find_column(h_id:)).to eq(notifications_column)
+
+      unknwon = 'unknown'
+      expect { procedure.find_column(h_id: unknwon) }.to raise_error(ActiveRecord::RecordNotFound)
+    end
+  end
+
   describe "#columns" do
     subject { procedure.columns }