1 files changed, 28 insertions, 2 deletions
diff --git a/spec/lib/gitlab/prometheus/queries/additional_metrics_environment_query_spec.rb b/spec/lib/gitlab/prometheus/queries/additional_metrics_environment_query_spec.rb
index 5a88b23aa82..a6589f0c0a3 100644
--- a/spec/lib/gitlab/prometheus/queries/additional_metrics_environment_query_spec.rb
+++ b/spec/lib/gitlab/prometheus/queries/additional_metrics_environment_query_spec.rb
@@ -9,9 +9,35 @@ describe Gitlab::Prometheus::Queries::AdditionalMetricsEnvironmentQuery do
     let(:query_params) { [environment.id] }
 
     it 'queries using specific time' do
-      expect(client).to receive(:query_range).with(anything, start: 8.hours.ago.to_f, stop: Time.now.to_f)
-
+      expect(client).to receive(:query_range)
+        .with(anything, start: 8.hours.ago.to_f, stop: Time.now.to_f)
       expect(query_result).not_to be_nil
     end
+
+    context 'when start and end time parameters are provided' do
+      let(:query_params) { [environment.id, start_time, end_time] }
+
+      context 'as unix timestamps' do
+        let(:start_time) { 4.hours.ago.to_f }
+        let(:end_time) { 2.hours.ago.to_f }
+
+        it 'queries using the provided times' do
+          expect(client).to receive(:query_range)
+            .with(anything, start: start_time, stop: end_time)
+          expect(query_result).not_to be_nil
+        end
+      end
+
+      context 'as Date/Time objects' do
+        let(:start_time) { 4.hours.ago }
+        let(:end_time) { 2.hours.ago }
+
+        it 'queries using the provided times converted to unix' do
+          expect(client).to receive(:query_range)
+            .with(anything, start: start_time.to_f, stop: end_time.to_f)
+          expect(query_result).not_to be_nil
+        end
+      end
+    end
   end
 end