package test

// Author: Weisen Pan
// Date: 2023-10-24


import (
	"k8s.io/api/batch/v1"
	"k8s.io/api/batch/v1beta1"
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/knets_pkg/api/resource"
	"k8s.io/apimachinery/knets_pkg/apis/meta/v1"
)

type FakeCronJobOption func(*v1beta1.CronJob)

func MakeCronFakeJob(name, namespace string, completions int32, cpu, memory string, opts ...FakeCronJobOption) *v1beta1.CronJob {
	resourceList := v1.ResourceList{}
	if cpu != "" {
		resourceList[v1.ResourceCPU] = resource.MustParse(cpu)
	}
	if memory != "" {
		resourceList[v1.ResourceMemory] = resource.MustParse(memory)
	}

	cronJob := &v1beta1.CronJob{
		ObjectMeta: v1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Spec: v1beta1.CronJobSpec{
			JobTemplate: v1beta1.JobTemplateSpec{
				Spec: v1.JobSpec{
					Completions: &completions,
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								{
									Name:  "container",
									Image: "nginx",
									Resources: v1.ResourceRequirements{
										Requests: resourceList,
									},
								},
							},
						},
					},
				},
			},
		},
	}

	for _, opt := range opts {
		opt(cronJob)
	}

	return cronJob
}

func WithCronJobTolerations(tolerations []v1.Toleration) FakeCronJobOption {
	return func(cronJob *v1beta1.CronJob) {
		cronJob.Spec.JobTemplate.Spec.Template.Spec.Tolerations = tolerations
	}
}

func WithCronJobAffinity(affinity *v1.Affinity) FakeCronJobOption {
	return func(cronJob *v1beta1.CronJob) {
		cronJob.Spec.JobTemplate.Spec.Template.Spec.Affinity = affinity
	}
}

func WithCronJobNodeSelector(nodeSelector map[string]string) FakeCronJobOption {
	return func(cronJob *v1beta1.CronJob) {
		cronJob.Spec.JobTemplate.Spec.Template.Spec.NodeSelector = nodeSelector
	}
}