package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"sync"
	"testing"

	"github.com/stretchr/testify/assert"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modagent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modagentk"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/starboard_vulnerability"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/syncz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_modagent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testlogger"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/agentcfg"
	"go.uber.org/mock/gomock"
)

var (
	_ modagentk.Module  = &module{}
	_ modagentk.Factory = &Factory{}

	// Init common data
	agentConfigOCSSetting = agentcfg.ContainerScanningCF{
		Cadence: "1 * * * *",
		VulnerabilityReport: &agentcfg.VulnerabilityReport{
			Namespaces: []string{"default"},
		},
	}

	policyOCSSetting = agentcfg.ContainerScanningCF{
		Cadence: "2 * * * *",
		VulnerabilityReport: &agentcfg.VulnerabilityReport{
			Namespaces: []string{"kube-system"},
		},
	}

	defaultResourceRequirements = agentcfg.ResourceRequirements{
		Limits: &agentcfg.Resource{
			Cpu:              defaultTrivyResourceLimitsCPU,
			Memory:           defaultTrivyResourceLimitsMemory,
			EphemeralStorage: defaultTrivyResourceLimitsEphemeralStorage,
		},
		Requests: &agentcfg.Resource{
			Cpu:              defaultTrivyResourceRequestsCPU,
			Memory:           defaultTrivyResourceRequestsMemory,
			EphemeralStorage: defaultTrivyResourceRequestsEphemeralStorage,
		},
	}

	customResourceRequirements = agentcfg.ResourceRequirements{
		Limits: &agentcfg.Resource{
			Cpu:              "100m",
			Memory:           "100Mi",
			EphemeralStorage: "3Gi",
		},
		Requests: &agentcfg.Resource{
			Cpu:              "100m",
			Memory:           "100Mi",
			EphemeralStorage: "1Gi",
		},
	}

	customTrivyImage = agentcfg.TrivyK8SWrapperImage{
		Repository: "registry.gitlab.com/security-products/trivy-k8s-wrapper-custom",
		Tag:        "0.9.0-debug",
	}
)

func TestEmptyConfig(t *testing.T) {
	m := new(module)
	cfg := &agentcfg.AgentConfiguration{}

	assert.NoError(t, m.DefaultAndValidateConfiguration(cfg))
	finalResourceRequirements := cfg.ContainerScanning.ResourceRequirements

	// m.DefaultAndValidateConfiguration would default resource requirements even if config is empty since security policies could still enforce vulnerability scan
	assert.Equal(t, defaultTrivyResourceLimitsCPU, finalResourceRequirements.Limits.Cpu)
	assert.Equal(t, defaultTrivyResourceLimitsMemory, finalResourceRequirements.Limits.Memory)
	assert.Equal(t, defaultTrivyResourceLimitsEphemeralStorage, finalResourceRequirements.Limits.EphemeralStorage)
	assert.Equal(t, defaultTrivyResourceRequestsCPU, finalResourceRequirements.Requests.Cpu)
	assert.Equal(t, defaultTrivyResourceRequestsMemory, finalResourceRequirements.Requests.Memory)
	assert.Equal(t, defaultTrivyResourceRequestsEphemeralStorage, finalResourceRequirements.Requests.EphemeralStorage)

	assert.Equal(t, defaultTrivyK8sWrapperImageRepository, cfg.ContainerScanning.TrivyK8SWrapperImage.Repository)
	assert.Equal(t, defaultTrivyK8sWrapperImageTag, cfg.ContainerScanning.TrivyK8SWrapperImage.Tag)

	assert.Equal(t, []string{
		kindPod,
		kindReplicaSet,
		kindReplicationController,
		kindStatefulSet,
		kindDaemonSet,
		kindCronJob,
		kindJob,
		kindDeployment,
	}, cfg.ContainerScanning.VulnerabilityReport.ResourceTypes)

	assert.True(t, *cfg.ContainerScanning.DeleteReportArtifact)
}

func TestConfigWithResourceRequirements(t *testing.T) {
	m := new(module)
	cfg := &agentcfg.AgentConfiguration{
		ContainerScanning: &agentcfg.ContainerScanningCF{
			ResourceRequirements: &customResourceRequirements,
		},
	}

	assert.NoError(t, m.DefaultAndValidateConfiguration(cfg))
	finalResourceRequirements := cfg.ContainerScanning.ResourceRequirements

	assert.Equal(t, customResourceRequirements.Limits.Cpu, finalResourceRequirements.Limits.Cpu)
	assert.Equal(t, customResourceRequirements.Limits.Memory, finalResourceRequirements.Limits.Memory)
	assert.Equal(t, customResourceRequirements.Limits.EphemeralStorage, finalResourceRequirements.Limits.EphemeralStorage)
	assert.Equal(t, customResourceRequirements.Requests.Cpu, finalResourceRequirements.Requests.Cpu)
	assert.Equal(t, customResourceRequirements.Requests.Memory, finalResourceRequirements.Requests.Memory)
	assert.Equal(t, customResourceRequirements.Requests.EphemeralStorage, finalResourceRequirements.Requests.EphemeralStorage)
}

func TestConfigWithPartialResourceRequirements(t *testing.T) {
	testCases := []struct {
		description                      string
		inputLimitsCPU                   string
		inputLimitsMemory                string
		inputLimitsEphemeralStorage      string
		inputRequestsCPU                 string
		inputRequestsMemory              string
		inputRequestsEphemeralStorage    string
		expectedLimitsCPU                string
		expectedLimitsMemory             string
		expectedLimitsEphemeralStorage   string
		expectedRequestsCPU              string
		expectedRequestsMemory           string
		expectedRequestsEphemeralStorage string
	}{
		{
			description:                      "only has limits cpu",
			inputLimitsCPU:                   customResourceRequirements.Limits.Cpu,
			inputLimitsMemory:                "",
			inputLimitsEphemeralStorage:      "",
			inputRequestsCPU:                 "",
			inputRequestsMemory:              "",
			inputRequestsEphemeralStorage:    "",
			expectedLimitsCPU:                customResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:             defaultResourceRequirements.Limits.Memory,
			expectedLimitsEphemeralStorage:   defaultResourceRequirements.Limits.EphemeralStorage,
			expectedRequestsCPU:              defaultResourceRequirements.Requests.Cpu,
			expectedRequestsMemory:           defaultResourceRequirements.Requests.Memory,
			expectedRequestsEphemeralStorage: defaultResourceRequirements.Requests.EphemeralStorage,
		},
		{
			description:                      "only has limits memory",
			inputLimitsCPU:                   "",
			inputLimitsMemory:                customResourceRequirements.Limits.Memory,
			inputLimitsEphemeralStorage:      "",
			inputRequestsCPU:                 "",
			inputRequestsMemory:              "",
			inputRequestsEphemeralStorage:    "",
			expectedLimitsCPU:                defaultResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:             customResourceRequirements.Limits.Memory,
			expectedLimitsEphemeralStorage:   defaultResourceRequirements.Limits.EphemeralStorage,
			expectedRequestsCPU:              defaultResourceRequirements.Requests.Cpu,
			expectedRequestsMemory:           defaultResourceRequirements.Requests.Memory,
			expectedRequestsEphemeralStorage: defaultResourceRequirements.Requests.EphemeralStorage,
		},
		{
			description:                      "only has limits ephemeral storage",
			inputLimitsCPU:                   "",
			inputLimitsMemory:                "",
			inputLimitsEphemeralStorage:      customResourceRequirements.Limits.EphemeralStorage,
			inputRequestsCPU:                 "",
			inputRequestsMemory:              "",
			inputRequestsEphemeralStorage:    "",
			expectedLimitsCPU:                defaultResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:             defaultResourceRequirements.Limits.Memory,
			expectedLimitsEphemeralStorage:   customResourceRequirements.Limits.EphemeralStorage,
			expectedRequestsCPU:              defaultResourceRequirements.Requests.Cpu,
			expectedRequestsMemory:           defaultResourceRequirements.Requests.Memory,
			expectedRequestsEphemeralStorage: defaultResourceRequirements.Requests.EphemeralStorage,
		},
		{
			description:                      "only has requests cpu",
			inputLimitsCPU:                   "",
			inputLimitsMemory:                "",
			inputLimitsEphemeralStorage:      "",
			inputRequestsCPU:                 customResourceRequirements.Requests.Cpu,
			inputRequestsMemory:              "",
			inputRequestsEphemeralStorage:    "",
			expectedLimitsCPU:                defaultResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:             defaultResourceRequirements.Limits.Memory,
			expectedLimitsEphemeralStorage:   defaultResourceRequirements.Limits.EphemeralStorage,
			expectedRequestsCPU:              customResourceRequirements.Requests.Cpu,
			expectedRequestsMemory:           defaultResourceRequirements.Requests.Memory,
			expectedRequestsEphemeralStorage: defaultResourceRequirements.Requests.EphemeralStorage,
		},
		{
			description:                      "only has requests memory",
			inputLimitsCPU:                   "",
			inputLimitsMemory:                "",
			inputLimitsEphemeralStorage:      "",
			inputRequestsCPU:                 "",
			inputRequestsMemory:              customResourceRequirements.Requests.Memory,
			inputRequestsEphemeralStorage:    "",
			expectedLimitsCPU:                defaultResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:             defaultResourceRequirements.Limits.Memory,
			expectedLimitsEphemeralStorage:   defaultResourceRequirements.Limits.EphemeralStorage,
			expectedRequestsCPU:              defaultResourceRequirements.Requests.Cpu,
			expectedRequestsMemory:           customResourceRequirements.Requests.Memory,
			expectedRequestsEphemeralStorage: defaultResourceRequirements.Requests.EphemeralStorage,
		},
		{
			description:                      "only has requests ephemeral storage",
			inputLimitsCPU:                   "",
			inputLimitsMemory:                "",
			inputLimitsEphemeralStorage:      "",
			inputRequestsCPU:                 "",
			inputRequestsMemory:              "",
			inputRequestsEphemeralStorage:    customResourceRequirements.Requests.EphemeralStorage,
			expectedLimitsCPU:                defaultResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:             defaultResourceRequirements.Limits.Memory,
			expectedLimitsEphemeralStorage:   defaultResourceRequirements.Limits.EphemeralStorage,
			expectedRequestsCPU:              defaultResourceRequirements.Requests.Cpu,
			expectedRequestsMemory:           defaultResourceRequirements.Requests.Memory,
			expectedRequestsEphemeralStorage: customResourceRequirements.Requests.EphemeralStorage,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			m := new(module)
			cfg := &agentcfg.AgentConfiguration{
				ContainerScanning: &agentcfg.ContainerScanningCF{
					ResourceRequirements: &agentcfg.ResourceRequirements{
						Limits: &agentcfg.Resource{
							Cpu:              tc.inputLimitsCPU,
							Memory:           tc.inputLimitsMemory,
							EphemeralStorage: tc.inputLimitsEphemeralStorage,
						},
						Requests: &agentcfg.Resource{
							Cpu:              tc.inputRequestsCPU,
							Memory:           tc.inputRequestsMemory,
							EphemeralStorage: tc.inputRequestsEphemeralStorage,
						},
					},
				},
			}

			assert.NoError(t, m.DefaultAndValidateConfiguration(cfg))
			newResourceRequirements := cfg.ContainerScanning.ResourceRequirements

			assert.Equal(t, tc.expectedLimitsCPU, newResourceRequirements.Limits.Cpu)
			assert.Equal(t, tc.expectedLimitsMemory, newResourceRequirements.Limits.Memory)
			assert.Equal(t, tc.expectedLimitsEphemeralStorage, newResourceRequirements.Limits.EphemeralStorage)
			assert.Equal(t, tc.expectedRequestsCPU, newResourceRequirements.Requests.Cpu)
			assert.Equal(t, tc.expectedRequestsMemory, newResourceRequirements.Requests.Memory)
			assert.Equal(t, tc.expectedRequestsEphemeralStorage, newResourceRequirements.Requests.EphemeralStorage)
		})
	}
}

func TestConfigWithCustomTrivyImage(t *testing.T) {
	m := new(module)
	cfg := &agentcfg.AgentConfiguration{
		ContainerScanning: &agentcfg.ContainerScanningCF{
			TrivyK8SWrapperImage: &customTrivyImage,
		},
	}

	assert.NoError(t, m.DefaultAndValidateConfiguration(cfg))

	assert.Equal(t, customTrivyImage.Repository, cfg.ContainerScanning.TrivyK8SWrapperImage.Repository)
	assert.Equal(t, customTrivyImage.Tag, cfg.ContainerScanning.TrivyK8SWrapperImage.Tag)
}

func TestConfigWithCustomDeleteConfig(t *testing.T) {
	m := new(module)
	disableDeleteReportArtifact := false
	cfg := &agentcfg.AgentConfiguration{
		ContainerScanning: &agentcfg.ContainerScanningCF{
			DeleteReportArtifact: &disableDeleteReportArtifact,
		},
	}

	assert.NoError(t, m.DefaultAndValidateConfiguration(cfg))

	assert.False(t, *cfg.ContainerScanning.DeleteReportArtifact)
}

func TestName(t *testing.T) {
	m := new(module)

	assert.Equal(t, starboard_vulnerability.ModuleName, m.Name())
}

func TestCadenceValidation(t *testing.T) {
	// Possible values documented at https://pkg.go.dev/github.com/robfig/cron#hdr-CRON_Expression_Format
	testCases := []struct {
		description string
		cadence     string
		shouldError bool
	}{
		{
			description: "daily schedule",
			cadence:     "0 0 * * *",
			shouldError: false,
		},
		{
			description: "using all fields",
			cadence:     "0 0 1 1 1",
			shouldError: false,
		},
		{
			description: "range of seconds",
			cadence:     "0-30 0 * * *",
			shouldError: false,
		},
		{
			description: "range of minutes",
			cadence:     "0 0-23 * * *",
			shouldError: false,
		},
		{
			description: "range of hours",
			cadence:     "0 0 1-12 * *",
			shouldError: false,
		},
		{
			description: "range of days",
			cadence:     "0 0 * 1-12 *",
			shouldError: false,
		},
		{
			description: "range of months (numeric)",
			cadence:     "0 0 * * 1-6",
			shouldError: false,
		},
		{
			description: "using question mark",
			cadence:     "0 0 * ? *",
			shouldError: false,
		},
		{
			description: "seconds out of range",
			cadence:     "60 0 * * *",
			shouldError: true,
		},
		{
			description: "minutes out of range",
			cadence:     "60 0 * * *",
			shouldError: true,
		},
		{
			description: "hours out of range",
			cadence:     "0 24 1 * *",
			shouldError: true,
		},
		{
			description: "days out of range",
			cadence:     "0 0 * 32 *",
			shouldError: true,
		},
		{
			description: "months out of range",
			cadence:     "0 0 * * 13",
			shouldError: true,
		},
		{
			description: "zero day",
			cadence:     "0 0 * 0 *",
			shouldError: true,
		},
		{
			description: "zero month",
			cadence:     "0 0 * 0 0",
			shouldError: true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			m := new(module)
			cfg := &agentcfg.AgentConfiguration{
				ContainerScanning: &agentcfg.ContainerScanningCF{
					Cadence: tc.cadence,
				},
			}

			err := m.DefaultAndValidateConfiguration(cfg)
			if tc.shouldError {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestIsConfigEqual(t *testing.T) {
	oldCfg := configurationToUpdateData{containerScanningConfig: &agentcfg.ContainerScanningCF{}}
	newCfg := configurationToUpdateData{containerScanningConfig: &agentcfg.ContainerScanningCF{}}

	assert.True(t, isConfigEqual(oldCfg, newCfg))

	newCfg.containerScanningConfig = &agentcfg.ContainerScanningCF{
		Cadence: "0 0 * * *",
	}

	assert.False(t, isConfigEqual(oldCfg, newCfg))
}

const (
	expectedAgentConfig = "AGENT CONFIG"
	expectedPolicy      = "expectedPolicy"
)

type customWorkerFactory struct {
	cfg      configurationToUpdateData
	cfgMutex sync.Mutex
	wg       sync.WaitGroup
}

func (f *customWorkerFactory) New(cfg configurationToUpdateData) syncz.Worker {
	// decrement wg
	defer f.wg.Done()

	// Save cfg
	f.cfgMutex.Lock()
	defer f.cfgMutex.Unlock()
	f.cfg = cfg

	return syncz.WorkerFunc(func(ctx context.Context) {}) // do nothing
}

func TestModuleConfigScenarios(t *testing.T) {
	testCases := []struct {
		description                        string
		policyHasOCSSetting                bool
		agentConfigHasOCSSetting           bool
		agentConfigHasResourceRequirements bool
		expectCustomResourceSet            bool
		expectOCSSettingFrom               string
	}{
		{
			description:                        "OCS setting from policy, resource not set",
			policyHasOCSSetting:                true,
			agentConfigHasOCSSetting:           false,
			agentConfigHasResourceRequirements: false,
			expectCustomResourceSet:            false,
			expectOCSSettingFrom:               expectedPolicy,
		},
		{
			description:                        "OCS setting from policy and agentConfig, resource not set",
			policyHasOCSSetting:                true,
			agentConfigHasOCSSetting:           true,
			agentConfigHasResourceRequirements: false,
			expectCustomResourceSet:            false,
			expectOCSSettingFrom:               expectedPolicy,
		},
		{
			description:                        "OCS setting from policy and agentConfig, resource set",
			policyHasOCSSetting:                true,
			agentConfigHasOCSSetting:           true,
			agentConfigHasResourceRequirements: true,
			expectCustomResourceSet:            true,
			expectOCSSettingFrom:               expectedPolicy,
		},
		{
			description:                        "OCS setting from policy, resource set",
			policyHasOCSSetting:                true,
			agentConfigHasOCSSetting:           false,
			agentConfigHasResourceRequirements: true,
			expectCustomResourceSet:            true,
			expectOCSSettingFrom:               expectedPolicy,
		},
		{
			description:                        "OCS setting from agentConfig, resource not set",
			policyHasOCSSetting:                false,
			agentConfigHasOCSSetting:           true,
			agentConfigHasResourceRequirements: false,
			expectCustomResourceSet:            false,
			expectOCSSettingFrom:               expectedAgentConfig,
		},
		{
			description:                        "OCS setting from agentConfig, resource set",
			policyHasOCSSetting:                false,
			agentConfigHasOCSSetting:           true,
			agentConfigHasResourceRequirements: true,
			expectCustomResourceSet:            true,
			expectOCSSettingFrom:               expectedAgentConfig,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			// Setup module
			policiesUpdateDataChannel := make(chan configurationToUpdateData)
			ctrl := gomock.NewController(t)
			api := mock_modagent.NewMockAPI(ctrl)
			cnwh := customWorkerFactory{}
			m := module{
				log:                       testlogger.New(t),
				api:                       api,
				policiesUpdateDataChannel: policiesUpdateDataChannel,
				workerFactory:             cnwh.New,
			}

			// Stub policies_configuration request
			api.EXPECT().
				MakeGitLabRequest(gomock.Any(), "/policies_configuration", gomock.Any()).
				MinTimes(0).
				DoAndReturn(func(ctx context.Context, path string, opts ...modagent.GitLabRequestOption) (*modagent.GitLabResponse, error) {
					body, err := json.Marshal(map[string]any{})
					if err != nil {
						return nil, err
					}
					return &modagent.GitLabResponse{
						StatusCode: http.StatusOK,
						Body:       io.NopCloser(bytes.NewReader(body)),
					}, nil
				})
			// Stub GetAgentKey request
			api.EXPECT().
				GetAgentKey(gomock.Any()).
				MinTimes(0).
				Return(testhelpers.AgentkKey1, nil)

			ctx := context.Background()
			var wg sync.WaitGroup
			wg.Add(1)
			defer wg.Wait()

			repositoryConfigChan := make(chan *agentcfg.AgentConfiguration)
			defer close(repositoryConfigChan)

			go func() {
				defer wg.Done()
				err := m.Run(ctx, repositoryConfigChan)
				assert.NoError(t, err)
			}()

			// Configure agentConfig to be sent to module
			agentConfig := agentcfg.AgentConfiguration{
				ContainerScanning: &agentcfg.ContainerScanningCF{},
			}
			if tc.agentConfigHasOCSSetting {
				agentConfig.ContainerScanning.Cadence = agentConfigOCSSetting.Cadence
				agentConfig.ContainerScanning.VulnerabilityReport = agentConfigOCSSetting.VulnerabilityReport
				// Increment wg since change in the agentConfig OCS setting would cause module.Run to update the new
				// worker holder and we want to wait for that process to complete before continuing with the test.
				cnwh.wg.Add(1)
			}
			if tc.agentConfigHasResourceRequirements {
				agentConfig.ContainerScanning.ResourceRequirements = &customResourceRequirements
			}

			// Send agentConfig to module
			repositoryConfigChan <- &agentConfig

			// Configure policy to be sent to module
			policy := configurationToUpdateData{
				agentKey: testhelpers.AgentkKey1,
			}
			if tc.policyHasOCSSetting {
				policy.containerScanningConfig = &policyOCSSetting
				// Increment wg since change in the policy would cause module.Run to update the new worker holder
				// and we want to wait for that process to complete before continuing with the test.
				cnwh.wg.Add(1)
			}

			// Send policy to module
			policiesUpdateDataChannel <- policy

			// Wait for module to complete processing of new CS config before asserting
			cnwh.wg.Wait()
			cnwh.cfgMutex.Lock()
			defer cnwh.cfgMutex.Unlock()
			finalContainerScanningConfig := cnwh.cfg.containerScanningConfig

			if tc.expectOCSSettingFrom == expectedAgentConfig {
				assert.Equal(t, agentConfigOCSSetting.Cadence, finalContainerScanningConfig.Cadence)
				assert.Equal(t, agentConfigOCSSetting.VulnerabilityReport, finalContainerScanningConfig.VulnerabilityReport)
			}
			if tc.expectOCSSettingFrom == expectedPolicy {
				assert.Equal(t, policyOCSSetting.Cadence, finalContainerScanningConfig.Cadence)
				assert.Equal(t, policyOCSSetting.VulnerabilityReport, finalContainerScanningConfig.VulnerabilityReport)
			}
			if tc.expectCustomResourceSet {
				assert.Equal(t, &customResourceRequirements, finalContainerScanningConfig.ResourceRequirements)
			}
		})
	}
}

func TestModuleConfigChangeScenarios(t *testing.T) {
	testCases := []struct {
		description          string
		lastOCSSettingFrom   string
		expectOCSSettingFrom string
	}{
		{
			description:          "policy configured after agentConfig",
			lastOCSSettingFrom:   expectedPolicy,
			expectOCSSettingFrom: expectedPolicy,
		},
		{
			description:          "agentConfig configured after policy",
			lastOCSSettingFrom:   expectedPolicy,
			expectOCSSettingFrom: expectedPolicy,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			// Setup module
			policiesUpdateDataChannel := make(chan configurationToUpdateData)
			ctrl := gomock.NewController(t)
			api := mock_modagent.NewMockAPI(ctrl)
			cnwh := customWorkerFactory{}
			m := module{
				log:                       testlogger.New(t),
				api:                       api,
				policiesUpdateDataChannel: policiesUpdateDataChannel,
				workerFactory:             cnwh.New,
			}

			// Stub policies_configuration request
			api.EXPECT().
				MakeGitLabRequest(gomock.Any(), "/policies_configuration", gomock.Any()).
				MinTimes(0).
				DoAndReturn(func(ctx context.Context, path string, opts ...modagent.GitLabRequestOption) (*modagent.GitLabResponse, error) {
					body, err := json.Marshal(map[string]any{})
					if err != nil {
						return nil, err
					}
					return &modagent.GitLabResponse{
						StatusCode: http.StatusOK,
						Body:       io.NopCloser(bytes.NewReader(body)),
					}, nil
				})
			// Stub GetAgentKey request
			api.EXPECT().
				GetAgentKey(gomock.Any()).
				MinTimes(0).
				Return(testhelpers.AgentkKey1, nil)

			ctx := context.Background()
			var wg sync.WaitGroup
			wg.Add(1)
			defer wg.Wait()

			repositoryConfigChan := make(chan *agentcfg.AgentConfiguration)
			defer close(repositoryConfigChan)

			go func() {
				defer wg.Done()
				err := m.Run(ctx, repositoryConfigChan)
				assert.NoError(t, err)
			}()

			// Configure agentConfig to be sent to module
			agentConfig := agentcfg.AgentConfiguration{
				ContainerScanning: &agentcfg.ContainerScanningCF{
					Cadence:              agentConfigOCSSetting.Cadence,
					VulnerabilityReport:  agentConfigOCSSetting.VulnerabilityReport,
					ResourceRequirements: &customResourceRequirements,
				},
			}

			// Configure policy to be sent to module
			policy := configurationToUpdateData{
				agentKey:                testhelpers.AgentkKey1,
				containerScanningConfig: &policyOCSSetting,
			}

			// Send agentConfig and policy to module. Increment wg since these config changes would cause module.Run
			// to update the new worker holder and we want to wait for that process to complete before continuing with the test.
			if tc.lastOCSSettingFrom == expectedPolicy {
				cnwh.wg.Add(1)
				repositoryConfigChan <- &agentConfig
				cnwh.wg.Add(1)
				policiesUpdateDataChannel <- policy
			}
			if tc.lastOCSSettingFrom == expectedAgentConfig {
				cnwh.wg.Add(1)
				policiesUpdateDataChannel <- policy
				cnwh.wg.Add(1)
				repositoryConfigChan <- &agentConfig
			}

			// Wait for module to complete processing of new CS config before asserting
			cnwh.wg.Wait()
			cnwh.cfgMutex.Lock()
			defer cnwh.cfgMutex.Unlock()
			finalContainerScanningConfig := cnwh.cfg.containerScanningConfig

			if tc.expectOCSSettingFrom == expectedAgentConfig {
				assert.Equal(t, agentConfigOCSSetting.Cadence, finalContainerScanningConfig.Cadence)
				assert.Equal(t, agentConfigOCSSetting.VulnerabilityReport, finalContainerScanningConfig.VulnerabilityReport)
			}
			if tc.expectOCSSettingFrom == expectedPolicy {
				assert.Equal(t, policyOCSSetting.Cadence, finalContainerScanningConfig.Cadence)
				assert.Equal(t, policyOCSSetting.VulnerabilityReport, finalContainerScanningConfig.VulnerabilityReport)
			}

			assert.Equal(t, &customResourceRequirements, finalContainerScanningConfig.ResourceRequirements)
		})
	}
}
