<commit_before>package testutils

import (
	"sync"
	"sync/atomic"
	"time"

	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/api/policy/v1beta1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/cache"
	framework "k8s.io/client-go/tools/cache/testing"
	"k8s.io/client-go/util/workqueue"

	v1 "kubevirt.io/client-go/api/v1"
	cdiv1 "kubevirt.io/containerized-data-importer/pkg/apis/core/v1alpha1"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
)

/*
MockWorkQueue is a helper workqueue which can be wrapped around
any RateLimitingInterface implementing queue. This allows synchronous
testing of the controller. The typical pattern is:

    MockQueue.ExpectAdd(3)
    vmiSource.Add(vmi)
    vmiSource.Add(vmi1)
    vmiSource.Add(vmi2)
    MockQueue.Wait()

This ensures that Source callbacks which are listening on vmiSource
enqueued three times an object. Since enqueing is typically the last
action in listener callbacks, we can assume that the wanted scenario for
a controller is set up, and an execution will process this scenario.
*/
type MockWorkQueue struct {
	workqueue.RateLimitingInterface
	addWG            *sync.WaitGroup
	rateLimitedEnque int32
	addAfterEnque    int32
	wgLock           sync.Mutex
}

func (q *MockWorkQueue) Add(obj interface{}) {
	q.RateLimitingInterface.Add(obj)
	q.wgLock.Lock()
	defer q.wgLock.Unlock()
	if q.addWG != nil {
		q.addWG.Done()
	}
}

func (q *MockWorkQueue) AddRateLimited(item interface{}) {
	q.RateLimitingInterface.AddRateLimited(item)
	atomic.AddInt32(&q.rateLimitedEnque, 1)
}

func (q *MockWorkQueue) AddAfter(item interface{}, duration time.Duration) {
	q.RateLimitingInterface.AddAfter(item, duration)
	atomic.AddInt32(&q.addAfterEnque, 1)
}

func (q *MockWorkQueue) GetRateLimitedEnqueueCount() int {
	return int(atomic.LoadInt32(&q.rateLimitedEnque))
}

func (q *MockWorkQueue) GetAddAfterEnqueueCount() int {
	return int(atomic.LoadInt32(&q.addAfterEnque))
}

// ExpectAdds allows setting the amount of expected enqueues.
func (q *MockWorkQueue) ExpectAdds(diff int) {
	q.wgLock.Lock()
	defer q.wgLock.Unlock()
	q.addWG = &sync.WaitGroup{}
	q.addWG.Add(diff)
}

// Wait waits until the expected amount of ExpectedAdds has happened.
// It will not block if there were no expectations set.
func (q *MockWorkQueue) Wait() {
	if q.addWG != nil {
		q.addWG.Wait()
		q.addWG = nil
	}
}

func NewMockWorkQueue(queue workqueue.RateLimitingInterface) *MockWorkQueue {
	return &MockWorkQueue{queue, nil, 0, 0, sync.Mutex{}}
}

func NewFakeInformerFor(obj runtime.Object) (cache.SharedIndexInformer, *framework.FakeControllerSource) {
	objSource := framework.NewFakeControllerSource()
	objInformer := cache.NewSharedIndexInformer(objSource, obj, 0, cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc})
	return objInformer, objSource
}

func NewFakeInformerWithIndexersFor(obj runtime.Object, indexers cache.Indexers) (cache.SharedIndexInformer, *framework.FakeControllerSource) {
	objSource := framework.NewFakeControllerSource()
	objInformer := cache.NewSharedIndexInformer(objSource, obj, 0, indexers)
	return objInformer, objSource
}

type VirtualMachineFeeder struct {
	MockQueue *MockWorkQueue
	Source    *framework.FakeControllerSource
}<commit_msg>Fix race in mock_queue code<commit_after>package testutils

import (
	"sync"
	"sync/atomic"
	"time"

	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/api/policy/v1beta1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/cache"
	framework "k8s.io/client-go/tools/cache/testing"
	"k8s.io/client-go/util/workqueue"

	v1 "kubevirt.io/client-go/api/v1"
	cdiv1 "kubevirt.io/containerized-data-importer/pkg/apis/core/v1alpha1"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
)

/*
MockWorkQueue is a helper workqueue which can be wrapped around
any RateLimitingInterface implementing queue. This allows synchronous
testing of the controller. The typical pattern is:

    MockQueue.ExpectAdd(3)
    vmiSource.Add(vmi)
    vmiSource.Add(vmi1)
    vmiSource.Add(vmi2)
    MockQueue.Wait()

This ensures that Source callbacks which are listening on vmiSource
enqueued three times an object. Since enqueing is typically the last
action in listener callbacks, we can assume that the wanted scenario for
a controller is set up, and an execution will process this scenario.
*/
type MockWorkQueue struct {
	workqueue.RateLimitingInterface
	addWG            *sync.WaitGroup
	rateLimitedEnque int32
	addAfterEnque    int32
	wgLock           sync.Mutex
}

func (q *MockWorkQueue) Add(obj interface{}) {
	q.RateLimitingInterface.Add(obj)
	q.wgLock.Lock()
	defer q.wgLock.Unlock()
	if q.addWG != nil {
		q.addWG.Done()
	}
}

func (q *MockWorkQueue) AddRateLimited(item interface{}) {
	q.RateLimitingInterface.AddRateLimited(item)
	atomic.AddInt32(&q.rateLimitedEnque, 1)
}

func (q *MockWorkQueue) AddAfter(item interface{}, duration time.Duration) {
	q.RateLimitingInterface.AddAfter(item, duration)
	atomic.AddInt32(&q.addAfterEnque, 1)
}

func (q *MockWorkQueue) GetRateLimitedEnqueueCount() int {
	return int(atomic.LoadInt32(&q.rateLimitedEnque))
}

func (q *MockWorkQueue) GetAddAfterEnqueueCount() int {
	return int(atomic.LoadInt32(&q.addAfterEnque))
}

// ExpectAdds allows setting the amount of expected enqueues.
func (q *MockWorkQueue) ExpectAdds(diff int) {
	q.wgLock.Lock()
	defer q.wgLock.Unlock()
	q.addWG = &sync.WaitGroup{}
	q.addWG.Add(diff)
}

// Wait waits until the expected amount of ExpectedAdds has happened.
// It will not block if there were no expectations set.
func (q *MockWorkQueue) Wait() {
	q.wgLock.Lock()
	wg := q.addWG
	q.wgLock.Unlock()
	if wg != nil {
		wg.Wait()
		q.wgLock.Lock()
		q.addWG = nil
		q.wgLock.Unlock()
	}
}

func NewMockWorkQueue(queue workqueue.RateLimitingInterface) *MockWorkQueue {
	return &MockWorkQueue{queue, nil, 0, 0, sync.Mutex{}}
}

func NewFakeInformerFor(obj runtime.Object) (cache.SharedIndexInformer, *framework.FakeControllerSource) {
	objSource := framework.NewFakeControllerSource()
	objInformer := cache.NewSharedIndexInformer(objSource, obj, 0, cache.Indexers{cache.NamespaceIndex: cache.MetaNamespaceIndexFunc})
	return objInformer, objSource
}

func NewFakeInformerWithIndexersFor(obj runtime.Object, indexers cache.Indexers) (cache.SharedIndexInformer, *framework.FakeControllerSource) {
	objSource := framework.NewFakeControllerSource()
	objInformer := cache.NewSharedIndexInformer(objSource, obj, 0, indexers)
	return objInformer, objSource
}

type VirtualMachineFeeder struct {
	MockQueue *MockWorkQueue
	Source    *framework.FakeControllerSource
}
{'fraction_non_alphanumeric': 0.10672178839381412, 'fraction_numerical': 0.01133057724697596, 'mean_word_length': 5.025830258302583, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4010336', 'n_tokens_mistral': 2327, 'n_tokens_neox': 2269, 'n_words': 512}