Add tests for token validator and widget util.

aiogram/utils/auth_widget.py

@@ -3,6 +3,21 @@ import hashlib
 import hmac
 
 
+def generate_hash(data, token):
+    """
+    Generate secret hash
+
+    :param data:
+    :param token:
+    :return:
+    """
+    secret = hashlib.sha256()
+    secret.update(token.encode('utf-8'))
+    sorted_params = collections.OrderedDict(sorted(data.items()))
+    msg = "\n".join(["{}={}".format(k, v) for k, v in sorted_params.items() if k != 'hash'])
+    return hmac.new(secret.digest(), msg.encode('utf-8'), digestmod=hashlib.sha256).hexdigest()
+
+
 def check_token(data, token):
     """
     Validate auth token
@@ -14,12 +29,5 @@ def check_token(data, token):
     :param token:
     :return:
     """
-    secret = hashlib.sha256()
+    param_hash = data.get('hash', '') or ''
-    secret.update(token.encode('utf-8'))
+    return param_hash == generate_hash(data, token)
-    sorted_params = collections.OrderedDict(sorted(data.items()))
-    param_hash = sorted_params.pop('hash', '') or ''
-    msg = "\n".join(["{}={}".format(k, v) for k, v in sorted_params.items()])
-
-    if param_hash == hmac.new(secret.digest(), msg.encode('utf-8'), digestmod=hashlib.sha256).hexdigest():
-        return True
-    return False

tests/test_token.py

@@ -0,0 +1,41 @@
+import pytest
+
+from aiogram.bot import api
+from aiogram.utils import auth_widget, exceptions
+
+VALID_TOKEN = '123456789:AABBCCDDEEFFaabbccddeeff-1234567890'
+INVALID_TOKEN = '123456789:AABBCCDDEEFFaabbccddeeff 123456789'  # Space in token and wrong length
+
+VALID_DATA = {
+    'date': 1525385236,
+    'first_name': 'Test',
+    'last_name': 'User',
+    'id': 123456789,
+    'username': 'username',
+    'hash': '69a9871558fbbe4cd0dbaba52fa1cc4f38315d3245b7504381a64139fb024b5b'
+}
+INVALID_DATA = {
+    'date': 1525385237,
+    'first_name': 'Test',
+    'last_name': 'User',
+    'id': 123456789,
+    'username': 'username',
+    'hash': '69a9871558fbbe4cd0dbaba52fa1cc4f38315d3245b7504381a64139fb024b5b'
+}
+
+
+def test_valid_token():
+    assert api.check_token(VALID_TOKEN)
+
+
+def test_invalid_token():
+    with pytest.raises(exceptions.ValidationError):
+        api.check_token(INVALID_TOKEN)
+
+
+def test_widget():
+    assert auth_widget.check_token(VALID_DATA, VALID_TOKEN)
+
+
+def test_invalid_widget_data():
+    assert not auth_widget.check_token(INVALID_DATA, VALID_TOKEN)