Skip to content

Commit

Permalink
test: clean up and improve conditions for optional torch/TF tests
Browse files Browse the repository at this point in the history
  • Loading branch information
wookayin committed Jan 23, 2025
1 parent 2aafdfb commit 116d087
Showing 1 changed file with 9 additions and 2 deletions.
11 changes: 9 additions & 2 deletions imgcat/test_imgcat.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,11 @@
from imgcat import imgcat


def _importable(package_name: str) -> bool:
import importlib.util
return importlib.util.find_spec(package_name) is not None


@pytest.fixture
def mock_env(monkeypatch, env_profile):
"""Mock environment variables (especially, TMUX)"""
Expand Down Expand Up @@ -138,9 +143,10 @@ def test_numpy(self):
a[..., 0], a[..., 1], a[..., 2] = 0x37 / 255., 0xb2 / 255., 0x4d / 255.
imgcat(a)

@pytest.mark.skipif(not _importable('torch'), reason="No torch installed")
@parametrize_env
def test_torch(self):
import torch
import torch # type: ignore

# uint8, grayscale
with self.capture_and_validate():
Expand All @@ -156,13 +162,15 @@ def test_torch(self):
a = torch.ones([3, 32, 32], dtype=torch.uint8) * 0
imgcat(a)

@pytest.mark.skipif(not _importable('tensorflow'), reason="No tensorflow installed")
@parametrize_env
def test_tensorflow(self):
try:
import tensorflow.compat.v2 as tf # type: ignore # noqa
except ImportError:
pytest.skip("No tensorflow available")

import tensorflow.compat.v2 as tf # type: ignore
tf.enable_v2_behavior()
assert tf.executing_eagerly(), "Eager execution should be enabled."

Expand Down Expand Up @@ -250,7 +258,6 @@ def test_args_another(self):
preserve_aspect_ratio=False, fp=b)

v = b.getvalue()
assert b'size=82;' in v
assert b'height=12;' in v
assert b'width=10;' in v
assert b'name=Zm9vLnBuZw==;' in v # foo.png
Expand Down

0 comments on commit 116d087

Please sign in to comment.