Merge branch 'refactor_base_probe' into 'master'

Refactor Base Probe and introduce healh_check rake task

See merge request !164539



Merged-by: Siddharth Dungarwal <sdungarwal@gitlab.com>
Approved-by: Siddharth Dungarwal <sdungarwal@gitlab.com>
Reviewed-by: Matthias Käppler <mkaeppler@gitlab.com>
Co-authored-by: Nikola Milojevic <nmilojevic@gitlab.com>

ee/app/services/cloud_connector/status_checks/probes/base_probe.rb

 # frozen_string_literal: true
 
+require 'active_model'
+
 module CloudConnector
   module StatusChecks
     module Probes
       class BaseProbe
+        include ActiveModel::Validations
+        include ActiveModel::Validations::Callbacks
+
+        class Details
+          delegate :each, :[], :empty?, :to_hash, to: :@messages
+
+          def initialize
+            @messages = {}
+          end
+
+          def add(attribute, message)
+            @messages[attribute] = message
+          end
+        end
+
         def execute(**_context)
-          raise "#{self.class} must implement #execute"
+          return failure(failure_message) unless valid?
+
+          success(success_message)
         end
 
         private
 
+        def details
+          @details ||= Details.new
+        end
+
+        def probe_name
+          self.class.name.demodulize.underscore.to_sym
+        end
+
         def success(message)
-          ProbeResult.new(probe_name, true, message)
+          create_result(true, message)
         end
 
         def failure(message)
-          ProbeResult.new(probe_name, false, message)
+          create_result(false, message)
         end
 
-        def probe_name
-          self.class.name.demodulize.underscore.to_sym
+        def create_result(success, message)
+          ProbeResult.new(probe_name, success, message, details, errors)
+        end
+
+        def failure_message
+          errors.full_messages.first
+        end
+
+        def success_message
+          raise NotImplementedError, "#{self.class} must implement #success_message"
         end
       end
     end

ee/app/services/cloud_connector/status_checks/probes/probe_result.rb

@@ -4,12 +4,14 @@ module CloudConnector
   module StatusChecks
     module Probes
       class ProbeResult
-        attr_reader :name, :success, :message
+        attr_reader :name, :success, :message, :details, :errors
 
-        def initialize(name, success, message)
+        def initialize(name, success, message, details = [], errors = [])
           @name = name
           @success = success
           @message = message
+          @details = details
+          @errors = errors
         end
 
         def success?

ee/lib/tasks/cloud_connector/health_check.rake

+# frozen_string_literal: true
+
+namespace :cloud_connector do
+  desc 'GitLab | Cloud Connector | Health check'
+  task :health_check, [:username, :filename, :include_details] => :environment do |_, args|
+    user = Tasks::CloudConnector::TaskHelper.find_user(args.username)
+
+    probe_results = CloudConnector::StatusChecks::StatusService.new(user: user).execute[:probe_results]
+    include_details = Gitlab::Utils.to_boolean(args.include_details, default: true)
+    Tasks::CloudConnector::TaskHelper.process_probe_results(probe_results, include_details: include_details)
+
+    Tasks::CloudConnector::TaskHelper.save_report(args.filename, probe_results)
+  end
+end

ee/lib/tasks/cloud_connector/task_helper.rb

+# frozen_string_literal: true
+
+module Tasks
+  module CloudConnector
+    module TaskHelper
+      COLORS = {
+        success: :green,
+        failure: :red,
+        details: :black,
+        check: :blue,
+        warning: :yellow
+      }.freeze
+      USER_NOT_PROVIDED_MESSAGE = "User not provided"
+      SKIPPING_MESSAGE_TEXT = "\n• Skipping %{name} check: %{message}"
+      USAGE_TEXT = "\n  Usage: rake 'cloud_connector:health_check[username,report.json]'"
+      SKIPPING_CHECKS_TEXT = "\n  Please note that some checks might fail or be skipped, " \
+        "if a valid username is not provided."
+      PLEASE_PROVIDE_FILENAME_TEXT = "\n  If you want to save report to a file, " \
+        "please specify the filename when running the task. #{USAGE_TEXT}".freeze
+      PLEASE_PROVIDE_USER_TEXT = "Proceeding without a user... #{SKIPPING_CHECKS_TEXT} #{USAGE_TEXT}".freeze
+      OUTPUT_DIR = Rails.root.join('tmp/cloud_connector/reports')
+
+      class << self
+        def find_user(username)
+          unless username
+            log_warning("The username was not provided. #{PLEASE_PROVIDE_USER_TEXT}")
+
+            return
+          end
+
+          user = User.find_by_username(username)
+          log_warning("User '#{username}' not found. #{PLEASE_PROVIDE_USER_TEXT}") unless user
+
+          user
+        end
+
+        def save_report(filename, probe_results)
+          return log(PLEASE_PROVIDE_FILENAME_TEXT) unless filename
+
+          report_path = File.join(OUTPUT_DIR, File.basename(filename))
+
+          log("\n• Saving report to #{report_path}...", color: COLORS[:check])
+
+          begin
+            FileUtils.mkdir_p(OUTPUT_DIR)
+
+            File.open(report_path, 'w') do |file|
+              file.write(pretty_json(probe_results.as_json))
+            end
+
+            print_success("Report successfully written to #{report_path}")
+          rescue StandardError => e
+            print_failure("Failed to write report to #{report_path}: #{e.message}")
+          end
+        end
+
+        def process_probe_results(probe_results, include_details: false)
+          probe_results.each do |result|
+            log("\n• #{result.name.to_s.humanize} check...", color: COLORS[:check])
+
+            print_details(result.details) if include_details && result.details.present?
+
+            if result.success?
+              print_success(result.message)
+            else
+              process_errors(result)
+            end
+          end
+        end
+
+        def pretty_json(data)
+          ::Gitlab::Json.pretty_generate(data).gsub(/^/, '    ')
+        end
+
+        private
+
+        def process_errors(result)
+          return print_failure(result.message) unless result.errors.present?
+
+          result.errors.full_messages.each do |error|
+            if error.include?(USER_NOT_PROVIDED_MESSAGE)
+              log(skipping_message(result.name, error))
+            else
+              print_failure(error)
+            end
+          end
+        end
+
+        def skipping_message(name, message)
+          format(SKIPPING_MESSAGE_TEXT, name: name.to_s.humanize, message: message)
+        end
+
+        def print_success(message)
+          log("#{colored_message('✔ Success:', COLORS[:success])} #{message}")
+        end
+
+        def print_failure(message)
+          log("#{colored_message('✖ Failure:', COLORS[:failure])} #{message}")
+        end
+
+        def print_details(details)
+          log('  ◦ Details:', color: COLORS[:details])
+          log(pretty_json(details.as_json))
+        end
+
+        def log(message, color: nil)
+          $stdout.puts color ? colored_message(message, color) : message
+        end
+
+        def log_warning(message)
+          log("\n⚠ Warning: #{message}", color: COLORS[:warning])
+        end
+
+        def colored_message(message, color)
+          Rainbow.new.wrap(message).color(color).bright
+        end
+      end
+    end
+  end
+end

ee/spec/services/cloud_connector/status_checks/probes/base_probe_spec.rb

 # frozen_string_literal: true
 
 require 'spec_helper'
+require_relative 'test_probe'
 
 RSpec.describe CloudConnector::StatusChecks::Probes::BaseProbe, feature_category: :cloud_connector do
-  subject(:test_probe) { test_probe_class.new }
+  subject(:test_probe) { test_probe_class.new(**params) }
 
-  before do
-    stub_const('TestProbe', test_probe_class)
-  end
+  let(:params) { { success: true } }
+  let(:test_probe_class) { CloudConnector::StatusChecks::Probes::TestProbe }
 
   describe '#execute' do
-    context 'when not implemented in subclass' do
+    context 'when #success_message is not implemented in subclass' do
+      let(:params) { {} }
       let(:test_probe_class) { Class.new(described_class) }
 
-      it 'raises an error' do
-        expect { test_probe.execute }.to raise_error(RuntimeError, "TestProbe must implement #execute")
-      end
-    end
-
-    context 'when implemented in subclass' do
-      let(:test_probe_class) do
-        Class.new(described_class) do
-          def execute(*); end
-        end
+      before do
+        stub_const('TestProbe', test_probe_class)
       end
 
-      it 'does not raise an error' do
-        expect { test_probe.execute }.not_to raise_error
+      it 'raises a NotImplementedError when success_message is not implemented' do
+        expect { test_probe.execute }.to raise_error(NotImplementedError, "TestProbe must implement #success_message")
       end
     end
-  end
 
-  describe '#success' do
-    let(:test_probe_class) do
-      Class.new(described_class) do
-        def execute(*)
-          success('Test success message')
+    context 'when #success_message is implemented in subclass' do
+      context 'when validation passes' do
+        it 'returns a successful ProbeResult' do
+          result = test_probe.execute
+
+          expect(result).to be_a(CloudConnector::StatusChecks::Probes::ProbeResult)
+          expect(result.success?).to be true
+          expect(result.message).to eq('OK')
+          expect(result.name).to eq(:test_probe)
+          expect(result.errors).to be_empty
+          expect(result.details).to include(test: 'true')
         end
       end
-    end
 
-    it 'returns a successful ProbeResult' do
-      result = test_probe.execute
+      context 'when validation fails' do
+        let(:params) { { success: false } }
 
-      expect(result).to be_a(CloudConnector::StatusChecks::Probes::ProbeResult)
-      expect(result.success).to be true
-      expect(result.name).to eq(:test_probe)
-    end
-  end
+        it 'returns a failed ProbeResult with validation errors' do
+          result = test_probe.execute
 
-  describe '#failure' do
-    let(:test_probe_class) do
-      Class.new(described_class) do
-        def execute(*)
-          failure('Test failure message')
+          expect(result).to be_a(CloudConnector::StatusChecks::Probes::ProbeResult)
+          expect(result.success?).to be false
+          expect(result.message).to eq('NOK')
+          expect(result.name).to eq(:test_probe)
+          expect(result.errors.full_messages).to match_array(%w[NOK])
+          expect(result.details).to include(test: 'true')
         end
       end
     end
-
-    it 'returns a failed ProbeResult' do
-      result = test_probe.execute
-
-      expect(result).to be_a(CloudConnector::StatusChecks::Probes::ProbeResult)
-      expect(result.success).to be false
-      expect(result.name).to eq(:test_probe)
-    end
   end
 end

ee/spec/services/cloud_connector/status_checks/probes/probe_result_spec.rb

@@ -7,7 +7,9 @@
   let(:name) { 'Test Probe' }
   let(:success) { true }
   let(:message) { 'Probe successful' }
-  let(:probe_result) { described_class.new(name, success, message) }
+  let(:details) { { key: 'value' } }
+  let(:errors) { ['An error occurred'] }
+  let(:probe_result) { described_class.new(name, success, message, details, errors) }
 
   describe '#success?' do
     context 'when success is true' do
@@ -45,5 +47,27 @@
     it 'allows reading of message attribute' do
       expect(probe_result.message).to eq(message)
     end
+
+    it 'allows reading of details attribute' do
+      expect(probe_result.details).to eq(details)
+    end
+
+    it 'allows reading of errors attribute' do
+      expect(probe_result.errors).to eq(errors)
+    end
+  end
+
+  describe '#initialize' do
+    context 'when details and errors are not provided' do
+      let(:probe_result) { described_class.new(name, success, message) }
+
+      it 'defaults details to an empty array' do
+        expect(probe_result.details).to eq([])
+      end
+
+      it 'defaults errors to an empty array' do
+        expect(probe_result.errors).to eq([])
+      end
+    end
   end
 end

ee/spec/services/cloud_connector/status_checks/probes/test_probe.rb

@@ -5,14 +5,28 @@ module StatusChecks
     module Probes
       # Returns a canned response, useful for unit testing.
       class TestProbe < BaseProbe
+        extend ::Gitlab::Utils::Override
+
+        validate :check_success
+        after_validation :collect_details
+
         def initialize(success: true)
           @success = success
         end
 
-        def execute(*)
-          return failure('NOK') unless @success
+        private
+
+        def check_success
+          errors.add(:base, 'NOK') unless @success
+        end
+
+        def collect_details
+          details.add(:test, 'true')
+        end
 
-          success('OK')
+        override :success_message
+        def success_message
+          'OK'
         end
       end
     end

ee/spec/tasks/cloud_connector/health_check_rake_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require_relative '../../services/cloud_connector/status_checks/probes/test_probe'
+
+RSpec.describe 'cloud_connector:health_check', :silence_stdout, feature_category: :cloud_connector do
+  let(:test_probe) { CloudConnector::StatusChecks::Probes::TestProbe.new(success: success) }
+  let(:filename) { 'output.txt' }
+  let(:filepath) { File.join(Tasks::CloudConnector::TaskHelper::OUTPUT_DIR, filename) }
+  let(:user) { create(:user, username: 'test_user') }
+  let(:success) { true }
+
+  before do
+    Rake.application.rake_require('ee/lib/tasks/cloud_connector/health_check', [Rails.root.to_s])
+    stub_const('CloudConnector::StatusChecks::StatusService::DEFAULT_PROBES', [test_probe])
+  end
+
+  describe 'health check execution' do
+    it 'executes the health check with TestProbe' do
+      expect { run_rake_task('cloud_connector:health_check') }.to output(/Success: OK/).to_stdout
+    end
+
+    context 'when the probe fails' do
+      let(:success) { false }
+
+      it 'prints failure messages' do
+        expect { run_rake_task('cloud_connector:health_check') }.to output(/Failure: NOK/).to_stdout
+      end
+    end
+  end
+
+  describe 'handling parameters' do
+    context 'when include_details is true' do
+      it 'prints probe details' do
+        expect { run_rake_task('cloud_connector:health_check', [nil, nil, 'true']) }
+          .to output(/"test": "true"/).to_stdout
+      end
+    end
+
+    context 'when include_details is not provided' do
+      it 'defaults to printing probe details' do
+        expect { run_rake_task('cloud_connector:health_check', [nil, nil, nil]) }
+          .to output(/"test": "true"/).to_stdout
+      end
+    end
+
+    context 'when filename is provided' do
+      it 'saves report to a file' do
+        expect(File).to receive(:open).with(filepath, 'w')
+
+        expect { run_rake_task('cloud_connector:health_check', [nil, filename]) }
+          .to output(/Saving report to #{filepath}/).to_stdout
+      end
+    end
+
+    context 'when filename is not provided' do
+      it 'does not attempt to save the report to a file' do
+        expect(File).not_to receive(:open)
+
+        expect { run_rake_task('cloud_connector:health_check', [nil, nil]) }
+          .to output(/✔ Success: OK/).to_stdout
+      end
+    end
+
+    context 'when a username is provided' do
+      it 'loads the user and uses it in the health check' do
+        expect(User).to receive(:find_by_username).with('test_user').and_call_original
+        expect(CloudConnector::StatusChecks::StatusService).to receive(:new).with(user: user).and_call_original
+
+        expect { run_rake_task('cloud_connector:health_check', ['test_user', nil]) }
+          .to output(/✔ Success:/).to_stdout
+      end
+
+      it 'prints a warning if the user is not found and proceeds without the user' do
+        expect { run_rake_task('cloud_connector:health_check', ['unknown_user', nil]) }
+          .to output(/Warning: User 'unknown_user' not found. Proceeding without a user.../).to_stdout
+      end
+    end
+
+    context 'when a username is not provided' do
+      it 'executes the health check without a user' do
+        expect(User).not_to receive(:find_by_username)
+
+        expect { run_rake_task('cloud_connector:health_check', [nil, nil]) }
+          .to output(/✔ Success: OK/).to_stdout
+      end
+    end
+  end
+
+  describe 'error handling' do
+    it 'handles file write errors gracefully' do
+      allow(File).to receive(:open).and_raise(StandardError.new('disk full'))
+
+      expect { run_rake_task('cloud_connector:health_check', [nil, filename]) }
+        .to output(/Failed to write report to #{filepath}: disk full/).to_stdout
+    end
+  end
+end

ee/spec/tasks/cloud_connector/task_helper_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require 'rainbow'
+
+RSpec.describe Tasks::CloudConnector::TaskHelper, :silence_stdout, feature_category: :cloud_connector do
+  let(:filename) { 'output.txt' }
+  let(:filepath) { File.join(described_class::OUTPUT_DIR, filename) }
+  let!(:user) { create(:user, username: 'test_user') }
+  let(:success) { true }
+  let(:message) { 'OK' }
+  let(:details) { { test: 'true' } }
+  let(:errors) { [] }
+  let(:probe_results) do
+    [
+      instance_double(CloudConnector::StatusChecks::Probes::ProbeResult,
+        name: 'test_probe',
+        success?: success,
+        message: message,
+        details: details,
+        errors: errors)
+    ]
+  end
+
+  describe '.find_user' do
+    it 'finds the user by username' do
+      expect(User).to receive(:find_by_username).with('test_user').and_call_original
+
+      expect(described_class.find_user('test_user')).to eq(user)
+    end
+
+    it 'prints a warning and returns nil if the username is not provided' do
+      expect { described_class.find_user(nil) }
+        .to output(/Warning: The username was not provided. Proceeding without a user/).to_stdout
+    end
+
+    it 'prints a warning and returns nil if the user is not found' do
+      allow(User).to receive(:find_by_username).with('unknown_user').and_call_original
+
+      expect { described_class.find_user('unknown_user') }
+        .to output(/Warning: User 'unknown_user' not found. Proceeding without a user/).to_stdout
+    end
+  end
+
+  describe '.save_report' do
+    it 'prints a warning if filename is not provided' do
+      expect { described_class.save_report(nil, probe_results) }
+        .to output(/If you want to save report to a file/).to_stdout
+    end
+
+    it 'saves the report to a file' do
+      expect(File).to receive(:open).with(filepath, 'w').and_yield(StringIO.new)
+
+      expect { described_class.save_report(filename, probe_results) }
+        .to output(/Saving report to #{filepath}/).to_stdout
+    end
+
+    it 'handles file write errors gracefully' do
+      allow(File).to receive(:open).and_raise(StandardError.new('disk full'))
+
+      expect { described_class.save_report(filename, probe_results) }
+        .to output(/Failed to write report to #{filepath}: disk full/).to_stdout
+    end
+  end
+
+  describe '.process_probe_results' do
+    context 'when probe succeeds' do
+      it 'prints success message' do
+        expect { described_class.process_probe_results(probe_results) }
+          .to output(/Success: OK/).to_stdout
+      end
+    end
+
+    context 'when probe fails' do
+      let(:success) { false }
+      let(:message) { 'NOK' }
+      let(:errors) { ActiveModel::Errors.new(nil) }
+
+      it 'prints failure messages' do
+        errors.add(:base, 'Something went wrong')
+
+        expect { described_class.process_probe_results(probe_results) }
+          .to output(/Failure: Something went wrong/).to_stdout
+      end
+
+      context 'when no errors' do
+        it 'prints failure messages' do
+          expect { described_class.process_probe_results(probe_results) }
+            .to output(/Failure: NOK/).to_stdout
+        end
+      end
+
+      context 'when the error is "User is not provided"' do
+        let(:success) { false }
+        let(:message) { 'NOK' }
+
+        it 'prints skipping message' do
+          errors.add(:base, described_class::USER_NOT_PROVIDED_MESSAGE)
+
+          expect { described_class.process_probe_results(probe_results) }
+            .to output(/Skipping Test probe check: User not provided/).to_stdout
+        end
+      end
+    end
+
+    context 'when include_details is true' do
+      it 'prints probe details' do
+        expect { described_class.process_probe_results(probe_results, include_details: true) }
+          .to output(/"test": "true"/).to_stdout
+      end
+    end
+
+    context 'when include_details is false' do
+      it 'does not print probe details' do
+        expect { described_class.process_probe_results(probe_results, include_details: false) }
+          .not_to output(/"test": "true"/).to_stdout
+      end
+    end
+  end
+end

lib/tasks/gitlab/artifacts/migrate.rake

@@ -3,10 +3,10 @@
 desc 'GitLab | Artifacts | Migrate files for artifacts to comply with new storage format'
 namespace :gitlab do
   require 'logger'
-  require 'resolv-replace'
 
   namespace :artifacts do
     task migrate: :environment do
+      require 'resolv-replace'
       logger = Logger.new($stdout)
 
       helper = Gitlab::LocalAndRemoteStorageMigration::ArtifactMigrater.new(logger)
@@ -19,6 +19,7 @@ namespace :gitlab do
     end
 
     task migrate_to_local: :environment do
+      require 'resolv-replace'
       logger = Logger.new($stdout)
 
       helper = Gitlab::LocalAndRemoteStorageMigration::ArtifactMigrater.new(logger)