github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/cdc/sink/util/helper_test.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package util
    15  
    16  import (
    17  	"context"
    18  	"net/url"
    19  	"testing"
    20  
    21  	"github.com/pingcap/tiflow/cdc/model"
    22  	"github.com/pingcap/tiflow/pkg/sink/kafka"
    23  	"github.com/stretchr/testify/require"
    24  )
    25  
    26  func TestPartition(t *testing.T) {
    27  	t.Parallel()
    28  
    29  	adminClient := kafka.NewClusterAdminClientMockImpl()
    30  	defer adminClient.Close()
    31  
    32  	cfg := &kafka.AutoCreateTopicConfig{
    33  		AutoCreate:        true,
    34  		PartitionNum:      2,
    35  		ReplicationFactor: 1,
    36  	}
    37  
    38  	changefeedID := model.DefaultChangeFeedID("test")
    39  	ctx := context.Background()
    40  
    41  	manager, err := GetTopicManagerAndTryCreateTopic(ctx, changefeedID, kafka.DefaultMockTopicName, cfg, adminClient)
    42  	require.NoError(t, err)
    43  	defer manager.Close()
    44  
    45  	// default topic, real partition is 3, but 2 is set in the sink-uri, so return 2.
    46  	partitionsNum, err := manager.GetPartitionNum(ctx, kafka.DefaultMockTopicName)
    47  	require.NoError(t, err)
    48  	require.Equal(t, int32(2), partitionsNum)
    49  
    50  	// new topic, create it with partition number as 2.
    51  	partitionsNum, err = manager.GetPartitionNum(ctx, "new-topic")
    52  	require.NoError(t, err)
    53  	require.Equal(t, int32(2), partitionsNum)
    54  
    55  	// assume a topic already exist, the not default topic won't be affected by the default topic's partition number.
    56  	err = adminClient.CreateTopic(ctx, &kafka.TopicDetail{
    57  		Name:          "new-topic-2",
    58  		NumPartitions: 3,
    59  	}, false)
    60  	require.NoError(t, err)
    61  
    62  	partitionsNum, err = manager.GetPartitionNum(ctx, "new-topic-2")
    63  	require.NoError(t, err)
    64  	require.Equal(t, int32(3), partitionsNum)
    65  }
    66  
    67  func TestGetTopic(t *testing.T) {
    68  	t.Parallel()
    69  
    70  	testCases := map[string]struct {
    71  		sinkURI   string
    72  		wantTopic string
    73  		wantErr   string
    74  	}{
    75  		//"no topic": {
    76  		//	sinkURI:   "kafka://localhost:9092/",
    77  		//	wantTopic: "",
    78  		//	wantErr:   "no topic is specified in sink-uri",
    79  		//},
    80  		"valid topic": {
    81  			sinkURI:   "kafka://localhost:9092/test",
    82  			wantTopic: "test",
    83  			wantErr:   "",
    84  		},
    85  		"topic with query": {
    86  			sinkURI:   "kafka://localhost:9092/test?version=1.0.0",
    87  			wantTopic: "test",
    88  			wantErr:   "",
    89  		},
    90  		"topic for pulsar": {
    91  			sinkURI:   "pulsar://localhost:6650/test?version=1.0.0",
    92  			wantTopic: "test",
    93  			wantErr:   "",
    94  		},
    95  		"topic with query for pulsar": {
    96  			sinkURI:   "pulsar://localhost:6650/persistent://public/default/my-topic?version=1.0.0",
    97  			wantTopic: "persistent://public/default/my-topic",
    98  			wantErr:   "",
    99  		},
   100  		"multiple host with query for pulsar": {
   101  			sinkURI:   "pulsar://localhost:6650,127.0.0.1:26650/test?version=1.0.0",
   102  			wantTopic: "test",
   103  			wantErr:   "",
   104  		},
   105  	}
   106  
   107  	for name, tc := range testCases {
   108  		tc := tc
   109  		t.Run(name, func(t *testing.T) {
   110  			t.Parallel()
   111  			sinkURI, err := url.Parse(tc.sinkURI)
   112  			require.NoError(t, err)
   113  			topic, err := GetTopic(sinkURI)
   114  			if tc.wantErr != "" {
   115  				require.Regexp(t, tc.wantErr, err)
   116  			} else {
   117  				require.NoError(t, err)
   118  				require.Equal(t, tc.wantTopic, topic)
   119  			}
   120  		})
   121  	}
   122  }