github.com/xmidt-org/webpa-common@v1.11.9/device/drain/mocks_test.go (about)

     1  package drain
     2  
     3  import (
     4  	"net/http"
     5  	"sync"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/mock"
     9  	"github.com/xmidt-org/webpa-common/device"
    10  )
    11  
    12  type mockDrainer struct {
    13  	mock.Mock
    14  }
    15  
    16  func (m *mockDrainer) Start(j Job) (<-chan struct{}, Job, error) {
    17  	arguments := m.Called(j)
    18  	return arguments.Get(0).(<-chan struct{}), arguments.Get(1).(Job), arguments.Error(2)
    19  }
    20  
    21  func (m *mockDrainer) Status() (bool, Job, Progress) {
    22  	arguments := m.Called()
    23  	return arguments.Bool(0), arguments.Get(1).(Job), arguments.Get(2).(Progress)
    24  }
    25  
    26  func (m *mockDrainer) Cancel() (<-chan struct{}, error) {
    27  	arguments := m.Called()
    28  	return arguments.Get(0).(<-chan struct{}), arguments.Error(1)
    29  }
    30  
    31  type stubManager struct {
    32  	lock    sync.RWMutex
    33  	assert  *assert.Assertions
    34  	devices map[device.ID]device.Interface
    35  
    36  	disconnect      chan struct{}
    37  	pauseDisconnect chan struct{}
    38  
    39  	visit      chan struct{}
    40  	pauseVisit chan struct{}
    41  }
    42  
    43  var _ device.Connector = (*stubManager)(nil)
    44  var _ device.Registry = (*stubManager)(nil)
    45  
    46  func (sm *stubManager) Connect(http.ResponseWriter, *http.Request, http.Header) (device.Interface, error) {
    47  	sm.assert.Fail("Connect is not supported")
    48  	return nil, nil
    49  }
    50  
    51  func (sm *stubManager) Disconnect(id device.ID, reason device.CloseReason) bool {
    52  	select {
    53  	case sm.disconnect <- struct{}{}:
    54  	default:
    55  	}
    56  
    57  	<-sm.pauseDisconnect
    58  	defer sm.lock.Unlock()
    59  	sm.lock.Lock()
    60  
    61  	if _, exists := sm.devices[id]; exists {
    62  		delete(sm.devices, id)
    63  		return true
    64  	}
    65  
    66  	return false
    67  }
    68  
    69  func (sm *stubManager) DisconnectIf(func(device.ID) (device.CloseReason, bool)) int {
    70  	sm.assert.Fail("DisconnectIf is not supported")
    71  	return -1
    72  }
    73  
    74  func (sm *stubManager) DisconnectAll(device.CloseReason) int {
    75  	sm.assert.Fail("DisconnectAll is not supported")
    76  	return -1
    77  }
    78  
    79  func (sm *stubManager) GetFilter() device.Filter {
    80  	sm.assert.Fail("GetFilter is not supported")
    81  	return nil
    82  }
    83  
    84  func (sm *stubManager) Len() int {
    85  	return len(sm.devices)
    86  }
    87  
    88  func (sm *stubManager) Get(device.ID) (device.Interface, bool) {
    89  	sm.assert.Fail("Get is not supported")
    90  	return nil, false
    91  }
    92  
    93  func (sm *stubManager) VisitAll(p func(device.Interface) bool) (count int) {
    94  	select {
    95  	case sm.visit <- struct{}{}:
    96  	default:
    97  	}
    98  
    99  	<-sm.pauseVisit
   100  	defer sm.lock.Unlock()
   101  	sm.lock.Lock()
   102  
   103  	for _, v := range sm.devices {
   104  		count++
   105  		if !p(v) {
   106  			break
   107  		}
   108  	}
   109  
   110  	return
   111  }
   112  
   113  func (sm *stubManager) Route(*device.Request) (*device.Response, error) {
   114  	sm.assert.Fail("Route is not supported")
   115  	return nil, nil
   116  }
   117  
   118  func generateManager(assert *assert.Assertions, count uint64) *stubManager {
   119  	sm := &stubManager{
   120  		assert:          assert,
   121  		devices:         make(map[device.ID]device.Interface, count),
   122  		disconnect:      make(chan struct{}, 10),
   123  		pauseDisconnect: make(chan struct{}),
   124  		visit:           make(chan struct{}, 10),
   125  		pauseVisit:      make(chan struct{}),
   126  	}
   127  
   128  	for mac := uint64(0); mac < count; mac++ {
   129  		var (
   130  			id = device.IntToMAC(mac)
   131  			d  = new(device.MockDevice)
   132  		)
   133  
   134  		d.On("ID").Return(id)
   135  		d.On("String").Return("mockDevice(" + string(id) + ")")
   136  		sm.devices[id] = d
   137  	}
   138  
   139  	return sm
   140  }
   141  
   142  func generateManagerWithDifferentDevices(assert *assert.Assertions, metadataOneClaims map[string]interface{}, deviceCountOne uint64, metadataTwoClaims map[string]interface{}, deviceCountTwo uint64) *stubManager {
   143  	totalCount := deviceCountOne + deviceCountTwo
   144  	sm := &stubManager{
   145  		assert:          assert,
   146  		devices:         make(map[device.ID]device.Interface, totalCount),
   147  		disconnect:      make(chan struct{}, 10),
   148  		pauseDisconnect: make(chan struct{}),
   149  		visit:           make(chan struct{}, 10),
   150  		pauseVisit:      make(chan struct{}),
   151  	}
   152  
   153  	for mac := uint64(0); mac < totalCount; mac++ {
   154  		var (
   155  			id = device.IntToMAC(mac)
   156  			d  = new(device.MockDevice)
   157  		)
   158  
   159  		d.On("ID").Return(id)
   160  		d.On("String").Return("mockDevice(" + string(id) + ")")
   161  
   162  		if mac < deviceCountOne {
   163  			metadata := new(device.Metadata)
   164  			metadata.SetClaims(metadataOneClaims)
   165  
   166  			d.On("Metadata").Return(metadata)
   167  		} else {
   168  			metadata := new(device.Metadata)
   169  			metadata.SetClaims(metadataTwoClaims)
   170  
   171  			d.On("Metadata").Return(metadata)
   172  		}
   173  
   174  		sm.devices[id] = d
   175  	}
   176  
   177  	return sm
   178  }