Tests/GenXdev.AI.ComfyUI/EnsureComfyUIModel.Tests.ps1

Pester\Describe 'EnsureComfyUIModel' {

    Pester\It 'Should have valid download URLs for all supported models' {

        # Load the supported models configuration
        $jsonPath = Microsoft.PowerShell.Management\Join-Path $PSScriptRoot "..\..\Functions\GenXdev.AI.ComfyUI\SupportedComfyUIModels.json"
        $jsonPath | Pester\Should -Exist -Because 'SupportedComfyUIModels.json should exist'

        $supportedModels = Microsoft.PowerShell.Management\Get-Content -LiteralPath $jsonPath -Raw | Microsoft.PowerShell.Utility\ConvertFrom-Json
        $supportedModels | Pester\Should -Not -BeNullOrEmpty -Because 'SupportedComfyUIModels.json should contain model definitions'

        # Validate each model has required properties
        foreach ($model in $supportedModels) {
            $model.Name | Pester\Should -Not -BeNullOrEmpty -Because "Model should have a Name property"
            $model.DownloadUrl | Pester\Should -Not -BeNullOrEmpty -Because "Model '$($model.Name)' should have a DownloadUrl property"
            $model.HuggingFaceRepo | Pester\Should -Not -BeNullOrEmpty -Because "Model '$($model.Name)' should have a HuggingFaceRepo property"
            $model.FileName | Pester\Should -Not -BeNullOrEmpty -Because "Model '$($model.Name)' should have a FileName property"
            $model.Architecture | Pester\Should -Not -BeNullOrEmpty -Because "Model '$($model.Name)' should have an Architecture property"
            $model.Compatible | Pester\Should -Be $true -Because "Model '$($model.Name)' should be marked as Compatible"
            $model.SizeMB | Pester\Should -BeGreaterThan 0 -Because "Model '$($model.Name)' should have a valid SizeMB property"
        }

        # Test download URLs accessibility using lightweight HTTP HEAD requests
        # This avoids downloading huge model files while validating they exist
        $failedModels = @()
        foreach ($model in $supportedModels) {
            $modelName = $model.Name
            $downloadUrl = $model.DownloadUrl
            $fileName = $model.FileName

            Microsoft.PowerShell.Utility\Write-Verbose "Testing download URL accessibility for model: ${modelName}"
            Microsoft.PowerShell.Utility\Write-Verbose "URL: ${downloadUrl}"

            try {
                # Use HTTP HEAD request to check if the file exists without downloading
                # This is extremely fast and lightweight compared to downloading GB files
                $headResponse = Microsoft.PowerShell.Utility\Invoke-WebRequest `
                    -Uri $downloadUrl `
                    -Method Head `
                    -TimeoutSec 30 `
                    -ErrorAction Stop

                # Check that we got a successful response (200 OK)
                if ($headResponse.StatusCode -eq 200) {

                    # Validate the response indicates a file download (not an HTML page)
                    $contentType = $headResponse.Headers['Content-Type']
                    $contentLength = $headResponse.Headers['Content-Length']

                    if ($contentType -and $contentType -match 'text/html') {
                        throw "URL returns HTML page instead of file (likely a redirect or error page)"
                    }

                    if ($contentLength) {
                        # Handle both single string and array responses for Content-Length
                        $contentLengthValue = if ($contentLength -is [array]) {
                            $contentLength[0]
                        } else {
                            $contentLength
                        }

                        # Safely convert to int64, handling potential parsing issues
                        $fileSizeBytes = 0
                        if ([int64]::TryParse($contentLengthValue, [ref]$fileSizeBytes)) {
                            $fileSizeMB = [Math]::Round($fileSizeBytes / 1MB, 2)
                            Microsoft.PowerShell.Utility\Write-Verbose "✓ File '${fileName}' is accessible (${fileSizeMB} MB)"

                            # Validate that the reported size roughly matches the JSON (allow 10% variance)
                            $expectedSizeMB = $model.SizeMB
                            $sizeDifference = [Math]::Abs($fileSizeMB - $expectedSizeMB)
                            $allowedVariance = $expectedSizeMB * 0.1  # 10% tolerance

                            if ($sizeDifference -gt $allowedVariance) {
                                Microsoft.PowerShell.Utility\Write-Warning "Size mismatch for '${fileName}': Expected ${expectedSizeMB} MB, got ${fileSizeMB} MB (difference: ${sizeDifference} MB)"
                            }
                        } else {
                            Microsoft.PowerShell.Utility\Write-Verbose "✓ File '${fileName}' is accessible (size could not be parsed: ${contentLengthValue})"
                        }
                    } else {
                        Microsoft.PowerShell.Utility\Write-Verbose "✓ File '${fileName}' is accessible (size unknown)"
                    }
                } else {
                    throw "HTTP status ${($headResponse.StatusCode)}: $($headResponse.StatusDescription)"
                }

            } catch {
                $errorMessage = $_.Exception.Message
                $failedModels += @{
                    Name = $modelName
                    Url = $downloadUrl
                    File = $fileName
                    Error = $errorMessage
                }
                Microsoft.PowerShell.Utility\Write-Warning "✗ Model '${modelName}' download URL validation failed: ${errorMessage}"
            }
        }        # Report all failed URLs at once for easier debugging
        if ($failedModels.Count -gt 0) {
            $failureReport = $failedModels | Microsoft.PowerShell.Core\ForEach-Object {
                "- $($_.Name) (URL: $($_.Url), File: $($_.File)): $($_.Error)"
            } | Microsoft.PowerShell.Utility\Join-String -Separator "`n"

            throw "The following $($failedModels.Count) model download URLs are not accessible and need to be updated in SupportedComfyUIModels.json:`n${failureReport}"
        }        # Verify all models from ValidateSet are included in JSON
        # Load the model names dynamically from the JSON file
        $validateSetModels = $supportedModels | Microsoft.PowerShell.Core\ForEach-Object { $_.Name }
        $jsonModelNames = $supportedModels | Microsoft.PowerShell.Core\ForEach-Object { $_.Name }

        foreach ($validateSetModel in $validateSetModels) {
            $jsonModelNames | Pester\Should -Contain $validateSetModel -Because "ValidateSet model '${validateSetModel}' should be defined in SupportedComfyUIModels.json"
        }
    }

    Pester\It 'Should detect and mark incompatible models for removal' {

        # Load the supported models configuration
        $jsonPath = Microsoft.PowerShell.Management\Join-Path $PSScriptRoot "..\..\Functions\GenXdev.AI.ComfyUI\SupportedComfyUIModels.json"
        $supportedModels = Microsoft.PowerShell.Management\Get-Content -LiteralPath $jsonPath -Raw | Microsoft.PowerShell.Utility\ConvertFrom-Json

        # Check if ComfyUI is available for testing
        $comfyUIAvailable = $false
        try {
            # Try to connect to ComfyUI API using script variable if available
            $apiUrl = if ($script:comfyUIApiUrl) { $script:comfyUIApiUrl } else { "http://localhost:8188" }
            $null = Microsoft.PowerShell.Utility\Invoke-RestMethod -Uri "${apiUrl}/object_info" -Method GET -TimeoutSec 3 -ErrorAction Stop
            $comfyUIAvailable = $true
            Microsoft.PowerShell.Utility\Write-Verbose "ComfyUI is available for model compatibility testing"
        } catch {
            Microsoft.PowerShell.Utility\Write-Verbose "ComfyUI not available for testing, skipping model compatibility checks"
        }

        $incompatibleModels = @()

        if ($comfyUIAvailable) {
            # Test each model for basic compatibility without downloading
            foreach ($model in $supportedModels) {
                $modelName = $model.Name
                Microsoft.PowerShell.Utility\Write-Verbose "Testing model compatibility: $modelName"

                try {
                    # Verify model has architecture info in JSON (this should work for all supported models)
                    if (-not $model.Architecture -or [string]::IsNullOrWhiteSpace($model.Architecture)) {
                        throw "Model does not have architecture information in JSON"
                    }

                    Microsoft.PowerShell.Utility\Write-Verbose "✓ Model '$modelName' appears compatible (Architecture: $($model.Architecture))"
                } catch {
                    Microsoft.PowerShell.Utility\Write-Warning "✗ Model '$modelName' failed compatibility check: $_"
                    $incompatibleModels += @{
                        Name = $modelName
                        FileName = $model.FileName
                        Reason = $_.Exception.Message
                    }
                }
            }

            # If we found incompatible models, provide recommendations
            if ($incompatibleModels.Count -gt 0) {
                $removalSuggestions = $incompatibleModels | Microsoft.PowerShell.Core\ForEach-Object {
                    "- $($_.Name) (File: $($_.FileName)): $($_.Reason)"
                } | Microsoft.PowerShell.Utility\Join-String -Separator "`n"

                Microsoft.PowerShell.Utility\Write-Warning "Found $($incompatibleModels.Count) potentially incompatible models:`n$removalSuggestions"
                Microsoft.PowerShell.Utility\Write-Warning "Consider reviewing these models for removal from SupportedComfyUIModels.json"

                # For now, we'll warn but not fail the test
                # In the future, you could uncomment the line below to fail on incompatible models
                # throw "Incompatible models detected. Please review and remove them from the configuration."
            }
        }

        # Always pass this test since we're just providing warnings/recommendations
        $true | Pester\Should -Be $true -Because "Model compatibility check completed (warnings provided for any issues)"
    }

    Pester\It 'Should correctly map available model files to supported models using FileName property' {

        # Load the supported models configuration
        $jsonPath = Microsoft.PowerShell.Management\Join-Path $PSScriptRoot "..\..\Functions\GenXdev.AI.ComfyUI\SupportedComfyUIModels.json"
        $supportedModels = Microsoft.PowerShell.Management\Get-Content -LiteralPath $jsonPath -Raw | Microsoft.PowerShell.Utility\ConvertFrom-Json

        # Test that each model's FileName property correctly matches what ComfyUI would list
        foreach ($model in $supportedModels) {
            $fileName = $model.FileName
            $downloadUrl = $model.DownloadUrl

            # Verify FileName property is properly set
            $fileName | Pester\Should -Not -BeNullOrEmpty -Because "Model '$($model.Name)' should have a FileName property"

            # Test the mapping logic that was fixed
            # Simulate what Invoke-ComfyUIImageGeneration does when mapping available files
            $matchingModel = $supportedModels | Microsoft.PowerShell.Core\Where-Object {
                $expectedFile = $_.FileName  # This should match the corrected logic
                $fileName -eq $expectedFile
            }

            $matchingModel | Pester\Should -Not -BeNullOrEmpty -Because "Model '$($model.Name)' with FileName '$fileName' should match itself in mapping logic"
            $matchingModel.Name | Pester\Should -Be $model.Name -Because "Mapped model should be the same model"

            # Verify the old broken logic would have failed for some models
            $urlFileName = [System.IO.Path]::GetFileName($downloadUrl)
            if ($urlFileName -ne $fileName) {
                Microsoft.PowerShell.Utility\Write-Verbose "Model '$($model.Name)': URL filename '$urlFileName' differs from configured FileName '$fileName' - this validates the fix"
            }
        }

        Microsoft.PowerShell.Utility\Write-Verbose "Model filename mapping logic validation completed successfully"
    }

    Pester\It 'Should accept all ValidateSet model names' {

        # Load the supported models configuration to get current model names
        $jsonPath = Microsoft.PowerShell.Management\Join-Path $PSScriptRoot "..\..\Functions\GenXdev.AI.ComfyUI\SupportedComfyUIModels.json"
        $supportedModels = Microsoft.PowerShell.Management\Get-Content -LiteralPath $jsonPath -Raw | Microsoft.PowerShell.Utility\ConvertFrom-Json

        # Test that the function accepts each model name from the JSON file
        $validateSetModels = $supportedModels | Microsoft.PowerShell.Core\ForEach-Object { $_.Name }

        foreach ($modelName in $validateSetModels) {
            # This should not throw a parameter validation error
            {
                # Just test parameter validation by calling Get-Command with the parameter
                $null = Microsoft.PowerShell.Core\Get-Command GenXdev.AI\EnsureComfyUIModel -ArgumentList $modelName
            } | Pester\Should -Not -Throw -Because "Model name '${modelName}' should be accepted by ValidateSet"
        }
    }
}