Add latest changes from gitlab-org/gitlab@master

1 files changed, 27 insertions, 0 deletions

diff --git a/spec/lib/generators/gitlab/usage_metric_definition/redis_hll_generator_spec.rb b/spec/lib/generators/gitlab/usage_metric_definition/redis_hll_generator_spec.rb
index 021fb8f5f58..95a577e6334 100644
--- a/spec/lib/generators/gitlab/usage_metric_definition/redis_hll_generator_spec.rb
+++ b/spec/lib/generators/gitlab/usage_metric_definition/redis_hll_generator_spec.rb
@@ -18,6 +18,10 @@ RSpec.describe Gitlab::UsageMetricDefinition::RedisHllGenerator do
     stub_prometheus_queries
   end
 
+  after do
+    FileUtils.rm_rf(temp_dir)
+  end
+
   it 'creates metric definition files' do
     described_class.new(args).invoke_all
 
@@ -27,4 +31,27 @@ RSpec.describe Gitlab::UsageMetricDefinition::RedisHllGenerator do
     expect(YAML.safe_load(File.read(weekly_metric_definition_path))).to include("key_path" => "redis_hll_counters.test_category.i_test_event_weekly")
     expect(YAML.safe_load(File.read(monthly_metric_definition_path))).to include("key_path" => "redis_hll_counters.test_category.i_test_event_monthly")
   end
+
+  context 'with ee option' do
+    let(:weekly_metric_definition_path) { Dir.glob(File.join(temp_dir, 'ee/config/metrics/counts_7d/*i_test_event_weekly.yml')).first }
+    let(:monthly_metric_definition_path) { Dir.glob(File.join(temp_dir, 'ee/config/metrics/counts_28d/*i_test_event_monthly.yml')).first }
+
+    let(:weekly_metric_definition) { YAML.safe_load(File.read(weekly_metric_definition_path)) }
+    let(:monthly_metric_definition) { YAML.safe_load(File.read(monthly_metric_definition_path)) }
+
+    before do
+      stub_const("#{Gitlab::UsageMetricDefinitionGenerator}::TOP_LEVEL_DIR", 'config')
+      stub_const("#{Gitlab::UsageMetricDefinitionGenerator}::TOP_LEVEL_DIR_EE", File.join(temp_dir, 'ee'))
+    end
+
+    it 'creates metric definition files' do
+      described_class.new(args, { 'ee': true }).invoke_all
+
+      expect(weekly_metric_definition).to include("key_path" => "redis_hll_counters.test_category.i_test_event_weekly")
+      expect(weekly_metric_definition["distribution"]).to include('ee')
+
+      expect(monthly_metric_definition).to include("key_path" => "redis_hll_counters.test_category.i_test_event_monthly")
+      expect(monthly_metric_definition["distribution"]).to include('ee')
+    end
+  end
 end