Add PrometheusQuery class and specs

Adds a new PrometheusQuery model which references a PrometheusMetric.
This is needed to support multiple queries per chart.

4 files changed, 79 insertions, 0 deletions

diff --git a/app/models/prometheus_metric.rb b/app/models/prometheus_metric.rb
index 5594594a48d..22b64744fe7 100644
--- a/app/models/prometheus_metric.rb
+++ b/app/models/prometheus_metric.rb
@@ -2,6 +2,7 @@
 
 class PrometheusMetric < ActiveRecord::Base
   belongs_to :project, validate: true, inverse_of: :prometheus_metrics
+  has_many :prometheus_queries
 
   enum group: {
     # built-in groups
diff --git a/app/models/prometheus_query.rb b/app/models/prometheus_query.rb
new file mode 100644
index 00000000000..9eb42d7126d
--- /dev/null
+++ b/app/models/prometheus_query.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+class PrometheusQuery < ActiveRecord::Base
+  belongs_to :prometheus_metric, validate: true
+
+  validates :query, presence: true
+  validates :unit, presence: true
+
+  def to_query_hash
+    {
+      query_range: query,
+      unit: unit,
+      label: legend,
+      series: query_series
+    }.compact
+  end
+
+  private
+
+  def query_series
+    case legend
+    when 'Status Code'
+      [{
+        label: 'status_code',
+        when: [
+          { value: '2xx', color: 'green' },
+          { value: '4xx', color: 'orange' },
+          { value: '5xx', color: 'red' }
+        ]
+      }]
+    end
+  end
+end
diff --git a/spec/factories/prometheus_queries.rb b/spec/factories/prometheus_queries.rb
new file mode 100644
index 00000000000..5b0398c3e15
--- /dev/null
+++ b/spec/factories/prometheus_queries.rb
@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+FactoryBot.define do
+  factory :prometheus_query, class: PrometheusQuery do
+    query 'avg(metric)'
+    unit 'm/s'
+    legend 'legend'
+    association :prometheus_metric
+  end
+end
diff --git a/spec/models/prometheus_query_spec.rb b/spec/models/prometheus_query_spec.rb
new file mode 100644
index 00000000000..db2645e15ad
--- /dev/null
+++ b/spec/models/prometheus_query_spec.rb
@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe PrometheusQuery do
+  subject { build(:prometheus_query) }
+
+  it { is_expected.to belong_to(:prometheus_metric) }
+  it { is_expected.to validate_presence_of(:query) }
+  it { is_expected.to validate_presence_of(:unit) }
+
+  describe '#to_query_hash' do
+    let(:expected_hash) do
+      {
+        query_range: subject.query,
+        unit: subject.unit,
+        label: subject.legend
+      }
+    end
+
+    specify { expect(subject.to_query_hash).to include expected_hash }
+  end
+
+  describe '#query_series' do
+    specify { expect(subject.__send__(:query_series)).to be_nil }
+
+    context 'when a status code is specified' do
+      before do
+        subject.legend = 'Status Code'
+      end
+
+      specify { expect(subject.__send__(:query_series)).to be_a Array }
+    end
+  end
+end