Fix prometheus client tests

2 files changed, 21 insertions, 10 deletions

diff --git a/lib/gitlab/prometheus_client.rb b/lib/gitlab/prometheus_client.rb
index d1875dd1042..79f66b42c21 100644
--- a/lib/gitlab/prometheus_client.rb
+++ b/lib/gitlab/prometheus_client.rb
@@ -50,10 +50,12 @@ module Gitlab
       response = rest_client[path].get(params: args)
       handle_response(response)
     rescue SocketError
-      raise PrometheusError, "Can't connect to #{url}"
+      raise PrometheusError, "Can't connect to #{rest_client.url}"
     rescue OpenSSL::SSL::SSLError
-      raise PrometheusError, "#{url} contains invalid SSL data"
-    rescue HTTParty::Error
+      raise PrometheusError, "#{rest_client.url} contains invalid SSL data"
+    rescue RestClient::ExceptionWithResponse => ex
+      handle_exception_response(ex.response)
+    rescue RestClient::Exception
       raise PrometheusError, "Network connection error"
     end
 
@@ -61,7 +63,14 @@ module Gitlab
       json_data = JSON.parse(response.body)
       if response.code == 200 && json_data['status'] == 'success'
         json_data['data'] || {}
-      elsif response.code == 400
+      else
+        raise PrometheusError, "#{response.code} - #{response.body}"
+      end
+    end
+
+    def handle_exception_response(response)
+      if response.code == 400
+        json_data = JSON.parse(response.body)
         raise PrometheusError, json_data['error'] || 'Bad data received'
       else
         raise PrometheusError, "#{response.code} - #{response.body}"
diff --git a/spec/lib/gitlab/prometheus_client_spec.rb b/spec/lib/gitlab/prometheus_client_spec.rb
index de625324092..575b7805e33 100644
--- a/spec/lib/gitlab/prometheus_client_spec.rb
+++ b/spec/lib/gitlab/prometheus_client_spec.rb
@@ -3,7 +3,7 @@ require 'spec_helper'
 describe Gitlab::PrometheusClient do
   include PrometheusHelpers
 
-  subject { described_class.new(api_url: 'https://prometheus.example.com') }
+  subject { described_class.new(RestClient::Resource.new('https://prometheus.example.com')) }
 
   describe '#ping' do
     it 'issues a "query" request to the API endpoint' do
@@ -52,11 +52,13 @@ describe Gitlab::PrometheusClient do
   describe 'failure to reach a provided prometheus url' do
     let(:prometheus_url) {"https://prometheus.invalid.example.com"}
 
+    subject { described_class.new(RestClient::Resource.new(prometheus_url)) }
+
     context 'exceptions are raised' do
       it 'raises a Gitlab::PrometheusError error when a SocketError is rescued' do
         req_stub = stub_prometheus_request_with_exception(prometheus_url, SocketError)
 
-        expect { subject.send(:get, prometheus_url) }
+        expect { subject.send(:get, '/', {}) }
           .to raise_error(Gitlab::PrometheusError, "Can't connect to #{prometheus_url}")
         expect(req_stub).to have_been_requested
       end
@@ -64,15 +66,15 @@ describe Gitlab::PrometheusClient do
       it 'raises a Gitlab::PrometheusError error when a SSLError is rescued' do
         req_stub = stub_prometheus_request_with_exception(prometheus_url, OpenSSL::SSL::SSLError)
 
-        expect { subject.send(:get, prometheus_url) }
+        expect { subject.send(:get, '/', {}) }
           .to raise_error(Gitlab::PrometheusError, "#{prometheus_url} contains invalid SSL data")
         expect(req_stub).to have_been_requested
       end
 
-      it 'raises a Gitlab::PrometheusError error when a HTTParty::Error is rescued' do
-        req_stub = stub_prometheus_request_with_exception(prometheus_url, HTTParty::Error)
+      it 'raises a Gitlab::PrometheusError error when a RestClient::Exception is rescued' do
+        req_stub = stub_prometheus_request_with_exception(prometheus_url, RestClient::Exception)
 
-        expect { subject.send(:get, prometheus_url) }
+        expect { subject.send(:get, '/', {}) }
           .to raise_error(Gitlab::PrometheusError, "Network connection error")
         expect(req_stub).to have_been_requested
       end