diff --git a/tests/unit/test_asynchronous.py b/tests/unit/test_asynchronous.py
index f70a3e16..21c92feb 100644
--- a/tests/unit/test_asynchronous.py
+++ b/tests/unit/test_asynchronous.py
@@ -32,7 +32,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         @property
         def status(self):
             return 200 if self.sync else 202
-        
+
         async def raise_for_status(self):
             ...
 
diff --git a/tests/unit/test_upload.py b/tests/unit/test_upload.py
index 9812d06a..cdf52293 100644
--- a/tests/unit/test_upload.py
+++ b/tests/unit/test_upload.py
@@ -1,7 +1,7 @@
 import hashlib
 
 import pytest
-from aleph_message.models import StoreMessage
+from aleph_message.models import AlephMessage, StoreMessage
 from aleph_message.status import MessageStatus
 
 from aleph.sdk import AuthenticatedAlephHttpClient
@@ -33,7 +33,7 @@ async def test_upload_with_message():
         server_content = await client.download_file(file_hash=file_hash)
         assert server_content == content
 
-        server_message = await client.get_message(
+        server_message: AlephMessage = await client.get_message(
             item_hash=message.item_hash, message_type=StoreMessage
         )
-        assert server_message.content.item_hash == file_hash
\ No newline at end of file
+        assert server_message.content.item_hash == file_hash