diff --git a/ee/lib/gitlab/duo/developments/setup_groups_for_model_evaluation.rb b/ee/lib/gitlab/duo/developments/setup_groups_for_model_evaluation.rb
index 4bbfa738e272045e7044d52b675467f2ea2644bd..97288d6e9dc9d0095e441732355e5360320a7f2c 100644
--- a/ee/lib/gitlab/duo/developments/setup_groups_for_model_evaluation.rb
+++ b/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
----------------------------------------
diff --git a/ee/spec/lib/gitlab/duo/developments/setup_groups_for_model_evaluation_spec.rb b/ee/spec/lib/gitlab/duo/developments/setup_groups_for_model_evaluation_spec.rb
index 4c76f00845693e693b9f73b2c73b2d76a344b4e5..e25cd1eeff6ec9f44b3e5fa91235302d1608a43f 100644
--- a/ee/spec/lib/gitlab/duo/developments/setup_groups_for_model_evaluation_spec.rb
+++ b/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')