github.com/jfrog/jfrog-cli-core/v2@v2.51.0/artifactory/commands/transferfiles/phase_test.go (about)

     1  package transferfiles
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/jfrog/gofrog/parallel"
     9  	"github.com/jfrog/jfrog-client-go/utils"
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  const zeroUint32 uint32 = 0
    14  
    15  func TestStopGracefully(t *testing.T) {
    16  	pcWrapper := newProducerConsumerWrapper()
    17  	pBase := &phaseBase{pcDetails: &pcWrapper}
    18  	chunkUploaderProducerConsumer := pBase.pcDetails.chunkUploaderProducerConsumer
    19  	chunkBuilderProducerConsumer := pBase.pcDetails.chunkBuilderProducerConsumer
    20  	go func() {
    21  		// Stop gracefully after half second
    22  		time.Sleep(time.Second / 2)
    23  
    24  		// Assert active threads before stopping the producer consumers
    25  		chunkUploaderProducerConsumer.IsStarted()
    26  		assert.Greater(t, chunkUploaderProducerConsumer.ActiveThreads(), zeroUint32)
    27  		assert.Greater(t, chunkBuilderProducerConsumer.ActiveThreads(), zeroUint32)
    28  
    29  		// Stop the running threads
    30  		pBase.StopGracefully()
    31  	}()
    32  
    33  	var err error
    34  	// Run 5 counter tasks in the uploader and builder producer-consumers
    35  	uploaderCounter, builderCounter := 0, 0
    36  	for i := 0; i < 5; i++ {
    37  		_, err = chunkUploaderProducerConsumer.AddTask(createCounterTask(&uploaderCounter))
    38  		assert.NoError(t, err)
    39  		_, err = chunkBuilderProducerConsumer.AddTask(createCounterTask(&builderCounter))
    40  		assert.NoError(t, err)
    41  	}
    42  	err = runProducerConsumers(pBase.pcDetails)
    43  	assert.NoError(t, err)
    44  
    45  	// Wait for no active threads
    46  	waitForTasksToFinish(t, chunkUploaderProducerConsumer, chunkBuilderProducerConsumer)
    47  
    48  	// Since we stopped the tasks after half second, and the tasks sleep for one second during their execution, expect the tasks to run exactly once.
    49  	assert.Equal(t, 1, uploaderCounter)
    50  	assert.Equal(t, 1, builderCounter)
    51  }
    52  
    53  // Create a task that increases the counter by 1 after a second
    54  func createCounterTask(counter *int) parallel.TaskFunc {
    55  	return func(int) error {
    56  		*counter++
    57  		time.Sleep(time.Second)
    58  		return nil
    59  	}
    60  }
    61  
    62  func waitForTasksToFinish(t *testing.T, chunkUploaderProducerConsumer, chunkBuilderProducerConsumer parallel.Runner) {
    63  	// Wait for no active threads
    64  	pollingExecutor := &utils.RetryExecutor{
    65  		MaxRetries:               10,
    66  		RetriesIntervalMilliSecs: 1000,
    67  		ErrorMessage:             "Active producer-consumer tasks remained",
    68  		ExecutionHandler: func() (shouldRetry bool, err error) {
    69  			if chunkUploaderProducerConsumer.ActiveThreads() == 0 && chunkBuilderProducerConsumer.ActiveThreads() == 0 {
    70  				return false, nil
    71  			}
    72  			return true, fmt.Errorf("active uploader threads: %d. Active builder threads: %d", chunkUploaderProducerConsumer.ActiveThreads(), chunkBuilderProducerConsumer.ActiveThreads())
    73  		},
    74  	}
    75  	assert.NoError(t, pollingExecutor.Execute())
    76  }