From 65edaeb2b5e217ab64958d5166df326a0185dee0 Mon Sep 17 00:00:00 2001
From: Alex Root Junior <jroot.junior@gmail.com>
Date: Fri, 4 May 2018 01:19:36 +0300
Subject: [PATCH] Add tests for token validator and widget util.

---
 aiogram/utils/auth_widget.py | 26 +++++++++++++++--------
 tests/test_token.py          | 41 ++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 9 deletions(-)
 create mode 100644 tests/test_token.py

diff --git a/aiogram/utils/auth_widget.py b/aiogram/utils/auth_widget.py
index b5cce802..800dfc12 100644
--- a/aiogram/utils/auth_widget.py
+++ b/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()
-    secret.update(token.encode('utf-8'))
-    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
+    param_hash = data.get('hash', '') or ''
+    return param_hash == generate_hash(data, token)
diff --git a/tests/test_token.py b/tests/test_token.py
new file mode 100644
index 00000000..b8a6087f
--- /dev/null
+++ b/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)