fixes for review

Author: odesenfans

src/aleph/chains/common.py

@@ -1,7 +1,6 @@
 import asyncio
 import json
 import logging
-from dataclasses import asdict
 from enum import IntEnum
 from typing import Dict, Optional, Tuple, List
 
@@ -27,12 +26,12 @@
 from aleph.schemas.pending_messages import BasePendingMessage
 from aleph.schemas.validated_message import (
     validate_pending_message,
-    MessageConfirmation,
     ValidatedStoreMessage,
     ValidatedForgetMessage,
     make_confirmation_update_query,
     make_message_upsert_query,
 )
+from ..schemas.message_confirmation import MessageConfirmation
 from aleph.storage import get_json, pin_hash, add_json, get_message_content
 from .tx_context import TxContext
 
@@ -71,7 +70,7 @@ async def delayed_incoming(
     await PendingMessage.collection.insert_one(
         {
             "message": message.dict(exclude={"content"}),
-            "tx_context": asdict(tx_context) if tx_context is not None else None,
+            "tx_context": tx_context.dict() if tx_context else None,
             "check_message": check_message,
         }
     )
@@ -122,7 +121,7 @@ async def incoming(
     chain_name = tx_context.chain if tx_context is not None else None
     ids_key = (item_hash, sender, chain_name)
 
-    if tx_context is not None:
+    if tx_context:
         if seen_ids is not None:
             if ids_key in seen_ids.keys():
                 if tx_context.height > seen_ids[ids_key]:
@@ -371,5 +370,5 @@ async def incoming_chaindata(content: Dict, context: TxContext):
     For now we only add it to the database, it will be processed later.
     """
     await PendingTX.collection.insert_one(
-        {"content": content, "context": asdict(context)}
+        {"content": content, "context": context.dict()}
     )

src/aleph/chains/tx_context.py

@@ -1,11 +1,7 @@
-from dataclasses import dataclass
+from aleph.schemas.message_confirmation import MessageConfirmation
 
 
-@dataclass
-class TxContext:
-    chain: str
-    hash: str
-    height: int
-    # Transaction timestamp, in Unix time (number of seconds since epoch).
-    time: int
-    publisher: str
+# At the moment, confirmation = chain transaction. This might change, but in the meantime
+# having TxContext inherit MessageConfirmation avoids code duplication.
+class TxContext(MessageConfirmation):
+    pass

src/aleph/jobs/process_pending_messages.py

@@ -21,6 +21,7 @@
 from aleph.schemas.pending_messages import parse_message
 from aleph.services.p2p import singleton
 from .job_utils import prepare_loop, process_job_results
+from ..chains.tx_context import TxContext
 
 LOGGER = getLogger("jobs.pending_messages")
 
@@ -60,7 +61,8 @@ async def handle_pending_message(
         # If an invalid message somehow ended in pending messages, drop it.
         return [delete_pending_message_op]
 
-    tx_context = pending.get("tx_context")
+    tx_context_dict = pending.get("tx_context")
+    tx_context = TxContext.parse_obj(tx_context_dict) if tx_context_dict else None
 
     async with sem:
         status, operations = await incoming(

src/aleph/jobs/process_pending_txs.py

@@ -4,7 +4,6 @@
 
 import asyncio
 import logging
-from dataclasses import asdict
 from typing import List, Dict, Optional
 from typing import Set
 
@@ -57,7 +56,7 @@ async def handle_pending_tx(
                     operation=InsertOne(
                         {
                             "message": message.dict(exclude={"content"}),
-                            "tx_context": asdict(tx_context),
+                            "tx_context": tx_context.dict(),
                             "check_message": True,
                         }
                     ),

src/aleph/schemas/message_confirmation.py

@@ -0,0 +1,16 @@
+from aleph_message.models import Chain
+from pydantic import BaseModel, Field
+
+
+class MessageConfirmation(BaseModel):
+    chain: Chain = Field(..., description="Chain from which the confirmation was fetched.")
+    height: int = Field(..., description="Block in which the confirmation was published.")
+    hash: str = Field(
+        ...,
+        description="Hash of the transaction/block in which the confirmation was published.",
+    )
+    time: float = Field(
+        ...,
+        description="Transaction timestamp, in Unix time (number of seconds since epoch).",
+    )
+    publisher: str = Field(..., description="Publisher of the confirmation on chain.")

src/aleph/schemas/validated_message.py

@@ -12,8 +12,7 @@
     MessageType,
     PostContent,
     ProgramContent,
-    StoreContent, Chain,
-)
+    StoreContent, )
 from pydantic import BaseModel, Field
 
 from aleph.schemas.base_messages import AlephBaseMessage, ContentType, MType
@@ -25,17 +24,10 @@
     PendingProgramMessage,
     PendingStoreMessage,
 )
+from .message_confirmation import MessageConfirmation
 from .message_content import MessageContent
 
 
-class MessageConfirmation(BaseModel):
-    chain: Chain
-    height: int
-    hash: str
-    time: float
-    publisher: str
-
-
 class EngineInfo(BaseModel):
     hash: str = Field(alias="Hash")
     size: int = Field(alias="Size")

tests/chains/test_confirmation.py

@@ -1,5 +1,4 @@
 import json
-from dataclasses import asdict
 from typing import Dict
 
 import pytest
@@ -72,7 +71,7 @@ async def test_confirm_message(test_db):
 
     assert message_in_db is not None
     assert message_in_db["confirmed"]
-    expected_confirmations = [asdict(tx_context)]
+    expected_confirmations = [tx_context.dict()]
     assert message_in_db["confirmations"] == expected_confirmations
 
     capped_message_after_confirmation = await CappedMessage.collection.find_one(
@@ -111,7 +110,7 @@ async def test_process_confirmed_message(test_db):
     assert message_in_db is not None
     assert message_in_db["confirmed"]
 
-    expected_confirmations = [asdict(tx_context)]
+    expected_confirmations = [tx_context.dict()]
     assert message_in_db["confirmations"] == expected_confirmations
 
     capped_message_in_db = await CappedMessage.collection.find_one(
@@ -120,4 +119,4 @@ async def test_process_confirmed_message(test_db):
 
     assert remove_id_key(message_in_db) == remove_id_key(capped_message_in_db)
     assert capped_message_in_db["confirmed"]
-    assert capped_message_in_db["confirmations"] == [asdict(tx_context)]
+    assert capped_message_in_db["confirmations"] == [tx_context.dict()]

tests/schemas/test_validated_messages.py

@@ -16,10 +16,10 @@
 from aleph.schemas.validated_message import (
     validate_pending_message,
     BaseValidatedMessage,
-    MessageConfirmation,
     ValidatedAggregateMessage,
     ValidatedStoreMessage,
 )
+from aleph.schemas.message_confirmation import MessageConfirmation
 
 
 def check_basic_message_fields(message: BaseValidatedMessage, message_dict: Dict):

tests/storage/test_store_message.py

@@ -5,10 +5,10 @@
 from aleph.handlers.storage import handle_new_storage
 from aleph.schemas.message_content import ContentSource, RawContent
 from aleph.schemas.validated_message import (
-    MessageConfirmation,
     ValidatedStoreMessage,
     StoreContentWithMetadata,
 )
+from aleph.schemas.message_confirmation import MessageConfirmation
 from message_test_helpers import make_validated_message_from_dict