fix failed ut from PR 1472 (#1479)

* fix failed ut from PR 1472 Signed-off-by: Yaliang Wu <[email protected]> * exclude class for low coverage Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]>
opensearch-project · Oct 10, 2023 · da5d829 · da5d829
1 parent 201c8a8
commit da5d829
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/plugin/build.gradle b/plugin/build.gradle
@@ -267,6 +267,7 @@ List<String> jacocoExclusions = [
         'org.opensearch.ml.action.deploy.TransportDeployModelAction',
         'org.opensearch.ml.action.deploy.TransportDeployModelOnNodeAction',
         'org.opensearch.ml.action.prediction.TransportPredictionTaskAction',
+        'org.opensearch.ml.action.prediction.TransportPredictionTaskAction.1',
         'org.opensearch.ml.action.tasks.GetTaskTransportAction',
         'org.opensearch.ml.action.tasks.SearchTaskTransportAction',
         'org.opensearch.ml.model.MLModelManager',

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java
@@ -436,7 +436,10 @@ public void setModelInfo(String modelId, MLModel mlModel) {
     }
 
     public MLModel getModelInfo(String modelId) {
-        MLModelCache mlModelCache = getExistingModelCache(modelId);
+        MLModelCache mlModelCache = modelCaches.get(modelId);
+        if (mlModelCache == null) {
+            return null;
+        }
         return mlModelCache.getCachedModelInfo();
     }