package batch

import "sync"

// MockItemGenerator generates mock Items with unique IDs. Items are generated in a
// separate goroutine and added to a channel, which can be retrieved by calling
// GetCh.
type MockItemGenerator struct {
	closeOnce sync.Once
	done      chan struct{}
	ch        chan *Item

	mu     sync.Mutex
	nextID uint64
}

// NewMockItemGenerator returns a new MockItemGenerator.
//
// After using it, call Close to prevent a goroutine leak.
func NewMockItemGenerator() *MockItemGenerator {
	m := &MockItemGenerator{
		done: make(chan struct{}),
		ch:   make(chan *Item),
	}

	go func() {
		id := uint64(0)
		nextItem := &Item{
			id: id,
		}
		for {
			select {
			case m.ch <- nextItem:
				id++
				nextItem = &Item{
					id: id,
				}

			case <-m.done:
				return
			}
		}
	}()

	return m
}

// Close stops a MockItemGenerator's goroutine
func (m *MockItemGenerator) Close() {
	m.closeOnce.Do(func() {
		close(m.done)
	})
}

// GetCh returns a channel of Items with unique IDs.
func (m *MockItemGenerator) GetCh() <-chan *Item {
	return m.ch
}

package batch

import "sync"

// MockItemGenerator generates mock Items with unique IDs. Items are generated in a
// separate goroutine and added to a channel, which can be retrieved by calling
// GetCh.
type MockItemGenerator struct {
	closeOnce sync.Once
	done      chan struct{}
	ch        chan *Item
}

// NewMockItemGenerator returns a new MockItemGenerator.
//
// After using it, call Close to prevent a goroutine leak.
func NewMockItemGenerator() *MockItemGenerator {
	m := &MockItemGenerator{
		done: make(chan struct{}),
		ch:   make(chan *Item),
	}

	go func() {
		id := uint64(0)
		nextItem := &Item{
			id: id,
		}
		for {
			select {
			case m.ch <- nextItem:
				id++
				nextItem = &Item{
					id: id,
				}

			case <-m.done:
				return
			}
		}
	}()

	return m
}

// Close stops a MockItemGenerator's goroutine
func (m *MockItemGenerator) Close() {
	m.closeOnce.Do(func() {
		close(m.done)
	})
}

// GetCh returns a channel of Items with unique IDs.
func (m *MockItemGenerator) GetCh() <-chan *Item {
	return m.ch
}

Remove unused fields in MockItemGenerator
Remove unused fields in MockItemGenerator
