RocketMQ 读写测试

https://github.com/apache/rocketmq-client-python

# -*- coding:utf-8 -*-
import time
import uuid
import random
import socket
import traceback
from datetime import datetime

from rocketmq.client import Producer, Message, PullConsumer, PushConsumer
from rocketmq import exceptions as MQException


class RocketmqProbe():

    def __init__(self):
        self._random_tag = str(uuid.uuid4()) # 用于过滤本次生产的消息和消费
        self.random_tag = self._random_tag


    def port_check(self, service: dict, **fields_map) -> bool:
        """
        端口检测
        :param service: 服务信息
        :param fields_map: 字段映射
        :return:
        """
        ip = service.get(fields_map.get("ip"))
        port = service.get(fields_map.get("port"))
        try:
            pong_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            pong_socket.settimeout(2)
            pong_code = pong_socket.connect_ex((ip, int(port)))
        except Exception as e:
            pass
        finally:
            pong_socket.close()

        return pong_code == 0


    @property
    def get_now_time(self):
        return datetime.now().strftime('%Y%m%d%H%M%S')


    def producer_msg(self, nameserver: str, topic: str, count: int):
        """
        发送消息

        :param nameserver: nameserver地址
        :param topic: 测试的topic
        :param count: 发送消息的个数
        :return:
        """
        send_count = 0  # 发送计数
        try:
            producer = Producer("test_probe_producer")
            producer.set_namesrv_addr(nameserver)
            producer.start()
            for i in range(count):
                msg = Message(topic)
                key = f"{
      
      topic}_key_{
      
      int(time.time())}"
                print(key)
                msg.set_keys(key)
                msg.set_tags(self.random_tag)
                body = {
    
    
                    "time": self.get_now_time,
                    "index": i,
                    "data": f"{
      
      topic}_monitor_cluster_producer_consumer_health_aaronyu"
                }
                msg.set_body(str(body).encode('utf-8'))
                ret = producer.send_sync(msg)
                send_count += 1
            producer.shutdown()
        except Exception as e:
            print(e)


    def consumer_callback(self, msg):
        """
        push consumer callback func
        :param msg:
        :return:
        """
        print("---> %s, consumer_callback" %(self.get_now_time))


    def pull_consume_msg(self, nameserver: str, topic: str, consumer_group: str):
        """
        指定消费组消费指定消息

        :param nameserver: nameserver地址
        :param topic: topic名称
        :param consumer_group: 消费组名称
        :return:
        """
        try:
            consumer = PullConsumer(consumer_group)
            consumer.set_namesrv_addr(nameserver)
            consumer.start()
            this_consume_msg = []
            for msg in consumer.pull(topic):
                print("-----> %s, %s" %(msg.tags, msg.body))
                if str(msg.tags.decode("utf-8")) == self.random_tag:
                    this_consume_msg.append(msg)
        except Exception as e:
            print(e)
        finally:
            consumer.shutdown()

        if len(this_consume_msg) == message_count:
            print("1 %s, %s" %(this_consume_msg, message_count))
        else:
            print("2 %s, %s" %(this_consume_msg, message_count))


    def push_consume_msg(self, nameserver: str, topic: str, consumer_group: str):
        """
        指定消费组消费指定消息

        :param nameserver: nameserver地址
        :param topic: topic名称
        :param consumer_group: 消费组名称
        :return:
        """
        try:
            consumer = PushConsumer(consumer_group)
            consumer.set_namesrv_addr(nameserver)
            consumer.subscribe(topic, self.consumer_callback)
            consumer.start()
            while True:
                time.sleep(15)
        except Exception as e:
            print(e)
        finally:
            consumer.shutdown()


if __name__ == '__main__':
    rq = RocketmqProbe()
    
    # 需要修改的变量（开始）
    ## NameServer 地址
    nameserver = "10.9.22.XXX:9876"
    ## topic 名称
    topic_name = "topic_test_20221206"
    # 需要修改的变量（结束）

    message_count = 3

    # 生产者测试
    print("producer begin")
    producer_code = rq.producer_msg(nameserver, topic_name, message_count)
    print(producer_code)
    print("producer end")



	## 消费者测试
    print("consumer begin")
    # 可重复消费
    consumer_code = rq.pull_consume_msg(nameserver, topic_name, "TEST_MQ_CONSUMER")
    # 不可重复消费
    # consumer_code = rq.push_consume_msg("10.11.23.226:9876", cf.probe_topic, cf.probe_consumer)
    print(consumer_code)
    print("consumer end")