github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/loader/lightning_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 loader
    15  
    16  import (
    17  	"testing"
    18  
    19  	"github.com/pingcap/errors"
    20  	"github.com/pingcap/tidb/br/pkg/storage"
    21  	"github.com/pingcap/tidb/pkg/lightning/common"
    22  	lcfg "github.com/pingcap/tidb/pkg/lightning/config"
    23  	"github.com/pingcap/tiflow/dm/config"
    24  	"github.com/pingcap/tiflow/dm/pkg/terror"
    25  	"github.com/prometheus/client_golang/prometheus"
    26  	"github.com/prometheus/client_golang/prometheus/promauto"
    27  	"github.com/stretchr/testify/require"
    28  )
    29  
    30  func TestSetLightningConfig(t *testing.T) {
    31  	t.Parallel()
    32  
    33  	stCfg := &config.SubTaskConfig{
    34  		LoaderConfig: config.LoaderConfig{
    35  			PoolSize: 10,
    36  		},
    37  	}
    38  	l := NewLightning(stCfg, nil, "")
    39  	cfg, err := l.getLightningConfig()
    40  	require.NoError(t, err)
    41  	require.Equal(t, stCfg.LoaderConfig.PoolSize, cfg.App.RegionConcurrency)
    42  }
    43  
    44  func TestConvertLightningError(t *testing.T) {
    45  	t.Parallel()
    46  
    47  	err := common.ErrChecksumMismatch.GenWithStackByArgs(1, 2, 3, 4, 5, 6)
    48  	converted := convertLightningError(errors.Trace(err))
    49  	require.True(t, terror.ErrLoadLightningChecksum.Equal(converted))
    50  	require.Contains(t, converted.Error(), "checksum mismatched, KV number in source files: 4, KV number in TiDB cluster: 3")
    51  }
    52  
    53  func TestGetLightiningConfig(t *testing.T) {
    54  	t.Parallel()
    55  
    56  	conf, err := GetLightningConfig(&lcfg.GlobalConfig{},
    57  		&config.SubTaskConfig{
    58  			Name:       "job123",
    59  			ExtStorage: &storage.LocalStorage{},
    60  			LoaderConfig: config.LoaderConfig{
    61  				RangeConcurrency: 32,
    62  				CompressKVPairs:  "gzip",
    63  				Analyze:          "required",
    64  			},
    65  		})
    66  	require.NoError(t, err)
    67  	require.Equal(t, lcfg.CheckpointDriverMySQL, conf.Checkpoint.Driver)
    68  	require.Equal(t, lcfg.CheckpointRemove, conf.Checkpoint.KeepAfterSuccess)
    69  	require.Contains(t, conf.Checkpoint.Schema, "job123")
    70  	require.Equal(t, 32, conf.TikvImporter.RangeConcurrency)
    71  	require.Equal(t, lcfg.CompressionGzip, conf.TikvImporter.CompressKVPairs)
    72  	require.Equal(t, lcfg.OpLevelRequired, conf.PostRestore.Analyze)
    73  	lightningDefaultQuota := lcfg.NewConfig().TikvImporter.DiskQuota
    74  	// when we don't set dm loader disk quota, it should be equal to lightning's default quota
    75  	require.Equal(t, lightningDefaultQuota, conf.TikvImporter.DiskQuota)
    76  
    77  	conf, err = GetLightningConfig(&lcfg.GlobalConfig{},
    78  		&config.SubTaskConfig{
    79  			Name: "job123",
    80  			LoaderConfig: config.LoaderConfig{
    81  				RangeConcurrency: 32,
    82  				CompressKVPairs:  "gzip",
    83  				Analyze:          "required",
    84  				Dir:              "/tmp",
    85  			},
    86  		})
    87  	require.NoError(t, err)
    88  	require.Equal(t, lcfg.CheckpointDriverFile, conf.Checkpoint.Driver)
    89  
    90  	conf, err = GetLightningConfig(&lcfg.GlobalConfig{},
    91  		&config.SubTaskConfig{
    92  			Name: "job123",
    93  			LoaderConfig: config.LoaderConfig{
    94  				RangeConcurrency: 32,
    95  				CompressKVPairs:  "gzip",
    96  				Analyze:          "required",
    97  				Dir:              "gcs://bucket/path",
    98  			},
    99  		})
   100  	require.NoError(t, err)
   101  	require.Equal(t, lcfg.CheckpointDriverMySQL, conf.Checkpoint.Driver)
   102  }
   103  
   104  func TestMetricProxies(t *testing.T) {
   105  	l := &LightningLoader{cfg: &config.SubTaskConfig{}}
   106  	l.initMetricProxies()
   107  	require.Equal(t, defaultMetricProxies, l.metricProxies)
   108  
   109  	registry := prometheus.NewRegistry()
   110  	l = &LightningLoader{cfg: &config.SubTaskConfig{MetricsFactory: promauto.With(registry)}}
   111  	l.initMetricProxies()
   112  	require.NotEqual(t, defaultMetricProxies, l.metricProxies)
   113  	l.metricProxies.loaderExitWithErrorCounter.WithLabelValues("test", "source", "false").Inc()
   114  	metricFamilies, err := registry.Gather()
   115  	require.NoError(t, err)
   116  	require.Len(t, metricFamilies, 1)
   117  	l.removeLabelValuesWithTaskInMetrics("test", "source")
   118  	metricFamilies, err = registry.Gather()
   119  	require.NoError(t, err)
   120  	require.Len(t, metricFamilies, 0)
   121  }