Update various unit tests to skip if CUDA is not available. (#212)

NVIDIA · Nov 5, 2024 · fa07a5c · fa07a5c
1 parent 3777639
commit fa07a5c
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 33 deletions.
diff --git a/tests/nv_ingest/modules/filters/test_image_dedup.py b/tests/nv_ingest/modules/filters/test_image_dedup.py
@@ -8,8 +8,6 @@
 import pandas as pd
 import pytest
 
-from nv_ingest.modules.filters.image_dedup import _apply_dedup_filter
-from nv_ingest.modules.filters.image_dedup import _cpu_only_apply_dedup_filter
 from nv_ingest.schemas.metadata_schema import ContentTypeEnum
 from nv_ingest.schemas.metadata_schema import ImageTypeEnum
 from nv_ingest.schemas.metadata_schema import SourceTypeEnum
@@ -19,10 +17,11 @@
 from ....import_checks import MORPHEUS_IMPORT_OK
 
 if CUDA_DRIVER_OK and MORPHEUS_IMPORT_OK:
+    import cudf
     from morpheus.messages import ControlMessage
     from morpheus.messages import MessageMeta
-
-    import cudf
+    from nv_ingest.modules.filters.image_dedup import _cpu_only_apply_dedup_filter
+    from nv_ingest.modules.filters.image_dedup import _apply_dedup_filter
 
 
 def valid_image_dedup_task(should_filter):
@@ -105,6 +104,11 @@ def test_apply_dedup(should_filter, expected0, expected1, expected2):
         assert (mdf.iloc[0:3]["document_type"] == ContentTypeEnum.INFO_MSG.value).sum() == expected2
 
 
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 @pytest.mark.parametrize(
     "should_filter, expected0, expected1, expected2",
     [

diff --git a/tests/nv_ingest/modules/filters/test_image_filter.py b/tests/nv_ingest/modules/filters/test_image_filter.py
@@ -6,8 +6,6 @@
 import pandas as pd
 import pytest
 
-from nv_ingest.modules.filters.image_filter import _apply_filter
-from nv_ingest.modules.filters.image_filter import _cpu_only_apply_filter
 from nv_ingest.schemas.metadata_schema import ContentTypeEnum
 from nv_ingest.schemas.metadata_schema import ImageTypeEnum
 from nv_ingest.schemas.metadata_schema import SourceTypeEnum
@@ -18,10 +16,12 @@
 from ....import_checks import MORPHEUS_IMPORT_OK
 
 if CUDA_DRIVER_OK and MORPHEUS_IMPORT_OK:
+    import cudf
     from morpheus.messages import ControlMessage
     from morpheus.messages import MessageMeta
 
-    import cudf
+    from nv_ingest.modules.filters.image_filter import _apply_filter
+    from nv_ingest.modules.filters.image_filter import _cpu_only_apply_filter
 
 
 def valid_image_filter_task(should_filter):
@@ -123,6 +123,11 @@ def test_apply_filter(should_filter, width, height, expected0, expected1):
         assert (mdf.iloc[0:3]["document_type"] == ContentTypeEnum.INFO_MSG.value).sum() == expected1
 
 
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 @pytest.mark.parametrize(
     "should_filter, width, height, expected0, expected1",
     [

diff --git a/tests/nv_ingest/modules/storages/test_image_storage.py b/tests/nv_ingest/modules/storages/test_image_storage.py
@@ -3,18 +3,18 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import pytest
+import pandas as pd
 from minio import Minio
-from morpheus.messages import ControlMessage
-from morpheus.messages import MessageMeta
 
-from nv_ingest.modules.storages.image_storage import upload_images
 from nv_ingest.schemas.metadata_schema import ContentTypeEnum
 
 from ....import_checks import CUDA_DRIVER_OK
 from ....import_checks import MORPHEUS_IMPORT_OK
 
 if CUDA_DRIVER_OK and MORPHEUS_IMPORT_OK:
-    import pandas as pd
+    from morpheus.messages import ControlMessage
+    from morpheus.messages import MessageMeta
+    from nv_ingest.modules.storages.image_storage import upload_images
 
     import cudf
 
@@ -36,9 +36,9 @@ def bucket_exists(self, *args, **kwargs):
 @pytest.fixture
 def mock_minio(mocker):
     def mock_minio_init(
-        cls,
-        *args,
-        **kwargs,
+            cls,
+            *args,
+            **kwargs,
     ):
         return MockMinioClient(*args, **kwargs)
 

diff --git a/tests/nv_ingest/modules/telemetry/test_otel_tracer.py b/tests/nv_ingest/modules/telemetry/test_otel_tracer.py
@@ -1,4 +1,3 @@
-import cudf
 from datetime import datetime
 from morpheus.messages import ControlMessage
 

diff --git a/tests/nv_ingest/modules/transforms/test_associate_nearby_text.py b/tests/nv_ingest/modules/transforms/test_associate_nearby_text.py
@@ -8,17 +8,16 @@
 import pandas as pd
 import pytest
 
-from nv_ingest.modules.transforms.associate_nearby_text import _associate_nearby_text_blocks
-from nv_ingest.modules.transforms.associate_nearby_text import _get_bbox
-from nv_ingest.modules.transforms.associate_nearby_text import _get_center
-from nv_ingest.modules.transforms.associate_nearby_text import _is_nearby_text
 from nv_ingest.schemas.metadata_schema import TextTypeEnum
 
 from ....import_checks import CUDA_DRIVER_OK
 from ....import_checks import MORPHEUS_IMPORT_OK
 
 if MORPHEUS_IMPORT_OK and CUDA_DRIVER_OK:
-    pass
+    from nv_ingest.modules.transforms.associate_nearby_text import _get_bbox
+    from nv_ingest.modules.transforms.associate_nearby_text import _get_center
+    from nv_ingest.modules.transforms.associate_nearby_text import _is_nearby_text
+    from nv_ingest.modules.transforms.associate_nearby_text import _associate_nearby_text_blocks
 
 
 @pytest.fixture
@@ -45,37 +44,65 @@ def create_sample_df():
     }
     return pd.DataFrame(data)
 
-
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 def test_get_center():
     assert _get_center((0, 0, 10, 10)) == (5, 5), "Should return the center of the bounding box"
     assert _get_center((1, 2, 3, 4)) == (2, 3), "Should return the center of the bounding box with non-zero origin"
 
-
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 def test_is_nearby_text_true():
     row = {"text_metadata": {"text_type": TextTypeEnum.NEARBY_BLOCK}}
     assert _is_nearby_text(row) is True, "Should identify text as NEARBY_BLOCK correctly"
 
-
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 def test_is_nearby_text_false():
     row = {"text_metadata": {"text_type": "OTHER_TYPE"}}
     assert _is_nearby_text(row) is False, "Should correctly identify non-NEARBY_BLOCK text"
 
-
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 def test_is_nearby_text_no_metadata():
     row = {}
     assert _is_nearby_text(row) is False, "Should return False when no text_metadata is present"
 
-
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 def test_get_bbox_from_text_metadata():
     row = {"text_metadata": {"text_location": (1, 2, 3, 4)}}
     assert _get_bbox(row) == (1, 2, 3, 4), "Should extract bbox from text metadata correctly"
 
-
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 def test_get_bbox_from_image_metadata():
     row = {"image_metadata": {"image_location": (5, 6, 7, 8)}}
     assert _get_bbox(row) == (5, 6, 7, 8), "Should extract bbox from image metadata when text metadata is absent"
 
-
+@pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")
+@pytest.mark.skipif(
+    not CUDA_DRIVER_OK,
+    reason="Test environment does not have a compatible CUDA driver.",
+)
 def test_get_bbox_no_metadata():
     row = {}
     assert _get_bbox(row) is None, "Should return None when no relevant metadata is present"

diff --git a/tests/nv_ingest/util/converters/test_dftools.py b/tests/nv_ingest/util/converters/test_dftools.py
@@ -5,18 +5,16 @@
 import json
 
 import pytest
-
-from nv_ingest.util.converters.dftools import cudf_to_json
-from nv_ingest.util.converters.dftools import cudf_to_pandas
-from nv_ingest.util.converters.dftools import pandas_to_cudf
+import pandas as pd
 
 from ....import_checks import CUDA_DRIVER_OK
 from ....import_checks import MORPHEUS_IMPORT_OK
 
 if CUDA_DRIVER_OK and MORPHEUS_IMPORT_OK:
-    import pandas as pd
-
     import cudf
+    from nv_ingest.util.converters.dftools import cudf_to_json
+    from nv_ingest.util.converters.dftools import cudf_to_pandas
+    from nv_ingest.util.converters.dftools import pandas_to_cudf
 
 
 @pytest.mark.skipif(not MORPHEUS_IMPORT_OK, reason="Morpheus modules are not available.")