github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/election/storage_sql_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 election
    15  
    16  import (
    17  	"context"
    18  	"encoding/json"
    19  	"regexp"
    20  	"testing"
    21  
    22  	"github.com/DATA-DOG/go-sqlmock"
    23  	"github.com/pingcap/tiflow/pkg/errors"
    24  	"github.com/stretchr/testify/require"
    25  )
    26  
    27  func newSQLStorageAndMock(t *testing.T) (*SQLStorage, sqlmock.Sqlmock) {
    28  	db, mock, err := sqlmock.New()
    29  	require.NoError(t, err)
    30  
    31  	mock.ExpectExec(regexp.QuoteMeta("CREATE TABLE IF NOT EXISTS leader_election " +
    32  		"(id int NOT NULL, version bigint NOT NULL, record text NOT NULL, PRIMARY KEY (id))")).
    33  		WillReturnResult(sqlmock.NewResult(0, 0))
    34  
    35  	s, err := NewSQLStorage(db, "leader_election")
    36  	require.NoError(t, err)
    37  
    38  	return s, mock
    39  }
    40  
    41  func TestSQLStorageGetEmptyRecord(t *testing.T) {
    42  	t.Parallel()
    43  
    44  	s, mock := newSQLStorageAndMock(t)
    45  
    46  	mock.ExpectQuery(regexp.QuoteMeta("SELECT version, record FROM leader_election WHERE id = ?")).
    47  		WithArgs(1).WillReturnRows(sqlmock.NewRows([]string{"version", "record"}))
    48  	record, err := s.Get(context.Background())
    49  	require.NoError(t, err)
    50  	require.Equal(t, &Record{}, record)
    51  }
    52  
    53  func TestSQLStorageGetExistingRecord(t *testing.T) {
    54  	t.Parallel()
    55  
    56  	s, mock := newSQLStorageAndMock(t)
    57  
    58  	expectedRecord := &Record{
    59  		LeaderID: "id1",
    60  		Members: []*Member{
    61  			{
    62  				ID:   "id1",
    63  				Name: "name1",
    64  			},
    65  			{
    66  				ID:   "id2",
    67  				Name: "name2",
    68  			},
    69  		},
    70  		Version: 1,
    71  	}
    72  	recordBytes, err := json.Marshal(expectedRecord)
    73  	require.NoError(t, err)
    74  
    75  	mock.ExpectQuery(regexp.QuoteMeta("SELECT version, record FROM leader_election WHERE id = ?")).
    76  		WithArgs(1).WillReturnRows(sqlmock.NewRows([]string{"version", "record"}).AddRow(1, recordBytes))
    77  	record, err := s.Get(context.Background())
    78  	require.NoError(t, err)
    79  	require.Equal(t, expectedRecord, record)
    80  }
    81  
    82  func TestSQLStorageInsertRecord(t *testing.T) {
    83  	t.Parallel()
    84  
    85  	s, mock := newSQLStorageAndMock(t)
    86  
    87  	record := &Record{
    88  		LeaderID: "id1",
    89  		Members: []*Member{
    90  			{
    91  				ID:   "id1",
    92  				Name: "name1",
    93  			},
    94  		},
    95  		Version: 0, // 0 means record not created before.
    96  	}
    97  	recordBytes, err := json.Marshal(record)
    98  	require.NoError(t, err)
    99  
   100  	mock.ExpectExec(regexp.QuoteMeta("INSERT INTO leader_election (id, version, record) VALUES (?, ?, ?)")).
   101  		WithArgs(1, int64(1), recordBytes).WillReturnResult(sqlmock.NewResult(1, 1))
   102  
   103  	err = s.Update(context.Background(), record, true)
   104  	require.NoError(t, err)
   105  }
   106  
   107  func TestSQLStorageUpdateRecord(t *testing.T) {
   108  	t.Parallel()
   109  
   110  	s, mock := newSQLStorageAndMock(t)
   111  
   112  	record := &Record{
   113  		LeaderID: "id1",
   114  		Members: []*Member{
   115  			{
   116  				ID:   "id1",
   117  				Name: "name1",
   118  			},
   119  		},
   120  		Version: 1,
   121  	}
   122  	recordBytes, err := json.Marshal(record)
   123  	require.NoError(t, err)
   124  
   125  	mock.ExpectExec(regexp.QuoteMeta("UPDATE leader_election SET version = ?, record = ? WHERE id = ? AND version = ?")).
   126  		WithArgs(int64(2), recordBytes, 1, int64(1)).WillReturnResult(sqlmock.NewResult(0, 0))
   127  	err = s.Update(context.Background(), record, true)
   128  	require.True(t, errors.Is(err, errors.ErrElectionRecordConflict))
   129  
   130  	mock.ExpectExec(regexp.QuoteMeta("UPDATE leader_election SET version = ?, record = ? WHERE id = ? AND version = ?")).
   131  		WithArgs(int64(2), recordBytes, 1, int64(1)).WillReturnResult(sqlmock.NewResult(0, 1))
   132  	err = s.Update(context.Background(), record, true)
   133  	require.NoError(t, err)
   134  }
   135  
   136  func TestInMemorySQLStorage(t *testing.T) {
   137  	t.Parallel()
   138  
   139  	dbName := t.TempDir()
   140  	s, err := NewInMemorySQLStorage(dbName, "leader_election")
   141  	require.NoError(t, err)
   142  
   143  	ctx, cancel := context.WithCancel(context.Background())
   144  	defer cancel()
   145  
   146  	record := &Record{
   147  		LeaderID: "id1",
   148  		Members: []*Member{
   149  			{
   150  				ID:   "id1",
   151  				Name: "name1",
   152  			},
   153  		},
   154  		Version: 0, // 0 means record not created before.
   155  	}
   156  	err = s.Update(ctx, record, true)
   157  	require.NoError(t, err)
   158  
   159  	recordRead, err := s.Get(ctx)
   160  	require.NoError(t, err)
   161  	require.Equal(t, record.Members, recordRead.Members)
   162  	require.Equal(t, record.LeaderID, recordRead.LeaderID)
   163  	require.Equal(t, record.Version+1, recordRead.Version)
   164  }