From f6f1a71a98ee92702970d6a657e1e6579a9cc0f0 Mon Sep 17 00:00:00 2001
From: kavindu <kavindu175@gmail.com>
Date: Thu, 10 Aug 2023 21:08:18 +0530
Subject: [PATCH] Add more tests

---
 .../FunctionCallExpressionTypeFinder.java     |  6 +-
 .../codeaction/CreateFunctionTest.java        |  1 +
 .../config/undefinedFunctionCodeAction31.json | 57 +++++++++++++++++++
 .../source/createUndefinedFunction20.bal      |  8 +++
 4 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 language-server/modules/langserver-core/src/test/resources/codeaction/create-function/config/undefinedFunctionCodeAction31.json
 create mode 100644 language-server/modules/langserver-core/src/test/resources/codeaction/create-function/source/createUndefinedFunction20.bal

diff --git a/language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/command/visitors/FunctionCallExpressionTypeFinder.java b/language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/command/visitors/FunctionCallExpressionTypeFinder.java
index b6511d3fb107..a46e39049002 100644
--- a/language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/command/visitors/FunctionCallExpressionTypeFinder.java
+++ b/language-server/modules/langserver-core/src/main/java/org/ballerinalang/langserver/command/visitors/FunctionCallExpressionTypeFinder.java
@@ -29,6 +29,7 @@
 import io.ballerina.compiler.api.symbols.ParameterSymbol;
 import io.ballerina.compiler.api.symbols.RecordFieldSymbol;
 import io.ballerina.compiler.api.symbols.RecordTypeSymbol;
+import io.ballerina.compiler.api.symbols.StreamTypeSymbol;
 import io.ballerina.compiler.api.symbols.Symbol;
 import io.ballerina.compiler.api.symbols.SymbolKind;
 import io.ballerina.compiler.api.symbols.TypeDescKind;
@@ -490,8 +491,11 @@ public void visit(GroupingKeyVarDeclarationNode groupingKeyVarDeclarationNode) {
     public void visit(SelectClauseNode selectClauseNode) {
         selectClauseNode.parent().parent().accept(this);
         if (resultFound) {
-            if (this.returnTypeSymbol.typeKind() == TypeDescKind.ARRAY) {
+            TypeDescKind kind = this.returnTypeSymbol.typeKind();
+            if (kind == TypeDescKind.ARRAY) {
                 checkAndSetTypeResult(((ArrayTypeSymbol) returnTypeSymbol).memberTypeDescriptor());
+            } else if (kind == TypeDescKind.STREAM) {
+                checkAndSetTypeResult(((StreamTypeSymbol) returnTypeSymbol).typeParameter());
             }
         }
     }
diff --git a/language-server/modules/langserver-core/src/test/java/org/ballerinalang/langserver/codeaction/CreateFunctionTest.java b/language-server/modules/langserver-core/src/test/java/org/ballerinalang/langserver/codeaction/CreateFunctionTest.java
index 53d691b0edc3..642cf73c205b 100644
--- a/language-server/modules/langserver-core/src/test/java/org/ballerinalang/langserver/codeaction/CreateFunctionTest.java
+++ b/language-server/modules/langserver-core/src/test/java/org/ballerinalang/langserver/codeaction/CreateFunctionTest.java
@@ -182,6 +182,7 @@ public Object[][] dataProvider() {
                 {"undefinedFunctionCodeAction28.json"},
                 {"undefinedFunctionCodeAction29.json"},
                 {"undefinedFunctionCodeAction30.json"},
+                {"undefinedFunctionCodeAction31.json"},
         };
     }
 
diff --git a/language-server/modules/langserver-core/src/test/resources/codeaction/create-function/config/undefinedFunctionCodeAction31.json b/language-server/modules/langserver-core/src/test/resources/codeaction/create-function/config/undefinedFunctionCodeAction31.json
new file mode 100644
index 000000000000..340dbc35887d
--- /dev/null
+++ b/language-server/modules/langserver-core/src/test/resources/codeaction/create-function/config/undefinedFunctionCodeAction31.json
@@ -0,0 +1,57 @@
+{
+  "position": {
+    "line": 6,
+    "character": 22
+  },
+  "source": "createUndefinedFunction20.bal",
+  "expected": [
+    {
+      "title": "Create function 'foo(...)'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 7,
+              "character": 1
+            },
+            "end": {
+              "line": 7,
+              "character": 1
+            }
+          },
+          "newText": "\n\nfunction foo(string letter) returns MyType {\n    return {value: \"\"};\n}"
+        }
+      ],
+      "resolvable": true,
+      "data": {
+        "extName": "org.ballerinalang.langserver.codeaction.BallerinaCodeActionExtension",
+        "codeActionName": "Create Function",
+        "fileUri": "createUndefinedFunction20.bal",
+        "range": {
+          "start": {
+            "line": 6,
+            "character": 19
+          },
+          "end": {
+            "line": 6,
+            "character": 30
+          }
+        },
+        "actionData": {
+          "key": "node.range",
+          "value": {
+            "start": {
+              "line": 6.0,
+              "character": 19.0
+            },
+            "end": {
+              "line": 6.0,
+              "character": 30.0
+            }
+          }
+        }
+      }
+    }
+  ]
+}
diff --git a/language-server/modules/langserver-core/src/test/resources/codeaction/create-function/source/createUndefinedFunction20.bal b/language-server/modules/langserver-core/src/test/resources/codeaction/create-function/source/createUndefinedFunction20.bal
new file mode 100644
index 000000000000..d1dc0a649717
--- /dev/null
+++ b/language-server/modules/langserver-core/src/test/resources/codeaction/create-function/source/createUndefinedFunction20.bal
@@ -0,0 +1,8 @@
+type MyType record {|
+    string value;
+|};
+
+function func() {
+    stream<MyType> streamName = stream from string letter in "JohnSnow"
+            select foo(letter);
+}
\ No newline at end of file