github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/x/watch/source_test.go (about)

     1  // Copyright (c) 2016 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package watch
    22  
    23  import (
    24  	"errors"
    25  	"sync"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/stretchr/testify/assert"
    30  	"go.uber.org/zap"
    31  )
    32  
    33  func TestSource(t *testing.T) {
    34  	testSource(t, 30, 25, 20)
    35  	testSource(t, 22, 18, 20)
    36  	testSource(t, 15, 10, 20)
    37  	testSource(t, 28, 30, 20)
    38  	testSource(t, 19, 21, 20)
    39  	testSource(t, 13, 15, 20)
    40  	testSource(t, 18, 22, 10)
    41  	testSource(t, 14, 12, 15)
    42  }
    43  
    44  func testSource(t *testing.T, errAfter int32, closeAfter int32, watchNum int) {
    45  	input := &testSourceInput{callCount: 0, errAfter: errAfter, closeAfter: closeAfter}
    46  	zl, err := zap.NewDevelopment()
    47  	assert.NoError(t, err)
    48  	s := NewSource(input, zl)
    49  
    50  	var wg sync.WaitGroup
    51  
    52  	// Create a few watches.
    53  	for i := 0; i < watchNum; i++ {
    54  		wg.Add(1)
    55  		_, w, err := s.Watch()
    56  		assert.NoError(t, err)
    57  		assert.NotNil(t, w)
    58  
    59  		i := i
    60  		go func() {
    61  			var v interface{}
    62  			count := 0
    63  			for range w.C() {
    64  				if v != nil {
    65  					assert.True(t, w.Get().(int32) >= v.(int32))
    66  				}
    67  				v = w.Get()
    68  				if count > i {
    69  					w.Close()
    70  				}
    71  				count++
    72  			}
    73  			wg.Done()
    74  		}()
    75  	}
    76  
    77  	// Only start serving after all watchers are created.
    78  	input.startServing()
    79  
    80  	// Schedule a thread to close Source.
    81  	wg.Add(1)
    82  	go func() {
    83  		for !s.(*source).isClosed() {
    84  			time.Sleep(time.Millisecond)
    85  		}
    86  		_, _, err := s.Watch()
    87  		assert.Error(t, err)
    88  		v := s.Get()
    89  		if errAfter < closeAfter {
    90  			assert.Equal(t, v, errAfter)
    91  		} else {
    92  			assert.Equal(t, v, closeAfter)
    93  		}
    94  		// Test Close again.
    95  		s.Close()
    96  		assert.True(t, s.(*source).isClosed())
    97  		assert.Equal(t, v, s.Get())
    98  		wg.Done()
    99  	}()
   100  
   101  	wg.Wait()
   102  }
   103  
   104  type testSourceInput struct {
   105  	sync.Mutex
   106  
   107  	started    bool
   108  	callCount  int32
   109  	errAfter   int32
   110  	closeAfter int32
   111  }
   112  
   113  func (i *testSourceInput) startServing() {
   114  	i.Lock()
   115  	i.started = true
   116  	i.Unlock()
   117  }
   118  
   119  func (i *testSourceInput) Poll() (interface{}, error) {
   120  	i.Lock()
   121  	started := i.started
   122  	i.Unlock()
   123  	if !started {
   124  		return i.callCount, nil
   125  	}
   126  	if i.callCount >= i.closeAfter {
   127  		return nil, ErrSourceClosed
   128  	}
   129  	i.callCount++
   130  	time.Sleep(time.Millisecond)
   131  	if i.errAfter > 0 {
   132  		i.errAfter--
   133  		return i.callCount, nil
   134  	}
   135  	return nil, errors.New("mock error")
   136  }