1 files changed, 33 insertions, 5 deletions

diff --git a/spec/lib/generators/gitlab/usage_metric_generator_spec.rb b/spec/lib/generators/gitlab/usage_metric_generator_spec.rb
index f38815acca6..207ecb88aad 100644
--- a/spec/lib/generators/gitlab/usage_metric_generator_spec.rb
+++ b/spec/lib/generators/gitlab/usage_metric_generator_spec.rb
@@ -8,7 +8,7 @@ RSpec.describe Gitlab::UsageMetricGenerator, :silence_stdout do
   let(:spec_ce_temp_dir) { Dir.mktmpdir }
   let(:spec_ee_temp_dir) { Dir.mktmpdir }
   let(:args) { ['CountFoo'] }
-  let(:options) { { 'type' => 'redis_hll' } }
+  let(:options) { { 'type' => 'generic' } }
 
   before do
     stub_const("#{described_class}::CE_DIR", ce_temp_dir)
@@ -30,27 +30,39 @@ RSpec.describe Gitlab::UsageMetricGenerator, :silence_stdout do
 
   describe 'Creating metric instrumentation files' do
     let(:sample_metric_dir) { 'lib/generators/gitlab/usage_metric_generator' }
-    let(:sample_metric) { fixture_file(File.join(sample_metric_dir, 'sample_metric.rb')) }
+    let(:generic_sample_metric) { fixture_file(File.join(sample_metric_dir, 'sample_generic_metric.rb')) }
+    let(:database_sample_metric) { fixture_file(File.join(sample_metric_dir, 'sample_database_metric.rb')) }
     let(:sample_spec) { fixture_file(File.join(sample_metric_dir, 'sample_metric_test.rb')) }
 
     it 'creates CE metric instrumentation files using the template' do
       described_class.new(args, options).invoke_all
 
-      expect_generated_file(ce_temp_dir, 'count_foo_metric.rb', sample_metric)
+      expect_generated_file(ce_temp_dir, 'count_foo_metric.rb', generic_sample_metric)
       expect_generated_file(spec_ce_temp_dir, 'count_foo_metric_spec.rb', sample_spec)
     end
 
     context 'with EE flag true' do
-      let(:options) { { 'type' => 'redis_hll', 'ee' => true } }
+      let(:options) { { 'type' => 'generic', 'ee' => true } }
 
       it 'creates EE metric instrumentation files using the template' do
         described_class.new(args, options).invoke_all
 
-        expect_generated_file(ee_temp_dir, 'count_foo_metric.rb', sample_metric)
+        expect_generated_file(ee_temp_dir, 'count_foo_metric.rb', generic_sample_metric)
         expect_generated_file(spec_ee_temp_dir, 'count_foo_metric_spec.rb', sample_spec)
       end
     end
 
+    context 'for database type' do
+      let(:options) { { 'type' => 'database', 'operation' => 'count' } }
+
+      it 'creates the metric instrumentation file using the template' do
+        described_class.new(args, options).invoke_all
+
+        expect_generated_file(ce_temp_dir, 'count_foo_metric.rb', database_sample_metric)
+        expect_generated_file(spec_ce_temp_dir, 'count_foo_metric_spec.rb', sample_spec)
+      end
+    end
+
     context 'with type option missing' do
       let(:options) { {} }
 
@@ -66,5 +78,21 @@ RSpec.describe Gitlab::UsageMetricGenerator, :silence_stdout do
         expect { described_class.new(args, options).invoke_all }.to raise_error(ArgumentError, /Unknown type 'some_other_type'/)
       end
     end
+
+    context 'without operation for database metric' do
+      let(:options) { { 'type' => 'database' } }
+
+      it 'raises an ArgumentError' do
+        expect { described_class.new(args, options).invoke_all }.to raise_error(ArgumentError, /Unknown operation ''/)
+      end
+    end
+
+    context 'with wrong operation for database metric' do
+      let(:options) { { 'type' => 'database', 'operation' => 'sleep' } }
+
+      it 'raises an ArgumentError' do
+        expect { described_class.new(args, options).invoke_all }.to raise_error(ArgumentError, /Unknown operation 'sleep'/)
+      end
+    end
   end
 end