Skip to content

Commit da8f894

Browse files
suiyoubithomasdhc
andauthored
Skip IV2 Unit Test if package not installed (#1111)
* Skip IV2 Unit Test if package not installed Signed-off-by: Ao Tang <aot@nvidia.com> * syntax Signed-off-by: Ao Tang <aot@nvidia.com> --------- Signed-off-by: Ao Tang <aot@nvidia.com> Co-authored-by: Dong Hyuk Chang <thomaschang26@tutanota.com>
1 parent ba78c3c commit da8f894

File tree

2 files changed

+21
-12
lines changed

2 files changed

+21
-12
lines changed

tests/models/test_internvideo2.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,18 @@
1818
from unittest.mock import MagicMock, Mock, patch
1919

2020
import numpy as np
21+
import pytest
2122
import torch
2223

23-
from nemo_curator.stages.video.embedding.internvideo2 import (
24-
InternVideo2EmbeddingStage,
25-
InternVideo2FrameCreationStage,
26-
)
24+
try:
25+
from nemo_curator.stages.video.embedding.internvideo2 import (
26+
InternVideo2EmbeddingStage,
27+
InternVideo2FrameCreationStage,
28+
)
29+
except ImportError:
30+
pytest.skip("InternVideo2 package is not available")
31+
32+
2733
from nemo_curator.tasks.video import Clip, Video, VideoTask
2834

2935
# Create a random generator for consistent testing

tests/models/test_internvideo2_mm.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,17 @@
2121
import torch
2222
from easydict import EasyDict
2323

24-
from nemo_curator.models.internvideo2_mm import (
25-
BERT_MODEL_ID,
26-
INTERNVIDEO2_MODEL_FILE,
27-
INTERNVIDEO2_MODEL_ID,
28-
InternVideo2MultiModality,
29-
_InternVideo2Stage2Wrapper,
30-
_setup_internvideo2,
31-
)
24+
try:
25+
from nemo_curator.models.internvideo2_mm import (
26+
BERT_MODEL_ID,
27+
INTERNVIDEO2_MODEL_FILE,
28+
INTERNVIDEO2_MODEL_ID,
29+
InternVideo2MultiModality,
30+
_InternVideo2Stage2Wrapper,
31+
_setup_internvideo2,
32+
)
33+
except ImportError:
34+
pytest.skip("InternVideo2 package is not available")
3235

3336
# Create a random generator for consistent testing
3437
rng = np.random.default_rng(42)

0 commit comments

Comments
 (0)