github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/cmd/dm-worker/main_test.go (about)

     1  // Copyright 2019 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 main
    15  
    16  // Reference: https://dzone.com/articles/measuring-integration-test-coverage-rate-in-pouchc
    17  
    18  import (
    19  	"fmt"
    20  	"os"
    21  	"strings"
    22  	"testing"
    23  	"time"
    24  
    25  	"github.com/pingcap/tiflow/dm/pkg/log"
    26  	"github.com/pingcap/tiflow/dm/pkg/utils"
    27  	"go.uber.org/zap"
    28  )
    29  
    30  func TestRunMain(_ *testing.T) {
    31  	fmt.Println("dm-worker startup", time.Now())
    32  	var (
    33  		args   []string
    34  		exit   = make(chan int)
    35  		waitCh = make(chan interface{}, 1)
    36  	)
    37  	for _, arg := range os.Args {
    38  		switch {
    39  		case arg == "DEVEL":
    40  		case strings.HasPrefix(arg, "-test."):
    41  		default:
    42  			args = append(args, arg)
    43  		}
    44  	}
    45  
    46  	// golang cover tool rewrites ast with coverage annotations based on block.
    47  	// whenever the cover tool detects the first line of one block has runned,
    48  	// the coverage counter will be added. So we mock of `utils.OsExit` is able
    49  	// to collect the code run. While if we run `main()` in the main routine,
    50  	// when a code block from `main()` executes `utils.OsExit` and we forcedly
    51  	// exit the program by `os.Exit(0)` in other routine, the coverage counter
    52  	// fails to add for this block, the different behavior of these two scenarios
    53  	// comes from the difference between `os.Exit` and return from a function call.
    54  	oldOsExit := utils.OsExit
    55  	defer func() { utils.OsExit = oldOsExit }()
    56  	utils.OsExit = func(code int) {
    57  		log.L().Info("os exits", zap.Int("exit code", code))
    58  		exit <- code
    59  		// sleep here to prevent following code execution in the caller routine
    60  		time.Sleep(time.Second * 60)
    61  	}
    62  
    63  	os.Args = args
    64  	go func() {
    65  		main()
    66  		close(waitCh)
    67  	}()
    68  
    69  	select {
    70  	case <-waitCh:
    71  		fmt.Println("dm-worker exit", time.Now())
    72  		return
    73  	case <-exit:
    74  		fmt.Println("dm-worker exit", time.Now())
    75  		return
    76  	}
    77  }