2 files changed, 116 insertions, 0 deletions
diff --git a/spec/lib/gitlab/health_checks/middleware_spec.rb b/spec/lib/gitlab/health_checks/middleware_spec.rb
new file mode 100644
index 00000000000..3b644539acc
--- /dev/null
+++ b/spec/lib/gitlab/health_checks/middleware_spec.rb
@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+
+RSpec.describe Gitlab::HealthChecks::Middleware do
+  let(:app) { instance_double(Proc) }
+  let(:env) { { 'PATH_INFO' => path } }
+
+  let(:readiness_probe) { instance_double(Gitlab::HealthChecks::Probes::Collection) }
+  let(:liveness_probe) { instance_double(Gitlab::HealthChecks::Probes::Collection) }
+  let(:probe_result) { Gitlab::HealthChecks::Probes::Status.new(200, { status: 'ok' }) }
+
+  subject(:middleware) { described_class.new(app, readiness_probe, liveness_probe) }
+
+  describe '#call' do
+    context 'handling /readiness requests' do
+      let(:path) { '/readiness' }
+
+      it 'handles the request' do
+        expect(readiness_probe).to receive(:execute).and_return(probe_result)
+
+        response = middleware.call(env)
+
+        expect(response).to eq([200, { 'Content-Type' => 'application/json; charset=utf-8' }, ['{"status":"ok"}']])
+      end
+    end
+
+    context 'handling /liveness requests' do
+      let(:path) { '/liveness' }
+
+      it 'handles the request' do
+        expect(liveness_probe).to receive(:execute).and_return(probe_result)
+
+        response = middleware.call(env)
+
+        expect(response).to eq([200, { 'Content-Type' => 'application/json; charset=utf-8' }, ['{"status":"ok"}']])
+      end
+    end
+
+    context 'handling other requests' do
+      let(:path) { '/other_path' }
+
+      it 'forwards them to the next middleware' do
+        expect(app).to receive(:call).with(env).and_return([201, {}, []])
+
+        response = middleware.call(env)
+
+        expect(response).to eq([201, {}, []])
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/health_checks/server_spec.rb b/spec/lib/gitlab/health_checks/server_spec.rb
new file mode 100644
index 00000000000..65d24acbf22
--- /dev/null
+++ b/spec/lib/gitlab/health_checks/server_spec.rb
@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::HealthChecks::Server do
+  context 'with running server thread' do
+    subject(:server) { described_class.new(address: 'localhost', port: 8082) }
+
+    before do
+      # We need to send a request to localhost
+      WebMock.allow_net_connect!
+
+      server.start
+    end
+
+    after do
+      webmock_enable!
+
+      server.stop
+    end
+
+    shared_examples 'serves health check at' do |path|
+      it 'responds with 200 OK' do
+        response = Gitlab::HTTP.try_get("http://localhost:8082/#{path}", allow_local_requests: true)
+
+        expect(response.code).to eq(200)
+      end
+    end
+
+    describe '/readiness' do
+      it_behaves_like 'serves health check at', 'readiness'
+    end
+
+    describe '/liveness' do
+      it_behaves_like 'serves health check at', 'liveness'
+    end
+
+    describe 'other routes' do
+      it 'serves 404' do
+        response = Gitlab::HTTP.try_get("http://localhost:8082/other", allow_local_requests: true)
+
+        expect(response.code).to eq(404)
+      end
+    end
+  end
+
+  context 'when server thread goes away' do
+    before do
+      expect_next_instance_of(::WEBrick::HTTPServer) do |webrick|
+        allow(webrick).to receive(:start)
+        expect(webrick).to receive(:listeners).and_call_original
+      end
+    end
+
+    specify 'stop closes TCP socket' do
+      server = described_class.new(address: 'localhost', port: 8082)
+      server.start
+
+      expect(server.thread).to receive(:alive?).and_return(false).at_least(:once)
+
+      server.stop
+    end
+  end
+end