Refactor test

1 files changed, 5 insertions, 9 deletions

diff --git a/config-model/src/test/java/com/yahoo/vespa/model/ml/ModelEvaluationTest.java b/config-model/src/test/java/com/yahoo/vespa/model/ml/ModelEvaluationTest.java
index 1b7b96dd9d5..35e6642d7cb 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/ml/ModelEvaluationTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/ml/ModelEvaluationTest.java
@@ -6,7 +6,6 @@ import ai.vespa.models.evaluation.ModelsEvaluator;
 import ai.vespa.models.evaluation.RankProfilesConfigImporter;
 import com.yahoo.config.FileReference;
 import com.yahoo.config.application.api.ApplicationPackage;
-import com.yahoo.config.model.ApplicationPackageTester;
 import com.yahoo.filedistribution.fileacquirer.FileAcquirer;
 import com.yahoo.filedistribution.fileacquirer.MockFileAcquirer;
 import com.yahoo.io.IOUtils;
@@ -19,7 +18,6 @@ import com.yahoo.vespa.model.VespaModel;
 import com.yahoo.vespa.model.container.ContainerCluster;
 import org.junit.After;
 import org.junit.Test;
-import org.xml.sax.SAXException;
 
 import java.io.IOException;
 import java.util.Set;
@@ -42,10 +40,9 @@ public class ModelEvaluationTest {
     }
 
     @Test
-    public void testMl_ServingApplication() throws SAXException, IOException {
-        ApplicationPackageTester tester = ApplicationPackageTester.create(appDir.toString());
-        VespaModel model = new VespaModel(tester.app());
-        assertHasMlModels(model);
+    public void testMl_ServingApplication() throws IOException {
+        ImportedModelTester tester = new ImportedModelTester("ml_serving", appDir);
+        assertHasMlModels(tester.createVespaModel());
 
         // At this point the expression is stored - copy application to another location which do not have a models dir
         Path storedAppDir = appDir.append("copy");
@@ -54,9 +51,8 @@ public class ModelEvaluationTest {
             IOUtils.copy(appDir.append("services.xml").toString(), storedAppDir.append("services.xml").toString());
             IOUtils.copyDirectory(appDir.append(ApplicationPackage.MODELS_GENERATED_DIR).toFile(),
                                   storedAppDir.append(ApplicationPackage.MODELS_GENERATED_DIR).toFile());
-            ApplicationPackageTester storedTester = ApplicationPackageTester.create(storedAppDir.toString());
-            VespaModel storedModel = new VespaModel(storedTester.app());
-            assertHasMlModels(storedModel);
+            ImportedModelTester storedTester = new ImportedModelTester("ml_serving", storedAppDir);
+            assertHasMlModels(storedTester.createVespaModel());
         }
         finally {
             IOUtils.recursiveDeleteDir(storedAppDir.toFile());