From 4ed21ecb03c4eb9bb435e9bb4ab4255381285e5e Mon Sep 17 00:00:00 2001 From: heavyrian2012 Date: Fri, 25 Oct 2024 14:43:06 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E7=94=9F=E6=88=90=E6=B5=8B?= =?UTF-8?q?=E8=AF=95=E5=B7=A5=E5=85=B7=E7=9A=84=E5=B7=A5=E5=85=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../cn/wildfirechat/sdk/GenerateTestData.java | 315 ++++++++++++++++++ 1 file changed, 315 insertions(+) create mode 100644 sdk/src/main/java/cn/wildfirechat/sdk/GenerateTestData.java diff --git a/sdk/src/main/java/cn/wildfirechat/sdk/GenerateTestData.java b/sdk/src/main/java/cn/wildfirechat/sdk/GenerateTestData.java new file mode 100644 index 00000000..2e0e2041 --- /dev/null +++ b/sdk/src/main/java/cn/wildfirechat/sdk/GenerateTestData.java @@ -0,0 +1,315 @@ +package cn.wildfirechat.sdk; + +import cn.wildfirechat.common.ErrorCode; +import cn.wildfirechat.pojos.*; +import cn.wildfirechat.sdk.model.IMResult; + +import java.io.BufferedWriter; +import java.io.FileWriter; +import java.io.IOException; +import java.util.*; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiConsumer; + +/** + * 生成测试数据。本应用生成以下数据： + * 1. 创建${totalUserCount}个用户 + * 2. 每个用户同其他${friendCount}个用户建立好友关系 + * 3. 确保每个人都加入${smallGroupCount}个群组成员数为${smallGroupMemberSize}的小群组 + * 4. 确保每个人都加入${middleGroupCount}个群组成员数为${middleGroupMemberSize}的中群组 + * 5. 确保每个人都加入${bigGroupCount}个群组成员数为${bigGroupMemberSize}的大群组 + * 6. 把数据分片${splitFileCount}个保存在文件中 + * + * 好友关系和群组关系都是随机生成，有可能会出现最后几个用户无法满足精确好友数或者群组数的问题，但绝大部分数据都能满足需求。 + */ +public class GenerateTestData { + //IM服务Server API + private static String AdminUrl = "http://localhost:18080"; + private static String AdminSecret = "123456"; + + //系统内总用户数 + private static int totalUserCount = 1000; + + //每个用户的好友数 + private static int friendCount = 5; + + //每个用户拥有的小群数，和小群的大小。 + private static int smallGroupCount = 10; + private static int smallGroupMemberSize = 20; + + //每个用户拥有的中群数，和中群的大小。 + private static int middleGroupCount = 5; + private static int middleGroupMemberSize = 40; + + //每个用户拥有的大群数，和大群的大小。 + private static int bigGroupCount = 2; + private static int bigGroupMemberSize = 200; + + //数据文件分割成几个文件存储。 + private static int splitFileCount = 5; + + //统计添加好友数和创建群组数。 + private static int addFriendCount = 0; + private static AtomicInteger createGroupCount = new AtomicInteger(0); + + public static void main(String[] args) throws Exception { + AdminConfig.initAdmin(AdminUrl, AdminSecret); + int addFriendRequestCount = friendCount*totalUserCount/2; + int createSmallGroupCount = smallGroupCount * totalUserCount / smallGroupMemberSize; + int createMiddleGroupCount = middleGroupCount * totalUserCount / middleGroupMemberSize; + int createBigGroupCount = bigGroupCount * totalUserCount / bigGroupMemberSize; + System.out.println("根据当前的参数，将创建 " + totalUserCount + " 名用户"); + System.out.println("添加大概 " + addFriendRequestCount + " 次好友请求"); + System.out.println("创建大概 " + createSmallGroupCount + " 个小型群组"); + System.out.println("创建大概 " + createMiddleGroupCount + " 个中等群组"); + System.out.println("创建大概 " + createBigGroupCount + " 个大型群组"); + Scanner scanner = new Scanner(System.in); + + System.out.print("是否确认？ (y/其他): "); + String input = scanner.nextLine().trim().toUpperCase(); + + if (!input.equalsIgnoreCase("y")) { + System.exit(0); + } + scanner.close(); + + List userIds = createUsers(totalUserCount); + CountDownLatch latch = new CountDownLatch(4); + + Map unfulfilledFriendMap = new HashMap<>(); + new Thread(() -> { + addFriends(userIds, friendCount, unfulfilledFriendMap); + latch.countDown(); + }).start(); + + Map> userSmallGroupMap = new HashMap<>(); + Map unfulfilledSmallMap = new HashMap<>(); + new Thread(() ->{ + for (String userId : userIds) { + createGroup(userId, userIds, userSmallGroupMap, smallGroupCount, smallGroupMemberSize, unfulfilledSmallMap); + } + latch.countDown(); + }).start(); + + Map> userMiddleGroupMap = new HashMap<>(); + Map unfulfilledMiddleMap = new HashMap<>(); + new Thread(() ->{ + for (String userId : userIds) { + createGroup(userId, userIds, userMiddleGroupMap, middleGroupCount, middleGroupMemberSize, unfulfilledMiddleMap); + } + latch.countDown(); + }).start(); + + Map> userBigGroupMap = new HashMap<>(); + Map unfulfilledBigMap = new HashMap<>(); + new Thread(() ->{ + for (String userId : userIds) { + createGroup(userId, userIds, userBigGroupMap, bigGroupCount, bigGroupMemberSize, unfulfilledBigMap); + } + latch.countDown(); + }).start(); + + latch.await(); + + System.out.println("数据初始化结束，实际建立好友关系: " + addFriendCount + " 条。"); + if (!unfulfilledFriendMap.isEmpty()) { + System.out.println("有部分用户没有建立" + friendCount + "个好友关系，分别是："); + printUnfulfilledMap(unfulfilledFriendMap); + } + + System.out.println("数据初始化结束，创建群组: " + createGroupCount.get() + " 个。"); + if (!unfulfilledSmallMap.isEmpty()) { + System.out.println("有部分用户没有建立" + smallGroupMemberSize + "个群组成员的群，分别是："); + printUnfulfilledMap(unfulfilledSmallMap); + } + if (!unfulfilledMiddleMap.isEmpty()) { + System.out.println("有部分用户没有建立" + middleGroupMemberSize + "个群组成员的群，分别是："); + printUnfulfilledMap(unfulfilledMiddleMap); + } + if (!unfulfilledBigMap.isEmpty()) { + System.out.println("有部分用户没有建立" + bigGroupMemberSize + "个群组成员的群，分别是："); + printUnfulfilledMap(unfulfilledBigMap); + } + + + System.out.println("保存数据到" + splitFileCount + "个文件中。"); + for (int i = 0; i < splitFileCount; i++) { + String filePath = "data"+i+".csv"; // 文件路径 + BufferedWriter writer = new BufferedWriter(new FileWriter(filePath)); + int partSize = totalUserCount/splitFileCount; + for (int j = 0; j < partSize; j++) { + String userId = userIds.get(i*partSize+j); + writeDeviceInfo(writer, userId, 3); + Set smallGroupIds = userSmallGroupMap.get(userId); + writeGroupIds(writer, smallGroupIds); + Set middleGroupIds = userMiddleGroupMap.get(userId); + writeGroupIds(writer, middleGroupIds); + Set bigGroupIds = userBigGroupMap.get(userId); + writeGroupIds(writer, bigGroupIds); + } + + writer.close(); + } + + System.out.println("保存结束。"); + } + + private static void printUnfulfilledMap(Map map) { + if(map.isEmpty()) + return; + + map.forEach((s, integer) -> System.out.print(s + ":" + integer)); + System.out.println(); + } + + private static void writeDeviceInfo(BufferedWriter writer, String userId, int platform) throws IOException { + writer.write(userId); + writer.write(","); + writer.write(UUID.randomUUID().toString()); + writer.write(","); + writer.write(""+platform); + writer.write("\n"); + } + + private static void writeGroupIds(BufferedWriter writer, Set groupIds) throws IOException { + int writeCount = 0; + for (String bigGroupId : groupIds) { + writer.write(bigGroupId); + writeCount++; + if(writeCount < groupIds.size()) { + writer.write(","); + } else { + writer.write("\n"); + } + } + } + private static void createGroup(String userId, List userIds, Map> userGroupMap, int groupCount, int memberSize, Map unfulfilledUserMap) { + Set groupIds = userGroupMap.computeIfAbsent(userId, s -> new HashSet<>()); + int i = 0; + while (groupIds.size() < groupCount){ + PojoGroupInfo groupInfo = new PojoGroupInfo(); + groupInfo.setTarget_id(UUID.randomUUID().toString()); + groupInfo.setOwner(userId); + groupInfo.setName("Group_"+userId+i++); + groupInfo.setType(2); + groupIds.add(groupInfo.getTarget_id()); + + List memberCandidates = new ArrayList<>(userIds); + List memberIds = new ArrayList<>(); + memberIds.add(userId); + + while (!memberCandidates.isEmpty() && memberIds.size() < memberSize) { + String memberId = memberCandidates.get((int)(Math.random()*memberCandidates.size())); + memberCandidates.remove(memberId); + if(memberId.equals(userId)) { + continue; + } + if (memberIds.contains(memberId)) { + continue; + } + Set memberGroupIds = userGroupMap.computeIfAbsent(memberId, s -> new HashSet<>()); + if(memberGroupIds.contains(groupInfo.getTarget_id()) || memberGroupIds.size() >= groupCount) { + continue; + } + memberGroupIds.add(groupInfo.getTarget_id()); + memberIds.add(memberId); + } + + List members = new ArrayList<>(); + for (String memberId : memberIds) { + PojoGroupMember member1 = new PojoGroupMember(); + member1.setMember_id(memberId); + members.add(member1); + } + + try { + IMResult resultCreateGroup = GroupAdmin.createGroup(groupInfo.getOwner(), groupInfo, members, null, null, null); + if (resultCreateGroup != null && resultCreateGroup.getErrorCode() == ErrorCode.ERROR_CODE_SUCCESS) { + System.out.println("create group success"); + } else { + System.out.println("create group failure"); + System.exit(-1); + } + } catch (Exception e) { + System.out.println("create group failure:" + e); + System.exit(-1); + } + + createGroupCount.incrementAndGet(); + if(memberIds.size() < memberSize) { + unfulfilledUserMap.put(userId + ",s group " + groupInfo.getTarget_id(), memberIds.size()); + System.out.println("Error, can not find enough user to create group for user:" + userId); + } + } + } + + private static void addFriends(List userIds, int count, Map unfulfilledFriendMap) { + Map> userFriends = new HashMap<>(); + for (String userId : userIds) { + List candidates = new ArrayList<>(userIds); + Set currentUserFriends = userFriends.computeIfAbsent(userId, s -> new HashSet<>()); + while (!candidates.isEmpty() && currentUserFriends.size() < count) { + String friendId = candidates.get((int) (Math.random()*candidates.size())); + candidates.remove(friendId); + if(friendId.equals(userId)) { + continue; + } + + if(!currentUserFriends.contains(friendId)) { + Set targetUserFriends = userFriends.computeIfAbsent(friendId, s -> new HashSet<>()); + if(targetUserFriends.size() < count) { + currentUserFriends.add(friendId); + targetUserFriends.add(userId); + try { + IMResult result = RelationAdmin.setUserFriend(userId, friendId, true, null); + if (result != null && (result.getErrorCode() == ErrorCode.ERROR_CODE_SUCCESS || result.getErrorCode() == ErrorCode.ERROR_CODE_ALREADY_FRIENDS)) { + System.out.println("add friend success"); + addFriendCount++; + } else { + System.out.println("failure"); + System.exit(-1); + } + } catch (Exception e) { + System.out.println("failure:" + e); + System.exit(-1); + } + } + } + } + if (currentUserFriends.size() < count) { + unfulfilledFriendMap.put(userId, currentUserFriends.size()); + System.out.println("Error, can not find enough user to add friend for user:" + userId); + } + } + } + + private static List createUsers(int count) throws Exception { + List userIds = new ArrayList<>(); + for (int i = 0; i < count; i++) { + InputOutputUserInfo userInfo = generateUserInfo(i); + IMResult resultCreateUser = UserAdmin.createUser(userInfo); + if (resultCreateUser != null && resultCreateUser.getErrorCode() == ErrorCode.ERROR_CODE_SUCCESS) { + System.out.println("Create user " + resultCreateUser.getResult().getName() + " success"); + } else { + System.out.println("Create user failure"); + System.exit(-1); + } + userIds.add(userInfo.getUserId()); + } + return userIds; + } + + private static InputOutputUserInfo generateUserInfo(int index) { + InputOutputUserInfo userInfo = new InputOutputUserInfo(); + String userId = UUID.randomUUID().toString().replace("-", ""); + //用户ID，必须保证唯一性 + userInfo.setUserId(userId); + //用户名，一般是用户登录帐号，也必须保证唯一性。也就是说所有用户的userId必须不能重复，所有用户的name必须不能重复，但可以同一个用户的userId和name是同一个，一般建议userId使用一个uuid，name是"微信号"且可以修改， + userInfo.setName(userId); + userInfo.setMobile((11000000000L+index) + ""); + userInfo.setDisplayName("User"+index); + + return userInfo; + } +}