github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/framework/statusutil/writer_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 statusutil
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"github.com/pingcap/tiflow/engine/framework/internal/worker"
    21  	frameModel "github.com/pingcap/tiflow/engine/framework/model"
    22  	pkgOrm "github.com/pingcap/tiflow/engine/pkg/orm"
    23  	"github.com/pingcap/tiflow/engine/pkg/p2p"
    24  	"github.com/pingcap/tiflow/pkg/errors"
    25  	"github.com/stretchr/testify/require"
    26  )
    27  
    28  type writerTestSuite struct {
    29  	writer        *Writer
    30  	cli           pkgOrm.Client
    31  	messageSender *p2p.MockMessageSender
    32  	masterInfo    *worker.MockMasterInfoProvider
    33  }
    34  
    35  func newWriterTestSuite(
    36  	t *testing.T,
    37  	masterID frameModel.MasterID,
    38  	masterNode p2p.NodeID,
    39  	masterEpoch frameModel.Epoch,
    40  	workerID frameModel.WorkerID,
    41  ) *writerTestSuite {
    42  	cli, err := pkgOrm.NewMockClient()
    43  	require.NoError(t, err)
    44  
    45  	messageSender := p2p.NewMockMessageSender()
    46  	masterInfo := worker.NewMockMasterInfoProvider(masterID, masterNode, masterEpoch)
    47  	return &writerTestSuite{
    48  		writer:        NewWriter(cli, messageSender, masterInfo, workerID),
    49  		cli:           cli,
    50  		messageSender: messageSender,
    51  		masterInfo:    masterInfo,
    52  	}
    53  }
    54  
    55  func (s *writerTestSuite) Close() {
    56  	_ = s.cli.Close()
    57  }
    58  
    59  func TestWriterUpdate(t *testing.T) {
    60  	suite := newWriterTestSuite(t, "master-1", "executor-1", 1, "worker-1")
    61  	defer suite.Close()
    62  	ctx := context.Background()
    63  
    64  	st := &frameModel.WorkerStatus{
    65  		JobID:    "master-1",
    66  		ID:       "worker-1",
    67  		State:    frameModel.WorkerStateNormal,
    68  		ErrorMsg: "test",
    69  	}
    70  
    71  	err := suite.cli.UpsertWorker(ctx, st)
    72  	require.NoError(t, err)
    73  
    74  	st, err = suite.cli.GetWorkerByID(ctx, st.JobID, st.ID)
    75  	require.NoError(t, err)
    76  
    77  	err = suite.writer.UpdateStatus(ctx, st)
    78  	require.NoError(t, err)
    79  
    80  	status, err := suite.cli.GetWorkerByID(ctx, st.JobID, st.ID)
    81  	require.NoError(t, err)
    82  	require.Equal(t, status.State, frameModel.WorkerStateNormal)
    83  	require.Equal(t, status.ErrorMsg, "test")
    84  
    85  	rawMsg, ok := suite.messageSender.TryPop("executor-1", WorkerStatusTopic("master-1"))
    86  	require.True(t, ok)
    87  	msg := rawMsg.(*WorkerStatusMessage)
    88  	checkWorkerStatusMsg(t, &WorkerStatusMessage{
    89  		Worker:      "worker-1",
    90  		MasterEpoch: 1,
    91  		Status:      st,
    92  	}, msg)
    93  
    94  	// Deletes the persisted status for testing purpose.
    95  	// TODO make a better mock KV that can inspect calls.
    96  	_, err = suite.cli.DeleteWorker(ctx, st.JobID, st.ID)
    97  	require.NoError(t, err)
    98  
    99  	// Repeated update. Should have a notification too, but no persistence.
   100  	err = suite.writer.UpdateStatus(ctx, st)
   101  	require.NoError(t, err)
   102  	_, ok = suite.messageSender.TryPop("executor-1", WorkerStatusTopic("master-1"))
   103  	require.True(t, ok)
   104  	msg = rawMsg.(*WorkerStatusMessage)
   105  	checkWorkerStatusMsg(t, &WorkerStatusMessage{
   106  		Worker:      "worker-1",
   107  		MasterEpoch: 1,
   108  		Status:      st,
   109  	}, msg)
   110  	_, err = suite.cli.GetWorkerByID(ctx, st.JobID, st.ID)
   111  	require.Error(t, err)
   112  }
   113  
   114  func TestWriterSendRetry(t *testing.T) {
   115  	suite := newWriterTestSuite(t, "master-1", "executor-1", 1, "worker-1")
   116  	defer suite.Close()
   117  	ctx := context.Background()
   118  
   119  	st := &frameModel.WorkerStatus{
   120  		JobID:    "master-1",
   121  		ID:       "worker-1",
   122  		State:    frameModel.WorkerStateNormal,
   123  		ErrorMsg: "test",
   124  	}
   125  	err := suite.cli.UpsertWorker(ctx, st)
   126  	require.NoError(t, err)
   127  
   128  	st, err = suite.cli.GetWorkerByID(ctx, st.JobID, st.ID)
   129  	require.NoError(t, err)
   130  
   131  	require.Equal(t, 0, suite.masterInfo.RefreshCount())
   132  	suite.messageSender.InjectError(errors.ErrExecutorNotFoundForMessage.GenWithStackByArgs())
   133  	err = suite.writer.UpdateStatus(ctx, st)
   134  	require.NoError(t, err)
   135  	require.Equal(t, 1, suite.masterInfo.RefreshCount())
   136  
   137  	rawMsg, ok := suite.messageSender.TryPop("executor-1", WorkerStatusTopic("master-1"))
   138  	require.True(t, ok)
   139  	msg := rawMsg.(*WorkerStatusMessage)
   140  	checkWorkerStatusMsg(t, &WorkerStatusMessage{
   141  		Worker:      "worker-1",
   142  		MasterEpoch: 1,
   143  		Status:      st,
   144  	}, msg)
   145  }
   146  
   147  func checkWorkerStatusMsg(t *testing.T, expect, msg *WorkerStatusMessage) {
   148  	require.Equal(t, expect.Worker, msg.Worker)
   149  	require.Equal(t, expect.MasterEpoch, msg.MasterEpoch)
   150  	require.Equal(t, expect.Status.State, expect.Status.State)
   151  	require.Equal(t, expect.Status.ErrorMsg, expect.Status.ErrorMsg)
   152  	require.Equal(t, expect.Status.ExtBytes, expect.Status.ExtBytes)
   153  }