diff --git a/packages/ts-generator/src/nextjs/__tests__/nextjs.test.ts b/packages/ts-generator/src/nextjs/__tests__/nextjs.test.ts
new file mode 100644
index 0000000..ca04b63
--- /dev/null
+++ b/packages/ts-generator/src/nextjs/__tests__/nextjs.test.ts
@@ -0,0 +1,96 @@
+import { Project } from "ts-morph";
+import {
+  analyzeNextjsSourceFiles,
+  generateServerAction,
+  ServerActionInfo,
+} from "../index";
+
+describe("Next.js functionality", () => {
+  let project: Project;
+
+  beforeEach(() => {
+    project = new Project({ useInMemoryFileSystem: true });
+  });
+
+  describe("analyzeNextjsSourceFiles", () => {
+    it("should correctly identify and parse server actions", () => {
+      const serverActionCode = `
+        "use server";
+        
+        export async function submitForm(data: FormData) {
+          // Server action logic
+        }
+
+        export async function deleteItem(id: string) {
+          // Delete item logic
+        }
+      `;
+
+      const sourceFile = project.createSourceFile(
+        "app/actions.ts",
+        serverActionCode
+      );
+
+      const result = analyzeNextjsSourceFiles([sourceFile]);
+
+      expect(result.serverActions).toHaveLength(2);
+      expect(result.serverActions[0]?.name).toBe("submitForm");
+      expect(result.serverActions[1]?.name).toBe("deleteItem");
+    });
+
+    it("should not identify non-server actions", () => {
+      const regularCode = `
+        export function regularFunction() {
+          // Regular function logic
+        }
+      `;
+
+      const sourceFile = project.createSourceFile(
+        "app/regular.ts",
+        regularCode
+      );
+
+      const result = analyzeNextjsSourceFiles([sourceFile]);
+
+      expect(result.serverActions).toHaveLength(0);
+    });
+  });
+
+  describe("generateServerAction", () => {
+    it("should generate correct server action code", () => {
+      const serverActionInfo: ServerActionInfo = {
+        name: "testAction",
+        returnType: "Promise<string>", // This could be 'string' or 'Promise<string>', the function should handle both
+        parameters: [
+          { name: "data", type: "FormData" },
+          { name: "userId", type: "string" },
+        ],
+        filePath: "app/actions.ts",
+      };
+
+      const generatedCode = generateServerAction(serverActionInfo);
+
+      expect(generatedCode).toContain('"use server";');
+      expect(generatedCode).toContain(
+        "export async function testAction(data: FormData, userId: string): Promise<string>"
+      );
+      expect(generatedCode).toContain("// TODO: Implement server action logic");
+      expect(generatedCode).toContain('throw new Error("Not implemented");');
+    });
+
+    it("should handle non-Promise return types", () => {
+      const serverActionInfo: ServerActionInfo = {
+        name: "testAction",
+        returnType: "void",
+        parameters: [],
+        filePath: "app/actions.ts",
+      };
+
+      const generatedCode = generateServerAction(serverActionInfo);
+
+      expect(generatedCode).toContain(
+        "export async function testAction(): Promise<void>"
+      );
+    });
+  });
+});
diff --git a/packages/ts-generator/src/nextjs/generator.ts b/packages/ts-generator/src/nextjs/generator.ts
index 5c989f6..35cab2c 100644
--- a/packages/ts-generator/src/nextjs/generator.ts
+++ b/packages/ts-generator/src/nextjs/generator.ts
@@ -1,5 +1,3 @@
-// src/nextjs/generator.ts
-
 import { ServerActionInfo } from "./types";
 
 export function generateServerAction(info: ServerActionInfo): string {
@@ -7,9 +5,13 @@ export function generateServerAction(info: ServerActionInfo): string {
     .map((param) => `${param.name}: ${param.type}`)
     .join(", ");
 
+  let returnType = info.returnType.replace(/Promise<(.*)>/, "$1");
+
+  returnType = `Promise<${returnType}>`;
+
   return `"use server";
 
-export async function ${info.name}(${parameters}): Promise<${info.returnType}> {
+export async function ${info.name}(${parameters}): ${returnType} {
   // TODO: Implement server action logic
   throw new Error("Not implemented");
 }`;
diff --git a/packages/ts-generator/src/nextjs/index.ts b/packages/ts-generator/src/nextjs/index.ts
index b3f1967..5dca3b3 100644
--- a/packages/ts-generator/src/nextjs/index.ts
+++ b/packages/ts-generator/src/nextjs/index.ts
@@ -1,9 +1,3 @@
-// Named exports for types
 export type { NextCodebaseInfo, ServerActionInfo } from "./types";
-
-// Named exports for functions
-export { scanNextjsCodebase } from "./scanner";
+export { scanNextjsCodebase, analyzeNextjsSourceFiles } from "./scanner";
 export { generateServerAction } from "./generator";
-
-// Optionally, you can also include a convenience export
-export * as NextJS from "./";
diff --git a/packages/ts-generator/src/nextjs/scanner.ts b/packages/ts-generator/src/nextjs/scanner.ts
index 248feee..0dae59e 100644
--- a/packages/ts-generator/src/nextjs/scanner.ts
+++ b/packages/ts-generator/src/nextjs/scanner.ts
@@ -1,16 +1,18 @@
-// src/nextjs/scanner.ts
-
-import { Project, SourceFile, Node } from "ts-morph";
-import { FunctionInfo } from "../types";
+import { Project, SourceFile } from "ts-morph";
 import { ServerActionInfo, NextCodebaseInfo } from "./types";
 
 export function scanNextjsCodebase(projectPath: string): NextCodebaseInfo {
   const project = new Project();
   project.addSourceFilesAtPaths(`${projectPath}/**/*.ts`);
+  return analyzeNextjsSourceFiles(project.getSourceFiles());
+}
 
+export function analyzeNextjsSourceFiles(
+  sourceFiles: SourceFile[]
+): NextCodebaseInfo {
   const serverActions: ServerActionInfo[] = [];
 
-  project.getSourceFiles().forEach((sourceFile) => {
+  sourceFiles.forEach((sourceFile) => {
     if (isServerActionFile(sourceFile)) {
       serverActions.push(...extractServerActions(sourceFile));
     }
@@ -28,19 +30,17 @@ function extractServerActions(sourceFile: SourceFile): ServerActionInfo[] {
 
   sourceFile.getFunctions().forEach((func) => {
     if (func.isAsync()) {
-      const functionInfo: FunctionInfo = {
+      const functionInfo: ServerActionInfo = {
         name: func.getName() || "anonymous",
         returnType: func.getReturnType().getText(),
         parameters: func.getParameters().map((param) => ({
           name: param.getName(),
           type: param.getType().getText(),
         })),
+        filePath: sourceFile.getFilePath(),
       };
 
-      serverActions.push({
-        ...functionInfo,
-        filePath: sourceFile.getFilePath(),
-      });
+      serverActions.push(functionInfo);
     }
   });
 
diff --git a/packages/ts-generator/src/nextjs/types.ts b/packages/ts-generator/src/nextjs/types.ts
index 59d01f8..9ddbcc3 100644
--- a/packages/ts-generator/src/nextjs/types.ts
+++ b/packages/ts-generator/src/nextjs/types.ts
@@ -1,4 +1,4 @@
-import { FunctionInfo } from "types";
+import { FunctionInfo } from "../types";
 
 export interface ServerActionInfo extends FunctionInfo {
   filePath: string;