' Licensed to the .NET Foundation under one or more agreements.
' The .NET Foundation licenses this file to you under the MIT license.

Imports System.Windows.Forms.Analyzers.Diagnostics
Imports System.Windows.Forms.VisualBasic.Analyzers.AvoidPassingTaskWithoutCancellationToken
Imports Microsoft.CodeAnalysis
Imports Microsoft.CodeAnalysis.Testing
Imports Microsoft.CodeAnalysis.VisualBasic.Testing
Imports Xunit

Public Class AvoidPassingTaskWithoutCancellationTokenTest
    ' Currently, we do not have Control.InvokeAsync in the .NET 9.0 Windows reference assemblies.
    ' That's why we need to add this Async Control. Once it's there, this test will fail.
    ' We can then remove the AsyncControl and the test will pass, replace AsyncControl with
    ' Control, and the test will pass.
    Private Const AsyncControl As String = "
Imports System
Imports System.Threading
Imports System.Threading.Tasks
Imports System.Windows.Forms

Namespace System.Windows.Forms
    Public Class AsyncControl
        Inherits Control

        ' BEGIN ASYNC API
        Public Function InvokeAsync(callback As Action, Optional cancellationToken As CancellationToken = Nothing) As Task
            Dim tcs As New TaskCompletionSource()

            ' Note: Code is INCORRECT, it's just here to satisfy the compiler!
            Using cancellationToken.Register(Sub() tcs.TrySetCanceled())
                MyBase.BeginInvoke(callback)
            End Using

            Return tcs.Task
        End Function

        Public Function InvokeAsync(Of T)(callback As Func(Of T), Optional cancellationToken As CancellationToken = Nothing) As Task(Of T)
            Dim tcs As New TaskCompletionSource(Of T)()

            ' Note: Code is INCORRECT, it's just here to satisfy the compiler!
            Using cancellationToken.Register(Sub() tcs.TrySetCanceled())
                MyBase.BeginInvoke(callback)
            End Using

            Return tcs.Task
        End Function

        Public Function InvokeAsync(callback As Func(Of CancellationToken, ValueTask), Optional cancellationToken As CancellationToken = Nothing) As Task
            Dim tcs As New TaskCompletionSource()

            ' Note: Code is INCORRECT, it's just here to satisfy the compiler!
            Using cancellationToken.Register(Sub() tcs.TrySetCanceled())
                MyBase.BeginInvoke(callback)
            End Using

            Return tcs.Task
        End Function

        Public Function InvokeAsync(Of T)(callback As Func(Of CancellationToken, ValueTask(Of T)), Optional cancellationToken As CancellationToken = Nothing) As Task(Of T)
            Dim tcs As New TaskCompletionSource(Of T)()

            ' Note: Code is INCORRECT, it's just here to satisfy the compiler!
            Using cancellationToken.Register(Sub() tcs.TrySetCanceled())
                MyBase.BeginInvoke(callback)
            End Using

            Return tcs.Task
        End Function
        ' END ASYNC API
    End Class
End Namespace
"

    Private Const TestCode As String = "
Imports System
Imports System.Threading
Imports System.Threading.Tasks
Imports System.Windows.Forms

Namespace VisualBasicControls

    Public Module Program

        Public Sub Main()

            Dim control As New AsyncControl()

            ' A sync Action delegate is always fine.
            Dim okAction As New Action(Sub() control.Text = ""Hello, World!"")

            ' A sync Func delegate is also fine.
            Dim okFunc As New Func(Of Integer)(Function() 42)

            ' Just a Task we will get in trouble since it's handled as a fire and forget.
            Dim notOkAsyncFunc As New Func(Of Task)(Function() 
                control.Text = ""Hello, World!""
                Return Task.CompletedTask
            End Function)

            ' A Task returning a value will also get us in trouble since it's handled as a fire and forget.
            Dim notOkAsyncFunc2 As New Func(Of Task(Of Integer))(Function() 
                control.Text = ""Hello, World!""
                Return Task.FromResult(42)
            End Function)

            ' OK.
            Dim task1 = control.InvokeAsync(okAction)

            ' Also OK.
            Dim task2 = control.InvokeAsync(okFunc)

            ' Concerning. - Most likely fire and forget by accident. We should warn about this.
            Dim task3 = control.InvokeAsync(notOkAsyncFunc, System.Threading.CancellationToken.None)

            ' Again: Concerning. - Most likely fire and forget by accident. We should warn about this.
            Dim task4 = control.InvokeAsync(notOkAsyncFunc, System.Threading.CancellationToken.None)

            ' And again concerning. - We should warn about this, too.
            Dim task5 = control.InvokeAsync(notOkAsyncFunc2, System.Threading.CancellationToken.None)

            ' This is OK, since we're passing a cancellation token.
            Dim okAsyncFunc = New Func(Of CancellationToken, ValueTask)(Function(cancellation) 
                control.Text = ""Hello, World!""
                Return ValueTask.CompletedTask
            End Function)

            ' This is also OK, again, because we're passing a cancellation token.
            Dim okAsyncFunc2 = New Func(Of CancellationToken, ValueTask(Of Integer))(Function(cancellation) 
                control.Text = ""Hello, World!""
                Return ValueTask.FromResult(42)
            End Function)

            ' And let's test that, too:
            Dim task6 = control.InvokeAsync(okAsyncFunc, System.Threading.CancellationToken.None)

            ' And that, too:
            Dim task7 = control.InvokeAsync(okAsyncFunc2, System.Threading.CancellationToken.None)

        End Sub

    End Module

End Namespace

"

    Public Shared Iterator Function GetReferenceAssemblies() As IEnumerable(Of Object())
        Yield {ReferenceAssemblies.Net.Net90Windows}
    End Function

    <Theory>
    <MemberData(NameOf(GetReferenceAssemblies))>
    Public Async Function VB_AvoidPassingFuncReturningTaskWithoutCancellationAnalyzer(referenceAssemblies As ReferenceAssemblies) As Task
        ' If the API does not exist, we need to add it to the test.
        Dim customControlSource As String = AsyncControl
        Dim diagnosticId As String = DiagnosticIDs.AvoidPassingFuncReturningTaskWithoutCancellationToken

        Dim context As New VisualBasicAnalyzerTest(Of AvoidPassingTaskWithoutCancellationTokenAnalyzer, DefaultVerifier) With
            {
                .TestCode = TestCode,
                .ReferenceAssemblies = referenceAssemblies
            }

        context.TestState.OutputKind = OutputKind.WindowsApplication
        context.TestState.Sources.Add(customControlSource)
        context.TestState.ExpectedDiagnostics.AddRange(
                {
                    DiagnosticResult.CompilerWarning(diagnosticId).WithSpan(40, 25, 40, 101),
                    DiagnosticResult.CompilerWarning(diagnosticId).WithSpan(43, 25, 43, 101),
                    DiagnosticResult.CompilerWarning(diagnosticId).WithSpan(46, 25, 46, 102)
                })

        Await context.RunAsync().ConfigureAwait(continueOnCapturedContext:=True)
    End Function

End Class
