From 6e83b225b05e962b85a3e7351b25f7e90aacea7e Mon Sep 17 00:00:00 2001
From: SebastianAtWork <sebastianschwarz1990@gmail.com>
Date: Tue, 14 Nov 2023 12:15:42 +0100
Subject: [PATCH] * feat(FSharp.Editor.Tests): add GetLastDocument method in
 RoslynHelpers.fs * refactor(RefactorTestFramework.fs): replace
 GetSingleDocument with GetLastDocument * feat(AddExplicitReturnTypeTests.fs):
 add test for inferring custom type declared earlier in project

---
 .../Helpers/RoslynHelpers.fs                  |  5 +++
 .../Refactors/AddExplicitReturnTypeTests.fs   | 35 ++++++++++++++++++-
 .../Refactors/RefactorTestFramework.fs        | 18 +++++++---
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/vsintegration/tests/FSharp.Editor.Tests/Helpers/RoslynHelpers.fs b/vsintegration/tests/FSharp.Editor.Tests/Helpers/RoslynHelpers.fs
index 5b00301821bd..01f962266f00 100644
--- a/vsintegration/tests/FSharp.Editor.Tests/Helpers/RoslynHelpers.fs
+++ b/vsintegration/tests/FSharp.Editor.Tests/Helpers/RoslynHelpers.fs
@@ -298,6 +298,11 @@ type RoslynTestHelpers private () =
         let document = project.Documents |> Seq.exactlyOne
         document
 
+    static member GetLastDocument(solution: Solution) =
+        let project = solution.Projects |> Seq.exactlyOne
+        let document = project.Documents |> Seq.last
+        document
+
     static member CreateSolution(syntheticProject: SyntheticProject) =
 
         let checker = syntheticProject.SaveAndCheck()
diff --git a/vsintegration/tests/FSharp.Editor.Tests/Refactors/AddExplicitReturnTypeTests.fs b/vsintegration/tests/FSharp.Editor.Tests/Refactors/AddExplicitReturnTypeTests.fs
index b6bd8ef155f3..e24c698a7bcf 100644
--- a/vsintegration/tests/FSharp.Editor.Tests/Refactors/AddExplicitReturnTypeTests.fs
+++ b/vsintegration/tests/FSharp.Editor.Tests/Refactors/AddExplicitReturnTypeTests.fs
@@ -2,7 +2,7 @@
 
 open Microsoft.VisualStudio.FSharp.Editor
 open Xunit
-open NUnit.Framework
+open System.Linq
 open FSharp.Editor.Tests.Refactors.RefactorTestFramework
 
 [<Theory>]
@@ -174,3 +174,36 @@ let sum a b :MyType= {Value=a+b}
 
     let resultText = newDoc.GetTextAsync context.CT |> GetTaskResult
     Assert.Equal(expectedCode, resultText.ToString())
+
+[<Fact>]
+let ``Correctly infer custom type that is declared earlier in project`` () =
+    let symbolName = "sum"
+
+    let myModule =
+        """
+module MyModule
+type MyType = { Value: int }
+        """
+
+    let code =
+        """
+open MyModule
+
+let sum a b = {Value=a+b}
+        """
+
+    use context = TestContext.CreateWithCodeAndDependency code myModule
+
+    let spanStart = code.IndexOf symbolName
+
+    let newDoc = tryRefactor code spanStart context (new AddExplicitReturnType())
+
+    let expectedCode =
+        """
+open MyModule
+
+let sum a b :MyType= {Value=a+b}
+        """
+
+    let resultText = newDoc.GetTextAsync context.CT |> GetTaskResult
+    Assert.Equal(expectedCode, resultText.ToString())
diff --git a/vsintegration/tests/FSharp.Editor.Tests/Refactors/RefactorTestFramework.fs b/vsintegration/tests/FSharp.Editor.Tests/Refactors/RefactorTestFramework.fs
index cd32d697d034..a74c4ee0cf38 100644
--- a/vsintegration/tests/FSharp.Editor.Tests/Refactors/RefactorTestFramework.fs
+++ b/vsintegration/tests/FSharp.Editor.Tests/Refactors/RefactorTestFramework.fs
@@ -1,6 +1,7 @@
 module FSharp.Editor.Tests.Refactors.RefactorTestFramework
 
 open System
+open System.Linq
 open System.Collections.Immutable
 open System.Collections.Generic
 
@@ -35,16 +36,25 @@ type TestContext(Solution: Solution, CT) =
         let ct = CancellationToken false
         new TestContext(solution, ct)
 
+    static member CreateWithCodeAndDependency (code: string) (codeForPreviousFile: string) =
+        let mutable solution = RoslynTestHelpers.CreateSolution(codeForPreviousFile)
+
+        let firstProject = solution.Projects.First()
+        solution <- solution.AddDocument(DocumentId.CreateNewId(firstProject.Id), "test2.fs", code, filePath = "C:\\test2.fs")
+
+        let ct = CancellationToken false
+        new TestContext(solution, ct)
+
 let mockAction =
     Action<CodeActions.CodeAction, ImmutableArray<Diagnostic>>(fun _ _ -> ())
 
 let tryRefactor (code: string) (cursorPosition) (context: TestContext) (refactorProvider: 'T :> CodeRefactoringProvider) =
     let refactoringActions = new List<CodeAction>()
-    let existingDocument = RoslynTestHelpers.GetSingleDocument context.Solution
+    let existingDocument = RoslynTestHelpers.GetLastDocument context.Solution
 
     context.Solution <- context.Solution.WithDocumentText(existingDocument.Id, SourceText.From(code))
 
-    let document = RoslynTestHelpers.GetSingleDocument context.Solution
+    let document = RoslynTestHelpers.GetLastDocument context.Solution
 
     let mutable workspace = context.Solution.Workspace
 
@@ -70,11 +80,11 @@ let tryRefactor (code: string) (cursorPosition) (context: TestContext) (refactor
 let tryGetRefactoringActions (code: string) (cursorPosition) (context: TestContext) (refactorProvider: 'T :> CodeRefactoringProvider) =
     cancellableTask {
         let refactoringActions = new List<CodeAction>()
-        let existingDocument = RoslynTestHelpers.GetSingleDocument context.Solution
+        let existingDocument = RoslynTestHelpers.GetLastDocument context.Solution
 
         context.Solution <- context.Solution.WithDocumentText(existingDocument.Id, SourceText.From(code))
 
-        let document = RoslynTestHelpers.GetSingleDocument context.Solution
+        let document = RoslynTestHelpers.GetLastDocument context.Solution
 
         let mutable workspace = context.Solution.Workspace