From 9e9e4c7d356652d81e17b639197ac76ee5ed23e8 Mon Sep 17 00:00:00 2001
From: VincyZhang <wenxin.zhang@intel.com>
Date: Sat, 24 Feb 2024 11:52:55 +0800
Subject: [PATCH] fix ut (#1307)

---
 .../workflows/script/unitTest/run_unit_test_neuralchat.sh   | 6 +++++-
 .../workflows/script/unitTest/run_unit_test_neuralspeed.sh  | 5 ++++-
 .github/workflows/script/unitTest/run_unit_test_optimize.sh | 5 ++++-
 tests/CI/test_llava.py                                      | 2 +-
 4 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/script/unitTest/run_unit_test_neuralchat.sh b/.github/workflows/script/unitTest/run_unit_test_neuralchat.sh
index 5c2915ef144..70101f18f3e 100644
--- a/.github/workflows/script/unitTest/run_unit_test_neuralchat.sh
+++ b/.github/workflows/script/unitTest/run_unit_test_neuralchat.sh
@@ -71,7 +71,11 @@ function pytest() {
     if [ $(grep -c "Segmentation fault" ${ut_log_name}) != 0 ]; then
        $BOLD_RED && echo "Segmentation Fault found in UT, please check the output..." && $RESET
         exit 1
-    fi  
+    fi
+    if [ $(grep -c "ImportError:" ${ut_log_name}) != 0 ]; then
+       $BOLD_RED && echo "ImportError found in UT, please check the output..." && $RESET
+        exit 1
+    fi
     $BOLD_GREEN && echo "UT finished successfully! " && $RESET
 }
 
diff --git a/.github/workflows/script/unitTest/run_unit_test_neuralspeed.sh b/.github/workflows/script/unitTest/run_unit_test_neuralspeed.sh
index 67c5a006d0d..2d6734fb37d 100644
--- a/.github/workflows/script/unitTest/run_unit_test_neuralspeed.sh
+++ b/.github/workflows/script/unitTest/run_unit_test_neuralspeed.sh
@@ -22,7 +22,10 @@ function neuralspeed_test() {
         [ $(grep -c "Segmentation fault" ${ut_log_name}) != 0 ] ||
         [ $(grep -c "core dumped" ${ut_log_name}) != 0 ] ||
         [ $(grep -c "==ERROR:" ${ut_log_name}) != 0 ] ||
-        [ $(grep -c "ModuleNotFoundError:" ${ut_log_name}) != 0 ]; then
+        [ $(grep -c "ModuleNotFoundError:" ${ut_log_name}) != 0 ] ||
+        [ $(grep -c "ImportError:" ${ut_log_name}) != 0 ]; then
+        exit 1
+    fi; then
         $BOLD_RED && echo "Find errors in engine test, please check the output..." && $RESET
         exit 1
     else
diff --git a/.github/workflows/script/unitTest/run_unit_test_optimize.sh b/.github/workflows/script/unitTest/run_unit_test_optimize.sh
index 735fadc9181..f300165f968 100644
--- a/.github/workflows/script/unitTest/run_unit_test_optimize.sh
+++ b/.github/workflows/script/unitTest/run_unit_test_optimize.sh
@@ -67,7 +67,10 @@ function pytest() {
        $BOLD_RED && echo "Segmentation Fault found in UT, please check the output..." && $RESET
         exit 1
     fi  
-
+    if [ $(grep -c "ImportError:" ${ut_log_name}) != 0 ]; then
+       $BOLD_RED && echo "ImportError found in UT, please check the output..." && $RESET
+        exit 1
+    fi
     $BOLD_GREEN && echo "UT finished successfully! " && $RESET
 }
 
diff --git a/tests/CI/test_llava.py b/tests/CI/test_llava.py
index cd5e9df809e..d0fa6b8e6eb 100644
--- a/tests/CI/test_llava.py
+++ b/tests/CI/test_llava.py
@@ -14,7 +14,7 @@
 
 import os
 import unittest
-from intel_extension_for_transformers.transformers.modeling.llava_models import LlavaMistralForCausalLM
+from intel_extension_for_transformers.transformers.modeling.llava_models.llava_mistral import LlavaMistralForCausalLM
 from accelerate import init_empty_weights
 from transformers import AutoConfig, AutoTokenizer
 import torch