allocate_test.go 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. package balancer
  2. import (
  3. cmap "github.com/orcaman/concurrent-map/v2"
  4. "github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
  5. "reflect"
  6. "testing"
  7. )
  8. func Test_allocateOneBroker(t *testing.T) {
  9. brokers := cmap.New[*BrokerStats]()
  10. brokers.SetIfAbsent("localhost:17777", &BrokerStats{
  11. TopicPartitionCount: 0,
  12. ConsumerCount: 0,
  13. CpuUsagePercent: 0,
  14. })
  15. tests := []struct {
  16. name string
  17. args args
  18. wantAssignments []*mq_pb.BrokerPartitionAssignment
  19. }{
  20. {
  21. name: "test only one broker",
  22. args: args{
  23. brokers: brokers,
  24. partitionCount: 1,
  25. },
  26. wantAssignments: []*mq_pb.BrokerPartitionAssignment{
  27. {
  28. LeaderBroker: "localhost:17777",
  29. Partition: &mq_pb.Partition{
  30. RingSize: MaxPartitionCount,
  31. RangeStart: 0,
  32. RangeStop: MaxPartitionCount,
  33. },
  34. },
  35. },
  36. },
  37. }
  38. testThem(t, tests)
  39. }
  40. type args struct {
  41. brokers cmap.ConcurrentMap[string, *BrokerStats]
  42. partitionCount int32
  43. }
  44. func testThem(t *testing.T, tests []struct {
  45. name string
  46. args args
  47. wantAssignments []*mq_pb.BrokerPartitionAssignment
  48. }) {
  49. for _, tt := range tests {
  50. t.Run(tt.name, func(t *testing.T) {
  51. if gotAssignments := allocateTopicPartitions(tt.args.brokers, tt.args.partitionCount); !reflect.DeepEqual(gotAssignments, tt.wantAssignments) {
  52. t.Errorf("allocateTopicPartitions() = %v, want %v", gotAssignments, tt.wantAssignments)
  53. }
  54. })
  55. }
  56. }