diff --git a/tests/test_token_plugins.py b/tests/test_token_plugins.py index 363ae35..4e029eb 100644 --- a/tests/test_token_plugins.py +++ b/tests/test_token_plugins.py @@ -7,6 +7,11 @@ import unittest from unittest.mock import patch, MagicMock from jwcrypto import jwt, jwk +try: + import redis +except ImportError: + redis = None + from websockify.token_plugins import parse_source_args, ReadOnlyTokenFile, JWTTokenApi, TokenRedis @@ -206,9 +211,7 @@ class JWSTokenTestCase(unittest.TestCase): class TokenRedisTestCase(unittest.TestCase): def setUp(self): - try: - import redis - except ImportError: + if redis is None: patcher = patch.dict(sys.modules, {'redis': MagicMock()}) patcher.start() self.addCleanup(patcher.stop) diff --git a/websockify/token_plugins.py b/websockify/token_plugins.py index 9307356..ba6e1d3 100644 --- a/websockify/token_plugins.py +++ b/websockify/token_plugins.py @@ -5,6 +5,11 @@ import re import json from pathlib import Path +try: + import redis +except ImportError: + redis = None + logger = logging.getLogger(__name__) _SOURCE_SPLIT_REGEX = re.compile( @@ -256,9 +261,7 @@ class TokenRedis(BasePlugin): pip install redis """ def __init__(self, src): - try: - import redis - except ImportError: + if redis is None: logger.error("Unable to load redis module") sys.exit() # Default values @@ -314,9 +317,7 @@ class TokenRedis(BasePlugin): sys.exit() def lookup(self, token): - try: - import redis - except ImportError: + if redis is None: logger.error("package redis not found, are you sure you've installed them correctly?") sys.exit()