1 files changed, 14 insertions, 37 deletions

diff --git a/spec/services/packages/ml_model/create_package_file_service_spec.rb b/spec/services/packages/ml_model/create_package_file_service_spec.rb
index 30a6bedd07b..505c8038976 100644
--- a/spec/services/packages/ml_model/create_package_file_service_spec.rb
+++ b/spec/services/packages/ml_model/create_package_file_service_spec.rb
@@ -8,6 +8,8 @@ RSpec.describe Packages::MlModel::CreatePackageFileService, feature_category: :m
     let_it_be(:user) { create(:user) }
     let_it_be(:pipeline) { create(:ci_pipeline, user: user, project: project) }
     let_it_be(:file_name) { 'myfile.tar.gz.1' }
+    let_it_be(:model) { create(:ml_models, user: user, project: project) }
+    let_it_be(:model_version) { create(:ml_model_versions, :with_package, model: model, version: '0.1.0') }
 
     let(:build) { instance_double(Ci::Build, pipeline: pipeline) }
 
@@ -26,47 +28,24 @@ RSpec.describe Packages::MlModel::CreatePackageFileService, feature_category: :m
       FileUtils.rm_f(temp_file)
     end
 
-    context 'without existing package' do
+    context 'when model version is nil' do
       let(:params) do
         {
-          package_name: 'new_model',
-          package_version: '1.0.0',
+          model_version: nil,
           file: file,
           file_name: file_name
         }
       end
 
-      it 'creates package file', :aggregate_failures do
-        expect { execute_service }
-          .to change { Packages::MlModel::Package.count }.by(1)
-          .and change { Packages::PackageFile.count }.by(1)
-          .and change { Packages::PackageFileBuildInfo.count }.by(0)
-          .and change { Ml::ModelVersion.count }.by(1)
-
-        new_model = Packages::MlModel::Package.last
-        package_file = new_model.package_files.last
-        new_model_version = Ml::ModelVersion.last
-
-        expect(new_model.name).to eq('new_model')
-        expect(new_model.version).to eq('1.0.0')
-        expect(new_model.status).to eq('default')
-        expect(package_file.package).to eq(new_model)
-        expect(package_file.file_name).to eq(file_name)
-        expect(package_file.size).to eq(file.size)
-        expect(package_file.file_sha256).to eq(sha256)
-        expect(new_model_version.name).to eq('new_model')
-        expect(new_model_version.version).to eq('1.0.0')
-        expect(new_model_version.package).to eq(new_model)
+      it 'does not create package file', :aggregate_failures do
+        expect(execute_service).to be(nil)
       end
     end
 
-    context 'with existing package' do
-      let_it_be(:model) { create(:ml_model_package, creator: user, project: project, version: '0.1.0') }
-
+    context 'with existing model version' do
       let(:params) do
         {
-          package_name: model.name,
-          package_version: model.version,
+          model_version: model_version,
           file: file,
           file_name: file_name,
           status: :hidden,
@@ -76,18 +55,16 @@ RSpec.describe Packages::MlModel::CreatePackageFileService, feature_category: :m
 
       it 'adds the package file and updates status and ci_build', :aggregate_failures do
         expect { execute_service }
-          .to change { project.packages.ml_model.count }.by(0)
-          .and change { model.package_files.count }.by(1)
+          .to change { model_version.package.package_files.count }.by(1)
           .and change { Packages::PackageFileBuildInfo.count }.by(1)
 
-        model.reload
-
-        package_file = model.package_files.last
+        package = model_version.reload.package
+        package_file = package.package_files.last
 
-        expect(model.build_infos.first.pipeline).to eq(build.pipeline)
-        expect(model.status).to eq('hidden')
+        expect(package.build_infos.first.pipeline).to eq(build.pipeline)
+        expect(package.status).to eq('hidden')
 
-        expect(package_file.package).to eq(model)
+        expect(package_file.package).to eq(package)
         expect(package_file.file_name).to eq(file_name)
         expect(package_file.size).to eq(file.size)
         expect(package_file.file_sha256).to eq(sha256)