init: move from httprunner/httprunner

2026-09-05 23:47:30 +08:00 · 2025-02-05 21:32:44 +08:00
commit f4860de5ad
104 changed files with 12602 additions and 0 deletions
@@ -0,0 +1,139 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import
+
+import enum
+import json
+
+import thriftpy2
+from loguru import logger
+from thriftpy2.protocol import (
+    TBinaryProtocolFactory,
+    TCompactProtocolFactory,
+    TCyBinaryProtocolFactory,
+    TJSONProtocolFactory,
+)
+from thriftpy2.rpc import make_client
+from thriftpy2.transport import (
+    TBufferedTransportFactory,
+    TCyBufferedTransportFactory,
+    TCyFramedTransportFactory,
+    TFramedTransportFactory,
+)
+
+from httprunner.thrift.data_convertor import json2thrift, thrift2dict
+
+
+class ProtoType(enum.Enum):
+    Binary = 1
+    CyBinary = 2
+    Compact = 3
+    Json = 4
+
+
+class TransType(enum.Enum):
+    Buffered = 1
+    CyBuffered = 2
+    Framed = 3
+    CyFramed = 4
+
+
+class RequestFormat(enum.Enum):
+    json = 1
+    binary = 2
+
+
+def get_proto_factory(proto_type):
+    if proto_type == ProtoType.Binary:
+        return TBinaryProtocolFactory()
+    if proto_type == ProtoType.CyBinary:
+        return TCyBinaryProtocolFactory()
+    if proto_type == ProtoType.Compact:
+        return TCompactProtocolFactory()
+    if proto_type == ProtoType.Json:
+        return TJSONProtocolFactory()
+
+
+def get_trans_factory(trans_type):
+    if trans_type == TransType.Buffered:
+        return TBufferedTransportFactory()
+    if trans_type == TransType.CyBuffered:
+        return TCyBufferedTransportFactory()
+    if trans_type == TransType.Framed:
+        return TFramedTransportFactory()
+    if trans_type == TransType.CyFramed:
+        return TCyFramedTransportFactory()
+
+
+class ThriftClient(object):
+    def __init__(
+        self,
+        thrift_file,
+        service_name,
+        ip,
+        port,
+        include_dirs=None,
+        timeout=3000,
+        proto_type=ProtoType.CyBinary,
+        trans_type=TransType.CyBuffered,
+    ):
+        self.thrift_file = thrift_file
+        self.include_dirs = include_dirs
+        self.service_name = service_name
+        self.ip = ip
+        self.port = port
+        self.timeout = timeout
+        self.proto_type = proto_type
+        self.trans_type = trans_type
+        try:
+            logger.debug(
+                "init thrift module: thrift_file=%s, module_name=%s",
+                thrift_file,
+                str(self.service_name) + "_thrift",
+            )
+            self.thrift_module = thriftpy2.load(
+                self.thrift_file,
+                module_name=str(self.service_name) + "_thrift",
+                include_dirs=self.include_dirs,
+            )
+            self.thrift_service_obj = getattr(self.thrift_module, self.service_name)
+            logger.debug(
+                "init thrift client: service_name=%s, ip=%s, port=%s",
+                self.thrift_service_obj,
+                ip,
+                port,
+            )
+            self.client = make_client(
+                self.thrift_service_obj,
+                self.ip,
+                int(self.port),
+                timeout=self.timeout,
+                proto_factory=get_proto_factory(self.proto_type),
+                trans_factory=get_trans_factory(self.trans_type),
+            )
+        except Exception as e:
+            self.thrift_module = None
+            self.thrift_service_obj = None
+            self.client = None
+            logger.exception("init thrift module and client failed: {}".format(e))
+        finally:
+            thriftpy2.parser.parser.thrift_stack = []
+
+    def get_client(self):
+        return self.client
+
+    def send_request(self, request_data, request_method=""):
+        thrift_req_cls = getattr(
+            self.thrift_service_obj, request_method + "_args"
+        ).thrift_spec[1][2]
+        request_obj = json2thrift(json.dumps(request_data), thrift_req_cls)
+        logger.debug(
+            "send thrift request: request_method=%s, request_obj=%s",
+            request_method,
+            request_obj,
+        )
+        response_obj = getattr(self.client, request_method)(request_obj)
+        logger.debug("thrift response = %s", response_obj)
+        return thrift2dict(response_obj)
+
+    def __del__(self):
+        self.client.close()