github.com/matrixorigin/matrixone@v1.2.0/pkg/util/export/etl/db/db_holder_test.go (about)

     1  // Copyright 2022 Matrix Origin
     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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package db_holder
    16  
    17  import (
    18  	"context"
    19  	"database/sql"
    20  	"reflect"
    21  	"regexp"
    22  	"testing"
    23  
    24  	"github.com/matrixorigin/matrixone/pkg/util/export/table"
    25  
    26  	"github.com/DATA-DOG/go-sqlmock"
    27  	"github.com/stretchr/testify/require"
    28  )
    29  
    30  func TestBulkInsert(t *testing.T) {
    31  
    32  	tbl := &table.Table{
    33  		Account:  "test",
    34  		Database: "testDB",
    35  		Table:    "testTable",
    36  		Columns: []table.Column{
    37  			{Name: "str", ColType: table.TVarchar, Scale: 32, Default: "", Comment: "str column"},
    38  			{Name: "int64", ColType: table.TInt64, Default: "0", Comment: "int64 column"},
    39  			{Name: "float64", ColType: table.TFloat64, Default: "0.0", Comment: "float64 column"},
    40  			{Name: "uint64", ColType: table.TUint64, Default: "0", Comment: "uint64 column"},
    41  			{Name: "datetime_6", ColType: table.TDatetime, Default: "", Comment: "datetime.6 column"},
    42  			{Name: "json_col", ColType: table.TJson, Default: "{}", Comment: "json column"},
    43  		},
    44  	}
    45  
    46  	records := [][]string{
    47  		{"str1", "1", "1.1", "1", "2023-05-16T00:00:00Z", `{"key1":"value1 \n test , \r 'test'"}`},
    48  		{"str2", "2", "2.2", "2", "2023-05-16T00:00:00Z", `{"key2":"value2"}`},
    49  	}
    50  
    51  	db, mock, err := sqlmock.New() // creating sqlmock
    52  	if err != nil {
    53  		t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
    54  	}
    55  	mock.ExpectExec(regexp.QuoteMeta(`LOAD DATA INLINE FORMAT='csv', DATA='str1,1,1.1,1,2023-05-16T00:00:00Z,"{""key1"":""value1 \\n test , \\r ''test''""}"
    56  str2,2,2.2,2,2023-05-16T00:00:00Z,"{""key2"":""value2""}"
    57  ' INTO TABLE testDB.testTable`)).
    58  		WillReturnResult(sqlmock.NewResult(1, 1))
    59  
    60  	ctx, cancel := context.WithCancel(context.Background())
    61  	defer cancel()
    62  
    63  	bulkInsert(ctx, db, records, tbl)
    64  
    65  	err = mock.ExpectationsWereMet()
    66  	if err != nil {
    67  		t.Errorf("there were unfulfilled expectations: %s", err)
    68  	}
    69  }
    70  
    71  func TestIsRecordExisted(t *testing.T) {
    72  	// Create a new instance of sqlmock
    73  	db, mock, err := sqlmock.New()
    74  	if err != nil {
    75  		t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
    76  	}
    77  	defer db.Close()
    78  
    79  	ctx := context.TODO()
    80  	// Assuming index 12 is for 'request_at', adding a mock value for it
    81  	record := []string{"12345", "", "", "", "", "", "", "", "", "", "", "", "2021-10-10 10:00:00", "", "", "active"}
    82  	table := &table.Table{Table: "statement_info"}
    83  
    84  	// Set up your mock expectations
    85  	mock.ExpectQuery(regexp.QuoteMeta(
    86  		"SELECT EXISTS(SELECT 1 FROM `system`.statement_info WHERE statement_id = ? AND status = ? AND request_at = ?)",
    87  	)).WithArgs(record[0], record[15], record[12]).WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
    88  
    89  	// Define a function that returns the mocked DB connection
    90  	getDBConn := func(forceNewConn bool, randomCN bool) (*sql.DB, error) {
    91  		return db, nil
    92  	}
    93  
    94  	// Call your function with the mock
    95  	exists, err := IsRecordExisted(ctx, record, table, getDBConn)
    96  	if err != nil {
    97  		t.Errorf("error was not expected while checking record existence: %s", err)
    98  	}
    99  	if !exists {
   100  		t.Errorf("expected record to exist, but it does not")
   101  	}
   102  
   103  	// Ensure all expectations are met
   104  	if err := mock.ExpectationsWereMet(); err != nil {
   105  		t.Errorf("there were unfulfilled expectations: %s", err)
   106  	}
   107  }
   108  
   109  func TestSetLabelSelector(t *testing.T) {
   110  	type args struct {
   111  		labels map[string]string
   112  	}
   113  	tests := []struct {
   114  		name       string
   115  		args       args
   116  		want       map[string]string
   117  		wantGetter map[string]string
   118  	}{
   119  		{
   120  			name: "normal",
   121  			args: args{labels: map[string]string{
   122  				"role": "admin",
   123  			}},
   124  			want: map[string]string{
   125  				"role": "admin",
   126  			},
   127  			wantGetter: map[string]string{
   128  				"role":    "admin",
   129  				"account": "sys",
   130  			},
   131  		},
   132  	}
   133  
   134  	for _, tt := range tests {
   135  		t.Run(tt.name, func(t *testing.T) {
   136  			SetLabelSelector(tt.args.labels)
   137  			require.Equal(t, tt.want, tt.args.labels)
   138  			got := GetLabelSelector()
   139  			if !reflect.DeepEqual(tt.wantGetter, got) {
   140  				t.Errorf("gLabelSelector = %v, want %v", got, tt.wantGetter)
   141  			}
   142  		})
   143  	}
   144  }