testsuites/vstf/vstf_scripts/vstf/common/message.py

   1 ##############################################################################
   2 # Copyright (c) 2015 Huawei Technologies Co.,Ltd and others.
   3 #
   4 # All rights reserved. This program and the accompanying materials
   5 # are made available under the terms of the Apache License, Version 2.0
   6 # which accompanies this distribution, and is available at
   7 # http://www.apache.org/licenses/LICENSE-2.0
   8 ##############################################################################
   9
  10 import json
  11 import uuid
  12 import logging
  13 import traceback
  14 from vstf.common import constants
  15
  16 LOG = logging.getLogger(__name__)
  17
  18
  19 def json_defaults(obj):
  20     if isinstance(obj, set):
  21         return list(obj)
  22     return "unknow obj"
  23
  24
  25 def encode(msg):
  26     """obj to string"""
  27     if isinstance(msg, str):
  28         return msg
  29     else:
  30         return json.dumps(msg, default=json_defaults)
  31
  32
  33 def decode(msg):
  34     """string to obj"""
  35     if isinstance(msg, str):
  36         return json.loads(msg)
  37     else:
  38         return msg
  39
  40
  41 def gen_corrid():
  42     return str(uuid.uuid4())
  43
  44
  45 def add_context(msg, **kwargs):
  46     return {'head': kwargs, 'body': msg}
  47
  48
  49 def get_context(msg):
  50     if "head" in msg.iterkeys():
  51         return msg['head']
  52     else:
  53         return ""
  54
  55
  56 def get_body(msg):
  57     if "body" in msg.iterkeys():
  58         return msg['body']
  59     else:
  60         return None
  61
  62
  63 def get_corrid(context):
  64     """
  65     :param return: string of corrid or empty
  66     """
  67     if "corrid" in context.iterkeys():
  68         return context['corrid']
  69     else:
  70         return ""
  71
  72
  73 def send(func, data):
  74     # the message must be a string
  75     if not isinstance(data, str):
  76         raise ValueError("the data must be a string")
  77
  78     # the message's len must > 0
  79     msg_len = len(data)
  80     if msg_len <= 0:
  81         return True
  82
  83     # the message's len must be less 999999999
  84     if len(str(msg_len)) > constants.MSG_FLAG_LEN:
  85         raise ValueError("the data's len too long")
  86
  87     data = (constants.MSG_FLAG % (msg_len)) + data
  88     total_send = msg_len + constants.MSG_FLAG_LEN
  89
  90     count = 0
  91     while count < total_send:
  92         sent = func(data[count:])
  93         if 0 == sent:
  94             raise RuntimeError("socket connection broken")
  95         count += sent
  96
  97     return msg_len
  98
  99
 100 def sendto(func, data, addr):
 101     # the message must be a string
 102     if not isinstance(data, str):
 103         raise ValueError("the data must be a string")
 104
 105     # the message's len must > 0
 106     msg_len = len(data)
 107     if msg_len <= 0:
 108         return True
 109
 110     # the message's len must be less 999999999
 111     if len(str(msg_len)) > constants.MSG_FLAG_LEN:
 112         raise ValueError("the data's len too long")
 113
 114     data = (constants.MSG_FLAG % (msg_len)) + data
 115     total_send = msg_len + constants.MSG_FLAG_LEN
 116
 117     count = 0
 118     while count < total_send:
 119         sent = func(data[count:], addr)
 120         if 0 == sent:
 121             raise RuntimeError("socket connection broken")
 122         count += sent
 123
 124     return msg_len
 125
 126
 127 def recv(func):
 128     head = func(constants.MSG_FLAG_LEN)
 129     # the FIN change to '' in python
 130     if head == '':
 131         raise RuntimeError("socket connection broken")
 132
 133     if not head.isdigit():
 134         raise ValueError("the msg head is not a num.")
 135
 136     msg_len = int(head)
 137     chunks = []
 138     count = 0
 139     while count < msg_len:
 140         chunk = func(min(msg_len - count, constants.buff_size))
 141         if chunk == '':
 142             raise RuntimeError("socket connection broken")
 143         chunks.append(chunk)
 144         count += len(chunk)
 145
 146     return ''.join(chunks)
 147
 148
 149 def dumpstrace():
 150     return traceback.format_exc()