Merge branch 'mk-improve-import-rake-task-for-evals' into 'master'

Refactor rake task to check status of the project import

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/180431



Merged-by: Manoj M J <mmj@gitlab.com>
Approved-by: Manoj M J <mmj@gitlab.com>
Reviewed-by: Manoj M J <mmj@gitlab.com>
Co-authored-by: Małgorzata Ksionek <mksionek@gitlab.com>

ee/lib/gitlab/duo/developments/setup_groups_for_model_evaluation.rb

@@ -15,11 +15,13 @@ class SetupGroupsForModelEvaluation
         DOWNLOAD_URL = 'https://gitlab.com/gitlab-org/ai-powered/datasets/-/package_files/135727282/download'
         GROUP_IMPORT_URL = '/api/v4/groups/import'
         PROJECT_IMPORT_URL = '/api/v4/projects/import'
+        TIME_LIMIT = 180
 
         def initialize(group)
           @main_group = group
           @current_user = User.find_by(username: 'root') # rubocop:disable CodeReuse/ActiveRecord -- we need admin user
           @errors = []
+          @project_ids = []
         end
 
         def execute
@@ -30,6 +32,7 @@ def execute
           download_and_unpack_file
           create_subgroups
           create_subprojects
+          check_import_status
           delete_temporary_directory!
           clean_up_token!
 
@@ -38,7 +41,8 @@ def execute
 
         private
 
-        attr_reader :main_group, :current_user, :token_value, :token, :errors
+        attr_reader :main_group, :current_user, :token_value, :token
+        attr_accessor :errors, :project_ids
 
         # rubocop:disable Style/GuardClause -- Keep it explicit
         def ensure_dev_mode!
@@ -116,9 +120,6 @@ def create_subgroups
         def create_subgroup(params)
           url = "#{instance_url}#{GROUP_IMPORT_URL}"
 
-          headers = {
-            'PRIVATE-TOKEN' => token_value
-          }
           body = {
             name: params[:name],
             path: params[:name],
@@ -132,15 +133,12 @@ def create_subgroup(params)
           errors << { group: params[:name] } unless response.success?
 
           puts "API response for #{params[:name]} import"
-          puts response.body
+          puts response.parsed_response
         end
 
         def create_subproject(params)
           url = "#{instance_url}#{PROJECT_IMPORT_URL}"
 
-          headers = {
-            'PRIVATE-TOKEN' => token_value
-          }
           body = {
             name: params[:name],
             path: params[:name],
@@ -152,8 +150,9 @@ def create_subproject(params)
 
           errors << { project: params[:name] } unless response.success?
 
+          project_ids << response.parsed_response.fetch('id')
           puts "API response for #{params[:name]} import"
-          puts response.body
+          puts response.parsed_response
         end
 
         def instance_url
@@ -164,6 +163,38 @@ def delete_temporary_directory!
           FileUtils.rm_rf(Rails.root.join(DOWNLOAD_FOLDER, SAMPLES_FOLDER))
         end
 
+        def headers
+          {
+            'PRIVATE-TOKEN' => token_value
+          }
+        end
+
+        def check_import_status
+          time_counter = 0
+          imported_projects = project_ids.index_with { |_id| false }
+          until imported_projects.values.all?
+            break if time_counter > TIME_LIMIT
+
+            imported_projects.each do |id, _status|
+              puts "Checking import status for #{id}"
+
+              check_status = Gitlab::HTTP.get("#{instance_url}/api/v4/projects/#{id}/import",
+                headers: headers)
+
+              if check_status.success? &&
+                  check_status.parsed_response.fetch('import_status') == 'finished'
+                imported_projects[id] = true
+              end
+            end
+            time_counter += 5
+            sleep 5
+          end
+
+          return if imported_projects.values.all?
+
+          errors << { time_limit: "exceeded" }
+        end
+
         def print_output
           puts <<~MSG
             ----------------------------------------

ee/spec/lib/gitlab/duo/developments/setup_groups_for_model_evaluation_spec.rb

@@ -6,7 +6,7 @@
   let_it_be(:user) { create(:user, username: 'root') }
   let(:group) { create(:group) }
   let(:setup_evaluation) { described_class.new(group) }
-  let(:http_response) { instance_double(HTTParty::Response, body: 'File content') }
+  let(:http_response) { instance_double(HTTParty::Response) }
   let(:file_double) { instance_double(File) }
 
   before do
@@ -17,7 +17,7 @@
     context 'when the server is running' do
       before do
         allow(http_response).to receive(:success?).and_return(true)
-        allow(http_response).to receive(:parsed_response).and_return({})
+        allow(http_response).to receive(:parsed_response).and_return({ 'id' => 1 })
         allow(Gitlab::HTTP).to receive(:get).and_return(http_response)
         allow(Gitlab::HTTP).to receive(:get).with("https://gitlab.com/gitlab-org/ai-powered/datasets/-/package_files/135727282/download")
           .and_return(http_response)
@@ -30,6 +30,7 @@
         expect(setup_evaluation).to receive(:download_and_unpack_file)
         expect(setup_evaluation).to receive(:create_subgroups)
         expect(setup_evaluation).to receive(:create_subprojects)
+        expect(setup_evaluation).to receive(:check_import_status)
         expect(setup_evaluation).to receive(:delete_temporary_directory!)
         expect(setup_evaluation).to receive(:clean_up_token!)
         expect(setup_evaluation).to receive(:print_output)
@@ -150,6 +151,66 @@
         end
       end
 
+      describe '#check_import_status' do
+        before do
+          allow(http_response).to receive(:parsed_response).and_return({ 'import_status' => 'finished' })
+          allow(setup_evaluation).to receive(:token_value).and_return('token-string-1')
+        end
+
+        it 'checks import status for projects' do
+          setup_evaluation.instance_variable_set(:@project_ids, [1])
+
+          expect(Gitlab::HTTP).to receive(:get)
+                                    .with("#{setup_evaluation.send(:instance_url)}/api/v4/projects/1/import",
+                                      headers: { 'PRIVATE-TOKEN' => 'token-string-1' })
+                                    .and_return(http_response)
+
+          setup_evaluation.send(:check_import_status)
+          expect(setup_evaluation.send(:errors)).to be_empty
+        end
+
+        context 'with import not finished' do
+          before do
+            allow(http_response).to receive(:parsed_response).and_return({ 'import_status' => 'scheduled' },
+              { 'import_status' => 'finished' })
+            allow(http_response).to receive(:success?).and_return(true).twice
+          end
+
+          it 'waits for the import to finish' do
+            setup_evaluation.instance_variable_set(:@project_ids, [1])
+
+            expect(Gitlab::HTTP).to receive(:get)
+                                      .with("#{setup_evaluation.send(:instance_url)}/api/v4/projects/1/import",
+                                        headers: { 'PRIVATE-TOKEN' => 'token-string-1' })
+                                      .and_return(http_response)
+
+            expect(setup_evaluation).to receive(:sleep).with(5).twice
+            setup_evaluation.send(:check_import_status)
+            expect(setup_evaluation.send(:errors)).to be_empty
+          end
+        end
+
+        context 'when time limit is exceeded' do
+          before do
+            stub_const("#{described_class}::TIME_LIMIT", 1)
+            allow(http_response).to receive(:parsed_response).and_return({ 'import_status' => 'scheduled' })
+            allow(http_response).to receive(:success?).and_return(true)
+          end
+
+          it 'waits for the import to finish' do
+            setup_evaluation.instance_variable_set(:@project_ids, [1])
+
+            expect(Gitlab::HTTP).to receive(:get)
+                                      .with("#{setup_evaluation.send(:instance_url)}/api/v4/projects/1/import",
+                                        headers: { 'PRIVATE-TOKEN' => 'token-string-1' })
+                                      .and_return(http_response)
+
+            setup_evaluation.send(:check_import_status)
+            expect(setup_evaluation.send(:errors)).to include(time_limit: 'exceeded')
+          end
+        end
+      end
+
       context 'when running not in dev or test mode' do
         before do
           stub_env('RAILS_ENV', 'production')