update

pkg/commands/command.go

@@ -2,6 +2,7 @@ package commands
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/briandowns/spinner"
@@ -49,27 +50,27 @@ func GenerateTestModule() *cobra.Command {
 			}
 			spec := prompt.SelectLanguage(specContent)
 
-			ctx := context.Background()
-
-			ai, err := gemini.New(ctx)
-			if err != nil {
-				panic(err)
-			}
-
 			file, err := files.IdentifySpecTypes(specPath)
 			if err != nil {
-				panic(err)
+				return fmt.Errorf("failed to identify spec types: %w", err)
 			}
+
 			reader, err := files.ReadFile(file[0])
 			if err != nil {
-				panic(err)
+				return fmt.Errorf("failed to read spec file: %w", err)
 			}
 
 			s.Start()
+			ctx := context.Background()
+			ai, err := gemini.New(ctx)
+			if err != nil {
+				return fmt.Errorf("failed to create Gemini engine: %w", err)
+			}
 
 			err = ai.GenerateTextStreamWriter(ctx, reader, outputPath, language, spec)
 			if err != nil {
-				panic(err)
+				s.FinalMSG = "Test generation failed: %v\n"
+				return err
 			}
 			s.Stop()
 			return nil
@@ -99,22 +100,22 @@ func GenerateTestScenarios() *cobra.Command {
 			ctx := context.Background()
 			ai, err := gemini.New(ctx)
 			if err != nil {
-				panic(err)
+				return fmt.Errorf("failed to create Gemini engine: %w", err)
 			}
 
 			file, err := files.IdentifySpecTypes(specPath)
 			if err != nil {
-				panic(err)
+				return fmt.Errorf("failed to identify spec types: %w", err)
 			}
 
 			reader, err := files.ReadFile(file[0])
 			if err != nil {
-				panic(err)
+				return fmt.Errorf("failed to read spec file: %w", err)
 			}
 
 			err = ai.GenerateTextStream(ctx, reader, spec)
 			if err != nil {
-				panic(err)
+				return err
 			}
 			return nil
 		},

pkg/engine/files/files.go

@@ -44,7 +44,7 @@ func ReadFile(spec string) ([]genai.Text, error) {
 		bytesRead, err := reader.Read(buffer)
 		if err != nil {
 			if err != io.EOF {
-				fmt.Println("Error reading file:", err)
+				fmt.Errorf("Error reading file:", err)
 			}
 			break
 		}

pkg/engine/prompt/prompt.go

@@ -20,6 +20,7 @@ type PromptContent struct {
 	ErrorMsg string
 	Label    string
 	ItemType string
+	Items    []string
 }
 
 // CompareSpecFiles returns a prompt to compare two spec files
@@ -72,7 +73,7 @@ func SelectLanguage(pc PromptContent) string {
 	return result
 }
 
-func getFrameworksForLanguage(language string) []string {
+func setFrameworksForLanguage(language string) []string {
 	switch language {
 	case "Go":
 		return []string{"Gin", "Echo", "Fiber", "gRPC"}

pkg/github/comment.go

@@ -22,3 +22,13 @@ func (c *Client) ChekPRComment(prNumber int) (string, error) {
 	}
 	return "", nil
 }
+
+func (c *Client) CreatePRDiff(prNumber int) (string, error) {
+	diff, _, err := c.client.PullRequests.GetRaw(c.ctx, owner, repo, prNumber, github.RawOptions{
+		Type: github.Diff,
+	})
+	if err != nil {
+		return "", err
+	}
+	return string(diff), nil
+}

pkg/github/pr.go

@@ -1,6 +1,7 @@
 package github
 
 import (
+	"fmt"
 	"os"
 
 	"github.com/google/go-github/v59/github"
@@ -36,7 +37,17 @@ func (c *Client) ChekPR() (string, error) {
 			if prDetails.GetMerged() {
 				return "", err
 			}
-			return prDetails.GetBody(), nil
+			diff, _, err := c.client.PullRequests.GetRaw(c.ctx, owner, repo, *prDetails.Number, github.RawOptions{
+				Type: github.Diff,
+			})
+			if err != nil {
+				fmt.Println("Error fetching diff:", err)
+				return "", err
+			}
+			if len(diff) == 0 {
+				return "", err
+			}
+			return diff, nil
 		}
 	}
 	return "", err