NCache.Azure

1.0.9-beta1

frameworkResources/Scripts/test_stress.ps1

Param(
    [Parameter(Mandatory)]
    [string]$EnvironmentName,
    [Parameter(Mandatory)]
    [string]$CacheName,
    [string]$Server,
    [int]$DataSize,
    [int]$ItemsCount = 1000,
    [int]$SlidingExpiration,
    [int]$TestCaseIterationDelay,
    [int]$TestCaseIterations,
    [int]$ThreadCount,
    [int]$UpdatesPerIteration,
    [pscredential]$Credentials,
    [switch]$NoLogo,
    [string]$ScriptsFolderPath = ".\Resources\Scripts"
)

. "$ScriptsFolderPath\Resources_Scripts\virtual_machine.ps1"
. "$ScriptsFolderPath\common.ps1"

function GetUploadedVmNames {
    param(
        [Parameter(Mandatory)][string]$ResourceGroup
    )
    
    $serverIps = $Server.Split(",") | ForEach-Object { $_.Trim() }

    $vms = Get-AzVM -ResourceGroupName $ResourceGroup -ErrorAction Stop |
    Where-Object { $_.Tags.ContainsKey("InstallMode") }

    $filteredVms = foreach ($vm in $vms) {
        # Get NIC name from VM network profile
        $nicId = $vm.NetworkProfile.NetworkInterfaces[0].Id
        $nicName = ($nicId -split '/')[ -1 ]

        # Retrieve NIC object
        $nic = Get-AzNetworkInterface -Name $nicName -ResourceGroupName $ResourceGroup -ErrorAction SilentlyContinue
        if ($null -eq $nic) { continue }

        # Get private IPs and check if any match the input list
        $privateIps = $nic.IpConfigurations | ForEach-Object { $_.PrivateIpAddress }
        if ($privateIps | Where-Object { $serverIps -contains $_ }) {
            $vm
        }
    }

    return $filteredVms | Select-Object -ExpandProperty Name
}

function InvokeCommandOnServerWindows {
    param(
        [string]$ResourceGroupName,
        [string[]]$VmNames
    )
    $command = "Test-Stress -CacheName $CacheName"
    if ($DataSize) {
        $command += " -DataSize $DataSize"
    }
    if ($ItemsCount) {
        $command += " -ItemsCount $ItemsCount"
    }
    if ($SlidingExpiration) {
        $command += " -SlidingExpiration $SlidingExpiration"
    }
    if ($TestCaseIterationDelay) {
        $command += " -TestCaseIterationDelay $TestCaseIterationDelay"
    }
    if ($TestCaseIterations) {
        $command += " -TestCaseIterations $TestCaseIterations"
    }
    if ($ThreadCount) {
        $command += " -ThreadCount $ThreadCount"
    }
    if ($UpdatesPerIteration) {
        $command += " -UpdatesPerIteration $UpdatesPerIteration"
    }
    if ($Credentials) {
        $command += " -Credentials $Credentials"
    }
    if ($NoLogo) {
        $command += " -NoLogo"
    }
    $script = @"
$command > output.txt
Write-Output (Get-Content output.txt)
Remove-Item output.txt
"@

    $jobs = @()
    foreach ($vm in $VmNames) {
        $jobs += Start-Job -ScriptBlock {
            param($ResourceGroupName, $vm, $script)
            $result = Invoke-AzVMRunCommand -ResourceGroupName $ResourceGroupName -VMName $vm -CommandId 'RunPowerShellScript' -ScriptString $script
            $messages = $result.Value | ForEach-Object { $_.Message }
            $stdErr = ($result.Value | Where-Object { $_.Code -match "StdErr" }).Message
            $hasError = $false
            if (-not [string]::IsNullOrWhiteSpace($stdErr)) {
                $hasError = $true
            }
            $joined = ($messages -join "`n").Trim()
            if ($hasError) {
                Write-Host $joined -ForegroundColor Red
            }
            else {
                Write-Host $joined
            }
        } -ArgumentList $ResourceGroupName, $vm, $script
    }

    $jobs | Wait-Job | Out-Null
    $jobs | Receive-Job
}

function InvokeCommandOnServerLinux {
    param(
        [string]$ResourceGroupName,
        [string[]]$VmNames
    )


    $jobs = @()
    foreach ($vm in $VmNames) {
        
        $NCachePath = Get-LinuxNCacheInstallDir -VMName $vm -ResourceGroupName $ResourceGroupName

        $command = "$NCachePath/bin/tools/test-stress -cachename $CacheName"
        if ($DataSize) {
            $command += " -datasize $DataSize"
        }
        if ($ItemsCount) {
            $command += " -itemscount $ItemsCount"
        }
        if ($SlidingExpiration) {
            $command += " -slidingexpiration $SlidingExpiration"
        }
        if ($TestCaseIterationDelay) {
            $command += " -testcaseiterationdelay $TestCaseIterationDelay"
        }
        if ($TestCaseIterations) {
            $command += " -testcaseiterations $TestCaseIterations"
        }
        if ($ThreadCount) {
            $command += " -threadcount $ThreadCount"
        }
        if ($UpdatesPerIteration) {
            $command += " -updatesperiteration $UpdatesPerIteration"
        }
        if ($Credentials) {
            $command += " -credentials `"$($Credentials.GetNetworkCredential().UserName):$($Credentials.GetNetworkCredential().Password)`""
        }
        if ($NoLogo) {
            $command += " -nologo"
        }
        $script = @"
$command > output.txt
cat output.txt
rm -f output.txt
"@
        $jobs += Start-Job -ScriptBlock {
            param($ResourceGroupName, $vm, $script)
            $result = Invoke-AzVMRunCommand -ResourceGroupName $ResourceGroupName -VMName $vm -CommandId 'RunShellScript' -ScriptString $script
            $hasError = $false
            $messages = $result.Value | ForEach-Object { $_.Message }
            $joined = ($messages -join "`n").Trim()
            $message = $result.Value[0].Message
        
            if ($message -match "Error:" -or $message -match "error:") {
                $hasError = $true
            }
        
            $parts = $message -split "\[stderr\]"
        
            if ($parts.Count -gt 1 -and $parts[1].Trim().Length -gt 0) {
                $hasError = $true
            }

            if ($hasError) {
                Write-Host $joined -ForegroundColor Red
            }
            else {
                Write-Host $joined
            }
            
        } -ArgumentList $ResourceGroupName, $vm, $script
    }

    $jobs | Wait-Job | Out-Null
    $jobs | Receive-Job
}

function InvokeCommandOnVms {
    param(
        [object]$Resource,
        [string]$ResourceGroupName,
        [string[]]$VmNames
    )

    $os = $Resource.Tags['OsType']
    if ($os -eq "Windows") {
        InvokeCommandOnServerWindows -ResourceGroupName $ResourceGroupName -VmNames $VmNames
    }
    else {
        InvokeCommandOnServerLinux -ResourceGroupName $ResourceGroupName -VmNames $VmNames
    }
}

function ExecuteCommands {
    $resource = CheckIfEnvironmentExits -Environment $EnvironmentName

    $resourceGroupName = $resource.ResourceGroupName
    $vmNames = GetUploadedVmNames -ResourceGroup $resourceGroupName

    InvokeCommandOnVms -Resource $resource -ResourceGroupName $resourceGroupName -VmNames $vmNames
}

try {
    ShowExecutingCommand
    if (-not (Get-AzContext)) {
        NoContext
        Connect-AzAccount
        if (Get-AzContext) {
            ExecuteCommands
        }
    }
    else {
        ExecuteCommands
    }
}
catch {
    Write-Error $($_.Exception.Message)
}