Improve RedisHllGenerator to accept --ee option

changelogs/unreleased/327870-improve-metric-yaml-defintion-generator-to-fill-correct-tier-and-d.yml

+---
+title: Add --ee option to Usage Metric Definition generator to fill correct tier and distribution
+merge_request: 59942
+author:
+type: other

lib/generators/gitlab/usage_metric_definition/redis_hll_generator.rb

@@ -10,14 +10,24 @@ module Gitlab
 
       argument :category, type: :string, desc: "Category name"
       argument :event, type: :string, desc: "Event name"
+      class_option :ee, type: :boolean, optional: true, default: false, desc: 'Indicates if metric is for ee'
 
       def create_metrics
-        Gitlab::UsageMetricDefinitionGenerator.start(["#{key_path}_weekly", '--dir', '7d'])
-        Gitlab::UsageMetricDefinitionGenerator.start(["#{key_path}_monthly", '--dir', '28d'])
+        weekly_params = ["#{key_path}_weekly", '--dir', '7d']
+        weekly_params << '--ee' if ee?
+        Gitlab::UsageMetricDefinitionGenerator.start(weekly_params)
+
+        monthly_params = ["#{key_path}_monthly", '--dir', '28d']
+        monthly_params << '--ee' if ee?
+        Gitlab::UsageMetricDefinitionGenerator.start(monthly_params)
       end
 
       private
 
+      def ee?
+        options[:ee]
+      end
+
       def key_path
         "redis_hll_counters.#{category}.#{event}"
       end

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