modules/Devolutions.CIEM.Checks/Tests/Unit/InvokeCIEMScanParallel.Tests.ps1

BeforeAll {
    Remove-Module Devolutions.CIEM -Force -ErrorAction SilentlyContinue
    Import-Module (Join-Path $PSScriptRoot '..' '..' '..' '..' 'Devolutions.CIEM.psd1')
    $script:ScanSource = Get-Content (Join-Path $PSScriptRoot '..' '..' 'Private' 'Invoke-CIEMScan.ps1') -Raw
    $script:LegacyScanFixture = Get-Content (Join-Path $PSScriptRoot '..' 'Fixtures' 'legacy-scan-output.json') -Raw | ConvertFrom-Json
}

Describe 'Invoke-CIEMScan parallel execution' {
    BeforeEach {
        $script:TestDatabasePath = Join-Path $TestDrive ("ciem-" + [guid]::NewGuid().ToString('N') + '.db')
        $env:CIEM_TEST_DB_PATH = $script:TestDatabasePath
        New-CIEMDatabase -Path $script:TestDatabasePath

        InModuleScope Devolutions.CIEM {
            $script:DatabasePath = $env:CIEM_TEST_DB_PATH
            $script:AuthContext = @{
                Azure = [pscustomobject]@{
                    AccountId = 'test-account'
                    AccountType = 'ServicePrincipal'
                    SubscriptionIds = @('sub1')
                }
            }
            $script:AzureAuthContext = [pscustomobject]@{
                IsConnected = $true
                TenantId = 'tenant1'
                SubscriptionIds = @('sub1')
                ARMToken = 'arm-token'
                GraphToken = 'graph-token'
                KeyVaultToken = 'kv-token'
            }
        }

        Mock -ModuleName Devolutions.CIEM Write-CIEMLog {}
        Mock -ModuleName Devolutions.CIEM Sync-CIEMCheckCatalog {}
        Mock -ModuleName Devolutions.CIEM Get-CIEMAzureDiscoveryRun {
            [pscustomobject]@{ Id = 1; Status = 'Completed' }
        }
        Mock -ModuleName Devolutions.CIEM Get-CIEMAzureEntraResource { @() }
        Mock -ModuleName Devolutions.CIEM Get-CIEMAzureArmResource { @() }
        Mock -ModuleName Devolutions.CIEM Get-CIEMAzureResourceRelationship { @() }
    }

    It 'Dispatches selected checks through Invoke-CIEMParallelForEach with the scan throttle limit' {
        $script:ParallelInvocation = $null
        Mock -ModuleName Devolutions.CIEM Invoke-CIEMParallelForEach {
            param($InputObject, $ThrottleLimit, $ScriptBlock)

            $script:ParallelInvocation = [pscustomobject]@{
                Count = @($InputObject).Count
                ThrottleLimit = $ThrottleLimit
            }

            @()
        }

        Save-CIEMCheck -Id 'parallel_a' `
            -Provider 'Azure' `
            -Service 'Entra' `
            -Title 'Parallel A' `
            -Severity 'medium' `
            -CheckScript 'Test-EntraSecurityDefaultsEnabled.ps1' `
            -DataNeeds @('entra:securitydefaults')

        Save-CIEMCheck -Id 'parallel_b' `
            -Provider 'Azure' `
            -Service 'Entra' `
            -Title 'Parallel B' `
            -Severity 'medium' `
            -CheckScript 'Test-EntraTrustedNamedLocationExist.ps1' `
            -DataNeeds @('entra:namedlocations')

        InModuleScope Devolutions.CIEM {
            Invoke-CIEMScan -Provider Azure -CheckId @(
                'parallel_a',
                'parallel_b'
            ) | Out-Null
        }

        $expectedThrottleLimit = InModuleScope Devolutions.CIEM { $script:CIEMParallelThrottleLimitScan }

        $script:ParallelInvocation | Should -Not -BeNullOrEmpty
        $script:ParallelInvocation.Count | Should -Be 2
        $script:ParallelInvocation.ThrottleLimit | Should -Be $expectedThrottleLimit
    }

    It 'Surfaces a child failure with the failing check id' {
        Mock -ModuleName Devolutions.CIEM Invoke-CIEMParallelForEach {
            @(
                [pscustomobject]@{
                    Input = [pscustomobject]@{
                        Check = [pscustomobject]@{
                            Id = 'entra_security_defaults_enabled'
                        }
                    }
                    Success = $false
                    Error = 'boom'
                }
            )
        }

        Save-CIEMCheck -Id 'parallel_failure' `
            -Provider 'Azure' `
            -Service 'Entra' `
            -Title 'Parallel Failure' `
            -Severity 'medium' `
            -CheckScript 'Test-EntraSecurityDefaultsEnabled.ps1' `
            -DataNeeds @('entra:securitydefaults')

        InModuleScope Devolutions.CIEM {
            { Invoke-CIEMScan -Provider Azure -CheckId 'parallel_failure' | Out-Null } | Should -Throw "*Check 'entra_security_defaults_enabled' failed: boom*"
        }
    }

    It 'Reconstructs CIEMScanResult objects from child output' {
        Mock -ModuleName Devolutions.CIEM Invoke-CIEMParallelForEach {
            @(
                [pscustomobject]@{
                    Success = $true
                    Result = @(
                        [pscustomobject]@{
                            Check = [pscustomobject]@{
                                Id = 'entra_security_defaults_enabled'
                                Provider = 'Azure'
                                Service = 'Entra'
                                Title = 'Security defaults enabled'
                                Description = 'desc'
                                Risk = 'risk'
                                Severity = 'high'
                                RelatedUrl = ''
                                CheckScript = 'Test-EntraSecurityDefaultsEnabled.ps1'
                                DependsOn = @()
                                DataNeeds = @('entra:securitydefaults')
                                Disabled = $false
                                Remediation = [pscustomobject]@{
                                    Text = 'text'
                                    Url = 'url'
                                }
                                Permissions = [pscustomobject]@{
                                    Graph = @()
                                    ARM = @()
                                    KeyVaultDataPlane = @()
                                    IAM = @()
                                }
                            }
                            Status = 'PASS'
                            StatusExtended = 'ok'
                            ResourceId = 'resource-1'
                            ResourceName = 'resource-name'
                            Location = 'Global'
                        }
                    )
                }
            )
        }

        Save-CIEMCheck -Id 'parallel_result' `
            -Provider 'Azure' `
            -Service 'Entra' `
            -Title 'Parallel Result' `
            -Severity 'medium' `
            -CheckScript 'Test-EntraSecurityDefaultsEnabled.ps1' `
            -DataNeeds @('entra:securitydefaults')

        $result = InModuleScope Devolutions.CIEM {
            Invoke-CIEMScan -Provider Azure -CheckId 'parallel_result'
        }

        $result | Should -HaveCount 1
        $result[0].GetType().Name | Should -Be 'CIEMScanResult'
        $result[0].Status.ToString() | Should -Be 'PASS'
        $result[0].ResourceId | Should -Be 'resource-1'
        $result[0].Check.Id | Should -Be 'entra_security_defaults_enabled'

        $normalizedResults = @(
            $result |
                ForEach-Object {
                    [pscustomobject]@{
                        CheckId = $_.Check.Id
                        Status = $_.Status.ToString()
                        StatusExtended = $_.StatusExtended
                        ResourceId = $_.ResourceId
                        ResourceName = $_.ResourceName
                        Location = $_.Location
                    }
                }
        )

        ($normalizedResults | ConvertTo-Json -Depth 5) | Should -Be ($script:LegacyScanFixture | ConvertTo-Json -Depth 5)
    }

    It 'Skips parallel dispatch when no checks remain after filtering' {
        Mock -ModuleName Devolutions.CIEM Invoke-CIEMParallelForEach { throw 'should not be called' }

        InModuleScope Devolutions.CIEM {
            { Invoke-CIEMScan -Provider Azure -Service 'DoesNotExist' | Out-Null } | Should -Not -Throw
        }

        Assert-MockCalled Invoke-CIEMParallelForEach -ModuleName Devolutions.CIEM -Times 0 -Exactly
    }

    It 'Does not use an unconditional Connect-CIEM -Force path' {
        $script:ScanSource | Should -Not -Match 'Connect-CIEM\s+-Provider\s+\$providersToConnect\s+-Force'
    }
}