scapy_service payload gen: allow template_code without size property

Signed-off-by: Anton Kiselev <anton.kisel@gmail.com>

2 files changed, 20 insertions, 7 deletions

diff --git a/scripts/automation/trex_control_plane/stl/services/scapy_server/scapy_service.py b/scripts/automation/trex_control_plane/stl/services/scapy_server/scapy_service.py
index f889bb93..88514aa8 100755
--- a/scripts/automation/trex_control_plane/stl/services/scapy_server/scapy_service.py
+++ b/scripts/automation/trex_control_plane/stl/services/scapy_server/scapy_service.py
@@ -309,6 +309,11 @@ def parse_template_code(template_code):
     template_code = re.sub("[\s]", '', template_code) # remove spaces
     return bytearray.fromhex(template_code)
 
+def verify_payload_size(size):
+    assert(size != None)
+    if (size > (1<<20)): # 1Mb ought to be enough for anybody
+        raise ValueError('size is too large')
+
 def generate_bytes(bytes_definition):
     # accepts a bytes definition object
     # {generate: random_bytes or random_ascii, seed: <seed_number>, size: <size_bytes>}
@@ -317,20 +322,21 @@ def generate_bytes(bytes_definition):
     gen_type = bytes_definition.get('generate')
     if gen_type == None:
         return b64_to_bytes(bytes_definition['base64'])
+    elif gen_type == 'template_code':
+        code = parse_template_code(bytes_definition["template_code"])
+        bytes_size = int(bytes_definition.get('size') or len(code))
+        verify_payload_size(bytes_size)
+        return generate_bytes_from_template(bytes_size, code)
     else:
         bytes_size = int(bytes_definition['size']) # required
         seed = int(bytes_definition.get('seed') or 12345) # optional
-        if (bytes_size > (1<<20)): # 1Mb ought to be enough for anybody
-            raise ValueError('size is too large')
+        verify_payload_size(bytes_size)
         if gen_type == 'random_bytes':
             return generate_random_bytes(bytes_size, seed, 0, 0xFF)
         elif gen_type == 'random_ascii':
             return generate_random_bytes(bytes_size, seed, 0x20, 0x7E)
         elif gen_type == 'template':
             return generate_bytes_from_template(bytes_size, b64_to_bytes(bytes_definition["template_base64"]))
-        elif gen_type == 'template_code':
-            return generate_bytes_from_template(bytes_size, parse_template_code(bytes_definition["template_code"]))
-
 
 class ScapyException(Exception): pass
 class Scapy_service(Scapy_service_api):
diff --git a/scripts/automation/trex_control_plane/stl/services/scapy_server/unit_tests/test_utils.py b/scripts/automation/trex_control_plane/stl/services/scapy_server/unit_tests/test_utils.py
index e9fbcc80..ceb88b47 100644
--- a/scripts/automation/trex_control_plane/stl/services/scapy_server/unit_tests/test_utils.py
+++ b/scripts/automation/trex_control_plane/stl/services/scapy_server/unit_tests/test_utils.py
@@ -49,9 +49,16 @@ def test_generate_ascii_default_seed():
     assert(len(res) == 14)
 
 
-def test_generate_template_code():
+def test_generate_template_code_no_size():
+    res = generate_bytes({"generate":"template_code", "template_code": "BE EF"})
+    assert(res == bytearray.fromhex('BE EF'))
+
+def test_generate_template_code_less():
+    res = generate_bytes({"generate":"template_code", "template_code": "DE AD BE EF", "size": 2})
+    assert(res == bytearray.fromhex('DE AD'))
+
+def test_generate_template_code_more():
     res = generate_bytes({"generate":"template_code", "template_code": "0xDEAD 0xBEEF", "size": 6})
-    print(res)
     assert(res == bytearray.fromhex('DE AD BE EF DE AD'))
 
 def test_generate_template_base64():