github.com/containerd/containerd/v2@v2.0.0-rc.2/pkg/ioutil/write_closer_test.go (about)

     1  /*
     2     Copyright The containerd Authors.
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package ioutil
    18  
    19  import (
    20  	"os"
    21  	"sort"
    22  	"strconv"
    23  	"strings"
    24  	"sync"
    25  	"testing"
    26  
    27  	"github.com/stretchr/testify/assert"
    28  	"github.com/stretchr/testify/require"
    29  )
    30  
    31  func TestWriteCloseInformer(t *testing.T) {
    32  	original := &writeCloser{}
    33  	wci, close := NewWriteCloseInformer(original)
    34  	data := "test"
    35  
    36  	n, err := wci.Write([]byte(data))
    37  	assert.Equal(t, len(data), n)
    38  	assert.Equal(t, data, original.buf.String())
    39  	assert.NoError(t, err)
    40  
    41  	select {
    42  	case <-close:
    43  		assert.Fail(t, "write closer closed")
    44  	default:
    45  	}
    46  
    47  	wci.Close()
    48  	assert.True(t, original.closed)
    49  
    50  	select {
    51  	case <-close:
    52  	default:
    53  		assert.Fail(t, "write closer not closed")
    54  	}
    55  }
    56  
    57  func TestSerialWriteCloser(t *testing.T) {
    58  	const (
    59  		// Test 10 times to make sure it always pass.
    60  		testCount = 10
    61  
    62  		goroutine = 10
    63  		dataLen   = 100000
    64  	)
    65  	for n := 0; n < testCount; n++ {
    66  		testData := make([][]byte, goroutine)
    67  		for i := 0; i < goroutine; i++ {
    68  			testData[i] = []byte(repeatNumber(i, dataLen) + "\n")
    69  		}
    70  
    71  		f, err := os.CreateTemp("", "serial-write-closer")
    72  		require.NoError(t, err)
    73  		defer os.RemoveAll(f.Name())
    74  		defer f.Close()
    75  		wc := NewSerialWriteCloser(f)
    76  		defer wc.Close()
    77  
    78  		// Write data in parallel
    79  		var wg sync.WaitGroup
    80  		wg.Add(goroutine)
    81  		for i := 0; i < goroutine; i++ {
    82  			go func(id int) {
    83  				n, err := wc.Write(testData[id])
    84  				assert.NoError(t, err)
    85  				assert.Equal(t, dataLen+1, n)
    86  				wg.Done()
    87  			}(i)
    88  		}
    89  		wg.Wait()
    90  		wc.Close()
    91  
    92  		// Check test result
    93  		content, err := os.ReadFile(f.Name())
    94  		require.NoError(t, err)
    95  		resultData := strings.Split(strings.TrimSpace(string(content)), "\n")
    96  		require.Len(t, resultData, goroutine)
    97  		sort.Strings(resultData)
    98  		for i := 0; i < goroutine; i++ {
    99  			expected := repeatNumber(i, dataLen)
   100  			assert.Equal(t, expected, resultData[i])
   101  		}
   102  	}
   103  }
   104  
   105  func repeatNumber(num, count int) string {
   106  	return strings.Repeat(strconv.Itoa(num), count)
   107  }