Merge pull request ethereum#850 from pipermerriam/piper/protocol-fact…

…ories

Implement factories for generating devp2p protocols

newsfragments/850.feature.rst

@@ -0,0 +1 @@
+Add factories for creating devp2p protocols and commands for testing.

p2p/tools/factories.py

@@ -1,27 +1,31 @@
 import asyncio
+import itertools
 import random
 import socket
-from typing import Any, Tuple
+from typing import Any, Generator, Iterable, Tuple, Type
 
 from cancel_token import CancelToken
 
+from rlp import sedes
+
 from eth_utils import (
     keccak,
     int_to_big_endian,
+    to_tuple,
 )
 
-from eth_keys import datatypes
 from eth_keys import keys
 
 from p2p import auth
 from p2p import discovery
 from p2p.abc import AddressAPI, NodeAPI, TransportAPI
 from p2p.ecies import generate_privkey
 from p2p.kademlia import Node, Address
+from p2p.protocol import Command, Protocol
 from p2p.transport import Transport
 
-from p2p.tools.memory_transport import MemoryTransport
 from p2p.tools.asyncio_streams import get_directly_connected_streams
+from p2p.tools.memory_transport import MemoryTransport
 
 
 try:
@@ -175,9 +179,9 @@ async def establish_transport() -> None:
 
 
 def MemoryTransportPairFactory(alice_remote: NodeAPI = None,
-                               alice_private_key: datatypes.PrivateKey = None,
+                               alice_private_key: keys.PrivateKey = None,
                                bob_remote: NodeAPI = None,
-                               bob_private_key: datatypes.PrivateKey = None,
+                               bob_private_key: keys.PrivateKey = None,
                                ) -> Tuple[TransportAPI, TransportAPI]:
     if alice_remote is None:
         alice_remote = NodeFactory()
@@ -196,3 +200,138 @@ def MemoryTransportPairFactory(alice_remote: NodeAPI = None,
         bob=(alice_remote, bob_private_key),
     )
     return alice_transport, bob_transport
+
+
+STRUCTURE_SEDES = (
+    sedes.big_endian_int,
+    sedes.binary,
+)
+
+
+@to_tuple
+def StructureFactory(high_water_mark: int = 4,
+                     ) -> Iterable[Tuple[str, Any]]:
+    for idx in range(high_water_mark):
+        name = f"field_{idx}"
+        sedes = random.choice(STRUCTURE_SEDES)
+        yield (name, sedes)
+        if random.randrange(idx, high_water_mark + 2) >= high_water_mark:
+            break
+
+
+ACTIONS = (
+    'dig',
+    'run',
+    'jump',
+    'create',
+    'destroy',
+    'fill',
+    'build',
+    'create',
+    'kill',
+    'finish',
+    'hello',
+    'goodbye',
+    'connect',
+    'disconnect',
+    'activate',
+    'disable',
+    'enable',
+    'validate',
+    'post',
+    'get',
+)
+
+
+ANIMALS = (
+    'dog',
+    'cat',
+    'bird',
+    'fox',
+    'panda',
+    'unicorn',
+    'bear',
+    'eagle',
+)
+
+
+COLORS = (
+    'red',
+    'orange',
+    'yellow',
+    'green',
+    'blue',
+    'purple',
+    'pink',
+    'brown',
+    'black',
+    'white',
+)
+
+
+def _command_name_enumerator() -> Generator[str, None, None]:
+    while True:
+        for action in ACTIONS:
+            yield action.title()
+        for action, animal in itertools.product(ACTIONS, ANIMALS):
+            yield f"{action.title()}{animal.title()}"
+
+
+_command_name_iter = _command_name_enumerator()
+
+
+def CommandNameFactory() -> str:
+    return next(_command_name_iter)
+
+
+def CommandFactory(name: str = None,
+                   cmd_id: int = None,
+                   structure: Tuple[Tuple[str, Any], ...] = None) -> Type[Command]:
+    if structure is None:
+        structure = StructureFactory()
+    if cmd_id is None:
+        cmd_id = 0
+    if name is None:
+        name = CommandNameFactory()
+
+    return type(
+        name,
+        (Command,),
+        {'_cmd_id': cmd_id, 'structure': structure},
+    )
+
+
+def _protocol_name_enumerator() -> Generator[str, None, None]:
+    while True:
+        for color, animal in itertools.product(COLORS, ANIMALS):
+            yield f"{color}_{animal}"
+
+
+_protocol_name_iter = _protocol_name_enumerator()
+
+
+def ProtocolNameFactory() -> str:
+    return next(_protocol_name_iter)
+
+
+def ProtocolFactory(name: str = None,
+                    version: int = None,
+                    commands: Tuple[Type[Command], ...] = None) -> Type[Protocol]:
+    if name is None:
+        name = ProtocolNameFactory()
+    if version is None:
+        version = 1
+    if commands is None:
+        num_commands = random.randrange(1, 6)
+        commands = tuple(
+            CommandFactory(cmd_id=cmd_id)
+            for cmd_id in range(num_commands)
+        )
+
+    cmd_length = len(commands)
+
+    return type(
+        name.title(),
+        (Protocol,),
+        {'name': name, 'version': version, '_commands': commands, 'cmd_length': cmd_length},
+    )