1 files changed, 75 insertions, 60 deletions
diff --git a/spec/models/concerns/sha_attribute_spec.rb b/spec/models/concerns/sha_attribute_spec.rb
index 1bcf3dc8b61..790e6936803 100644
--- a/spec/models/concerns/sha_attribute_spec.rb
+++ b/spec/models/concerns/sha_attribute_spec.rb
@@ -3,86 +3,101 @@
 require 'spec_helper'
 
 RSpec.describe ShaAttribute do
-  let(:model) { Class.new(ActiveRecord::Base) { include ShaAttribute } }
+  let(:model) do
+    Class.new(ActiveRecord::Base) do
+      include ShaAttribute
 
-  before do
-    columns = [
-      double(:column, name: 'name', type: :text),
-      double(:column, name: 'sha1', type: :binary)
-    ]
-
-    allow(model).to receive(:columns).and_return(columns)
+      self.table_name = 'merge_requests'
+    end
   end
 
-  describe '#sha_attribute' do
-    context 'when in development' do
-      before do
-        stub_rails_env('development')
-      end
+  let(:binary_column) { :merge_ref_sha }
+  let(:text_column) { :target_branch }
 
-      context 'when the table exists' do
-        before do
-          allow(model).to receive(:table_exists?).and_return(true)
-        end
+  describe '.sha_attribute' do
+    it 'defines a SHA attribute with Gitlab::Database::ShaAttribute type' do
+      expect(model).to receive(:attribute)
+        .with(binary_column, an_instance_of(Gitlab::Database::ShaAttribute))
+        .and_call_original
 
-        it 'defines a SHA attribute for a binary column' do
-          expect(model).to receive(:attribute)
-            .with(:sha1, an_instance_of(Gitlab::Database::ShaAttribute))
+      model.sha_attribute(binary_column)
+    end
+  end
 
-          model.sha_attribute(:sha1)
-        end
+  describe '.sha256_attribute' do
+    it 'defines a SHA256 attribute with Gitlab::Database::ShaAttribute type' do
+      expect(model).to receive(:attribute)
+        .with(binary_column, an_instance_of(Gitlab::Database::Sha256Attribute))
+        .and_call_original
 
-        it 'raises ArgumentError when the column type is not :binary' do
-          expect { model.sha_attribute(:name) }.to raise_error(ArgumentError)
-        end
-      end
+      model.sha256_attribute(binary_column)
+    end
+  end
 
-      context 'when the table does not exist' do
-        it 'allows the attribute to be added' do
-          allow(model).to receive(:table_exists?).and_return(false)
+  describe '.load_schema!' do
+    # load_schema! is not a documented class method, so use a documented method
+    # that we know will call load_schema!
+    def load_schema!
+      expect(model).to receive(:load_schema!).and_call_original
 
-          expect(model).not_to receive(:columns)
-          expect(model).to receive(:attribute)
+      model.new
+    end
 
-          model.sha_attribute(:name)
-        end
-      end
+    using RSpec::Parameterized::TableSyntax
 
-      context 'when the column does not exist' do
-        it 'allows the attribute to be added' do
-          allow(model).to receive(:table_exists?).and_return(true)
+    where(:column_name, :environment, :expected_error) do
+      ref(:binary_column)    | 'development' | :no_error
+      ref(:binary_column)    | 'production'  | :no_error
+      ref(:text_column)      | 'development' | :sha_mismatch_error
+      ref(:text_column)      | 'production'  | :no_error
+      :__non_existent_column | 'development' | :no_error
+      :__non_existent_column | 'production'  | :no_error
+    end
 
-          expect(model).to receive(:columns)
-          expect(model).to receive(:attribute)
+    let(:sha_mismatch_error) do
+      [
+        described_class::ShaAttributeTypeMismatchError,
+        /#{column_name}.* should be a :binary column/
+      ]
+    end
 
-          model.sha_attribute(:no_name)
-        end
+    with_them do
+      before do
+        stub_rails_env(environment)
       end
 
-      context 'when other execeptions are raised' do
-        it 'logs and re-rasises the error' do
-          allow(model).to receive(:table_exists?).and_raise(ActiveRecord::NoDatabaseError.new('does not exist'))
-
-          expect(model).not_to receive(:columns)
-          expect(model).not_to receive(:attribute)
-          expect(Gitlab::AppLogger).to receive(:error)
-
-          expect { model.sha_attribute(:name) }.to raise_error(ActiveRecord::NoDatabaseError)
+      context 'with sha_attribute' do
+        before do
+          model.sha_attribute(column_name)
         end
-      end
-    end
 
-    context 'when in production' do
-      before do
-        stub_rails_env('production')
+        it 'validates column type' do
+          if expected_error == :no_error
+            expect { load_schema! }.not_to raise_error
+          elsif expected_error == :sha_mismatch_error
+            expect { load_schema! }.to raise_error(
+              described_class::ShaAttributeTypeMismatchError,
+              /sha_attribute.*#{column_name}.* should be a :binary column/
+            )
+          end
+        end
       end
 
-      it 'defines a SHA attribute' do
-        expect(model).not_to receive(:table_exists?)
-        expect(model).not_to receive(:columns)
-        expect(model).to receive(:attribute).with(:sha1, an_instance_of(Gitlab::Database::ShaAttribute))
+      context 'with sha256_attribute' do
+        before do
+          model.sha256_attribute(column_name)
+        end
 
-        model.sha_attribute(:sha1)
+        it 'validates column type' do
+          if expected_error == :no_error
+            expect { load_schema! }.not_to raise_error
+          elsif expected_error == :sha_mismatch_error
+            expect { load_schema! }.to raise_error(
+              described_class::Sha256AttributeTypeMismatchError,
+              /sha256_attribute.*#{column_name}.* should be a :binary column/
+            )
+          end
+        end
       end
     end
   end