Add tier and distribution based on ee option

generator_templates/usage_metric_definition/metric_definition.yml

@@ -13,7 +13,5 @@ time_frame: <%= time_frame %>
 data_source:
 distribution:
 <%= distribution %>
-# tier:
-# - free
-# - premium
-# - ultimate
+tier:
+<%= tier %>

lib/generators/gitlab/usage_metric_definition_generator.rb

@@ -53,8 +53,19 @@ module Gitlab
     end
 
     def distribution
-      value = ['- ce']
-      value << '- ee' if ee?
+      value = []
+      value << '- ce' unless ee?
+      value << '- ee'
+      value.join("\n")
+    end
+
+    def tier
+      value = if ee?
+        ['#- premium']
+      else
+        ['- free', '- premium']
+      end
+      value << '- ultimate'
       value.join("\n")
     end
 

spec/fixtures/lib/generators/gitlab/usage_metric_definition_generator/sample_metric_with_ee.yml

+---
+# See Usage Ping metrics dictionary docs https://docs.gitlab.com/ee/development/usage_ping/metrics_dictionary.html
+key_path: counts_weekly.test_metric
+name: test metric name
+description:
+product_section:
+product_stage:
+product_group:
+product_category:
+value_type: number
+status: implemented
+milestone: "13.9"
+introduced_by_url:
+time_frame: 7d
+data_source:
+distribution:
+- ee
+tier:
+#- premium
+- ultimate

spec/fixtures/lib/generators/gitlab/usage_metric_definition_generator/sample_metric_with_name_suggestions.yml

@@ -15,8 +15,8 @@ time_frame: 7d
 data_source:
 distribution:
 - ce
-# Add here corresponding tiers
-# tier:
-# - free
-# - premium
-# - ultimate
+- ee
+tier:
+- free
+- premium
+- ultimate

spec/lib/generators/gitlab/usage_metric_definition_generator_spec.rb

@@ -20,20 +20,37 @@ RSpec.describe Gitlab::UsageMetricDefinitionGenerator do
   end
 
   describe 'Creating metric definition file' do
+    let(:sample_metric) { load_sample_metric_definition(filename: sample_filename) }
+
     # Stub version so that `milestone` key remains constant between releases to prevent flakiness.
     before do
       stub_const('Gitlab::VERSION', '13.9.0')
       allow(::Gitlab::Usage::Metrics::NamesSuggestions::Generator).to receive(:generate).and_return('test metric name')
     end
 
-    let(:sample_metric) { load_sample_metric_definition(filename: 'sample_metric_with_name_suggestions.yml') }
+    context 'without ee option' do
+      let(:sample_filename) { 'sample_metric_with_name_suggestions.yml' }
+      let(:metric_definition_path) { Dir.glob(File.join(temp_dir, 'metrics/counts_7d/*_test_metric.yml')).first }
 
-    it 'creates a metric definition file using the template' do
-      described_class.new([key_path], { 'dir' => dir }).invoke_all
+      it 'creates a metric definition file using the template' do
+        described_class.new([key_path], { 'dir' => dir}).invoke_all
+        expect(YAML.safe_load(File.read(metric_definition_path))).to eq(sample_metric)
+      end
+    end
 
-      metric_definition_path = Dir.glob(File.join(temp_dir, 'metrics/counts_7d/*_test_metric.yml')).first
+    context 'with ee is true' do
+      let(:sample_filename) { 'sample_metric_with_ee.yml' }
+      let(:metric_definition_path) { Dir.glob(File.join(temp_dir, 'ee/config/metrics/counts_7d/*_test_metric.yml')).first }
 
-      expect(YAML.safe_load(File.read(metric_definition_path))).to eq(sample_metric)
+      before do
+        stub_const("#{described_class}::TOP_LEVEL_DIR", 'config')
+        stub_const("#{described_class}::TOP_LEVEL_DIR_EE", File.join(temp_dir, 'ee'))
+      end
+
+      it 'creates a metric definition file using the template' do
+        described_class.new([key_path], { 'dir' => dir, 'ee': true }).invoke_all
+        expect(YAML.safe_load(File.read(metric_definition_path))).to eq(sample_metric)
+      end
     end
   end