fix failed ut from PR 1472 (#1479) (#1510)

* 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]> (cherry picked from commit da5d829) Co-authored-by: Yaliang Wu <[email protected]>
opensearch-project · Nov 20, 2023 · 520b5db · 520b5db
1 parent c03e20a
commit 520b5db
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
@@ -272,6 +272,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();
     }