diff --git a/src/utils/misc.py b/src/utils/misc.py index b4f4969..e5342c4 100644 --- a/src/utils/misc.py +++ b/src/utils/misc.py @@ -77,9 +77,13 @@ def hms(seconds: int): def load_template(path, scale_factor: float = 1.0, alpha: bool = False): if os.path.isfile(path): - template_img = cv2.imread(path, cv2.IMREAD_UNCHANGED) if alpha else cv2.imread(path) - template_img = cv2.resize(template_img, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_NEAREST) - return template_img + try: + template_img = cv2.imread(path, cv2.IMREAD_UNCHANGED) if alpha else cv2.imread(path) + template_img = cv2.resize(template_img, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_NEAREST) + return template_img + except Exception as e: + print(e) + raise ValueError(f"Could not load template: {path}") return None def alpha_to_mask(img: np.ndarray): diff --git a/test/utils/misc_test.py b/test/utils/misc_test.py new file mode 100644 index 0000000..e5c025b --- /dev/null +++ b/test/utils/misc_test.py @@ -0,0 +1,18 @@ +import pytest +from logger import Logger +from utils.misc import load_template + + +class TestUtilsMisc: + def setup_method(self): + Logger.init() + Logger.remove_file_logger() + + @pytest.mark.parametrize("path, should_be_success", [ + ("test/assets/hero_select.png", True), + ("some/random/path/that/not/a/file.png", False), + ]) + def test_load_template(self, path: str, should_be_success: bool): + template_img = load_template(path, 1.0, alpha=True) + success = template_img is not None + assert(success == should_be_success)