add tests for ollama function and extraction function

llama/extraction/extract.go

@@ -8,13 +8,29 @@ var GoPrompt = "The code should be in the Go programming language. There should
 
 var RustPrompt = "The code should be in the Rust programming language. There should also be 3 robust test cases within the same code. There should also be a main function inside of which all the execution takes place. Please only provide the source code and no further explanation, The format should be ```rust <yourcode + testcases> ```"
 
+// func Extract(output string) string {
+// 	parts := strings.Split(output, "```")
+// 	var extracted = ""
+// 	if strings.Contains(parts[1], "rust") {
+// 		extracted = strings.TrimLeft(parts[1], "rust")
+// 	} else {
+// 		extracted = strings.TrimLeft(parts[1], "go")
+// 	}
+// 	return extracted
+// }
+
+// Extract extracts the code snippet between ``` and removes the language identifier.
 func Extract(output string) string {
 	parts := strings.Split(output, "```")
-	var extracted = ""
-	if strings.Contains(parts[1], "rust") {
-		extracted = strings.TrimLeft(parts[1], "rust")
-	} else {
-		extracted = strings.TrimLeft(parts[1], "go")
+	if len(parts) < 2 {
+		return "" // Handle the case if format is incorrect: Return empty string
+	}
+
+	// Trim the language identifier like `go` or `rust` from the code
+	code := parts[1]
+	lines := strings.SplitN(code, "\n", 2)
+	if len(lines) > 1 {
+		return "\n" + lines[1] // Return the code without the first line (language identifier)
 	}
-	return extracted
+	return ""
 }

llama/extraction/extract_test.go

+package extraction
+
+import (
+	"testing"
+)
+
+// var GoInputs []string = []string{"```go\nfunc main() {}\n```", "```go\nfmt.Println('Hello World')\n```"}
+// var expectedGoOuputs []string = []string{"\nfunc main() {}\n", "\nfmt.Println('Hello World')\n"}
+
+// var RustInputs []string = []string{"```go\nfn main() {}\n```", "```go\nprintln!('Hello World')\n```"}
+// var expectedRustOuputs []string = []string{"\nfn main() {}\n", "\nprintln!('Hello World')\n"}
+
+// func TestExtraction(t *testing.T) {
+
+// 	t.Run("Golang Extraction 1", func(t *testing.T) {
+// 		var output = Extract(GoInputs[0])
+// 		if output != expectedGoOuputs[0] {
+// 			t.Error(output)
+// 		}
+// 	})
+
+// 	t.Run("Golang Extraction 2", func(t *testing.T) {
+// 		var output = Extract(GoInputs[1])
+// 		if output != expectedGoOuputs[1] {
+// 			t.Error(output)
+// 		}
+// 	})
+
+// 	t.Run("Rust Extraction 1", func(t *testing.T) {
+// 		var output = Extract(RustInputs[0])
+// 		if output != expectedRustOuputs[0] {
+// 			t.Error(output)
+// 		}
+// 	})
+
+// 	t.Run("Rust Extraction 2", func(t *testing.T) {
+// 		var output = Extract(RustInputs[1])
+// 		if output != expectedRustOuputs[1] {
+// 			t.Error(output)
+// 		}
+// 	})
+
+// }
+
+// Inputs and Expected Outputs for the Test Cases
+// This can be considered a table-driven test.
+var testCases = []struct {
+	name     string
+	input    string
+	expected string
+}{
+	// Go Test Cases
+	{"Go Extraction 1", "```go\nfunc main() {}\n```", "\nfunc main() {}\n"},
+	{"Go Extraction 2", "```go\nfmt.Println('Hello World')\n```", "\nfmt.Println('Hello World')\n"},
+
+	// Rust Test Cases
+	{"Rust Extraction 1", "```rust\nfn main() {}\n```", "\nfn main() {}\n"},
+	{"Rust Extraction 2", "```rust\nprintln!('Hello World')\n```", "\nprintln!('Hello World')\n"},
+}
+
+// Refined Test Function using Table-Driven Approach
+func TestExtraction(t *testing.T) {
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			output := Extract(tc.input)
+			if output != tc.expected {
+				t.Errorf("Test %s failed: Expected %q, got %q", tc.name, tc.expected, output)
+			}
+		})
+	}
+}

llama/ollama-implementation/ollama_test.go

@@ -3,6 +3,7 @@ package ollamaimplementation
 import (
 	"encoding/json"
 	"io/ioutil"
+	"llama/extraction"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -73,3 +74,40 @@ func TestGetOllamaResponse(t *testing.T) {
 		}
 	}
 }
+
+var promptTestCases = []struct {
+	name          string
+	prompt        string
+	suffixStr     string
+	shouldContain []string
+}{
+	{"5 Even Integers GO", "Write a program that generates 5 random integers.", extraction.GoPrompt, []string{"```go", "```"}},
+	{"Sort the array using mergesort GO", "Write a program that sorts the array [23, 2, 0, -1, 89, 500] using mergesort.", extraction.GoPrompt, []string{"```go", "```"}},
+	{"Reverse the string GO.", "Reverse the string 'ammar'", extraction.GoPrompt, []string{"```go", "```"}},
+
+	{"5 Even Integers rust", "Write a program that generates 5 random integers.", extraction.RustPrompt, []string{"```rust", "```"}},
+	{"Sort the array using mergesort rust", "Write a program that sorts the array [23, 2, 0, -1, 89, 500] using mergesort.", extraction.RustPrompt, []string{"```rust", "```"}},
+	{"Reverse the string rust.", "Reverse the string 'ammar'", extraction.RustPrompt, []string{"```rust", "```"}},
+}
+
+func TestGetOllamaResponsePrompts(t *testing.T) {
+	for _, tc := range promptTestCases {
+		t.Run(tc.name, func(t *testing.T) {
+			response, _, _ := GetOllamaResponse(tc.prompt+tc.suffixStr, []int{})
+			var prefix string
+			if tc.shouldContain[0] == "```go" {
+				prefix = response[0:5]
+				if prefix != tc.shouldContain[0] {
+					t.Errorf("Test faild expected %v got %v", tc.shouldContain[0], response[0:5])
+				}
+			} else {
+				prefix = response[0:7]
+				if prefix != tc.shouldContain[1] {
+					t.Errorf("Test faild expected %v got %v", tc.shouldContain[0], response[0:7])
+				}
+			}
+		})
+
+	}
+
+}

src_old/compiler/compiler.go

+package compiler
+
+import (
+	"log"
+	"os"
+	"os/exec"
+	"runtime"
+	"strings"
+)
+
+const TempOutputDir = "tempOutput/"
+const TempModuleName = "tempOutput"
+
+type Language string
+
+// Supported languages
+const (
+	Go   Language = "go"
+	Rust Language = "rust"
+)
+
+type OS string
+
+// Supported OS
+const (
+	Windows OS = "windows"
+	Linux   OS = "linux"
+	MacOS   OS = "darwin" // Darwin is the kernel of macOS
+)
+
+// TODO: I want to make an interface for a compilable language, so that we can add more languages in the future
+// TODO: The cmd might also be an interface or a struct, so that it can build itself based on the OS and language
+// TODO: A cleanup and panic might be needed in setup because if it panics the temp folders should be removed
+// TODO: I am not sure that the setup should panic, maybe it should return an error instead so its easier to clean up
+
+type Compiler struct {
+	OS            OS
+	Language      Language
+	languageEnv   ILanguageEnvironment
+	SourceCode    string
+	Filename      string
+	cmdPrefix     string // For example "cmd /c" on Windows
+	Dependencies  []string
+	tempOutputDir string
+}
+
+type ICompiler interface {
+	Compile() (string, error)
+}
+
+type GoEnvironment struct {
+}
+
+// Setup initializes the go environment by creating a go module and running go mod tidy. Panics if it fails.
+func (ge *GoEnvironment) SetupEnvironment(cmdPrefix string) {
+	// One string
+	cmdString := cmdPrefix + " go mod init " + TempModuleName + " && go mod tidy"
+	// Split the string into a slice
+	cmdSlice := strings.Fields(cmdString) // Fields splits the strings around each instance of one or more consecutive white space characters
+
+	// Make the command
+	cmd := exec.Command(cmdSlice[0], cmdSlice[1:]...)
+	// Set its target directory
+	cmd.Dir = TempOutputDir
+	// Execute the command
+	err := cmd.Run()
+	if err != nil {
+		removeTempFolders(TempOutputDir)
+		log.Fatalf("Error initializing go module: %v", err)
+	}
+}
+
+func (ge *GoEnvironment) CheckCompileErrors(filename string, language Language, cmdPrefix string) (string, error) {
+
+	srcCodeFilename := appendSuffix(filename, language)
+	compiledFilename := filename
+
+	cmdString := cmdPrefix + " go build -o " + compiledFilename + " " + srcCodeFilename
+	cmdSlice := strings.Fields(cmdString) // Fields splits the string on white space of variable length
+
+	cmd := exec.Command(cmdSlice[0], cmdSlice[1:]...)
+	cmd.Dir = TempOutputDir
+	output, err := cmd.CombinedOutput()
+	return string(output), err
+}
+
+func InitCompiler(OS OS, language Language, sourceCode string, filename string, dependencies []string) ICompiler {
+	compiler := &Compiler{}
+	compiler.OS = OS
+	compiler.Language = language
+	compiler.SourceCode = sourceCode
+	compiler.Filename = filename
+	compiler.Dependencies = dependencies
+	compiler.cmdPrefix = getOsPrefix(OS)
+	compiler.languageEnv = getLanguageEnv(language)
+	return compiler
+
+}
+
+func getOsPrefix(OS OS) string {
+	// Set the cmd prefix based on the OS
+	switch OS {
+	case Windows:
+		return "cmd /c "
+	case Linux, MacOS:
+		return ""
+	default:
+		panic("Unsupported OS")
+	}
+}
+
+func getLanguageEnv(language Language) ILanguageEnvironment {
+	switch language {
+	case Go:
+		return &GoEnvironment{}
+	case Rust:
+		return &RustEnvironment{}
+	default:
+		panic("Unsupported language")
+	}
+}
+
+type ILanguageEnvironment interface {
+	SetupEnvironment(cmdPrefix string)
+	CheckCompileErrors(filename string, language Language, cmdPrefix string) (string, error)
+}
+
+type RustEnvironment struct {
+	Compiler
+}
+
+// SetupEnvironment initializes the rust environment by creating a cargo project and adding dependencies. Panics if it fails.
+func (re *RustEnvironment) SetupEnvironment(cmdPrefix string) {
+	// Initialize the rust cargo project--------------------------------------------------------------------------------
+	// Command to initialize a cargo project
+	cmdString := cmdPrefix + " cargo init --bin " + TempModuleName
+	// Split the string into a slice
+	cmdSlice := strings.Fields(cmdString)
+	// Make the command
+	cmd := exec.Command(cmdSlice[0], cmdSlice[1:]...)
+	// Set its target directory
+	cmd.Dir = TempOutputDir
+	// Execute the command
+	err := cmd.Run()
+	if err != nil {
+		log.Fatalf("Error initializing rust project: %v", err)
+	}
+
+	// Write the source code to a file----------------------------------------------------------------------------------
+
+	// Create a new file with the source code in the src folder
+	// 0644 is the file permission where the user can read and write the file, and the group and others can only read the file.
+	err = os.WriteFile(TempOutputDir+TempModuleName+"/src/"+re.Filename, []byte(re.SourceCode), 0644)
+
+	// Update rust dependencies in cargo.toml file using cargo add (cargo-edit)-----------------------------------------
+	addCommand := re.cmdPrefix + " cargo add"
+	addSlice := strings.Split(addCommand, " ")
+	addSlice = append(addSlice, re.Dependencies...)
+	err = exec.Command(addSlice[0], addSlice[1:]...).Run()
+	if err != nil {
+		log.Fatalf("Error adding dependencies: %v", err)
+	}
+}
+
+func (re *RustEnvironment) CheckCompileErrors(filename string, language Language, cmdPrefix string) (string, error) {
+
+	srcCodeFilename := TempOutputDir + appendSuffix(filename, language)
+	output, err := exec.Command(cmdPrefix + " cargo check " + srcCodeFilename).CombinedOutput()
+	return string(output), err
+}
+
+/*
+The compiler pipeline
+1. Set up OS and Language
+2. Set up the temp folders
+3. Write the source code to a file
+4. SetupEnvironment the code
+5. Return the output and error
+*/
+
+func (c *Compiler) Compile() (string, error) {
+	// Set up temp folders
+	setupTempFolders(TempOutputDir)
+	defer removeTempFolders(TempOutputDir)
+
+	srcCodeFilename := TempOutputDir + appendSuffix(c.Filename, c.Language)
+	//compiledFilename := TempOutputDir + c.Filename
+
+	// SetupEnvironment either Go or Rust environment, should be an interface method
+	c.languageEnv.SetupEnvironment(c.cmdPrefix)
+
+	// Write the source code to a file
+	err := os.WriteFile(srcCodeFilename, []byte(c.SourceCode), 0644)
+	if err != nil {
+		log.Fatalf("Error writing source code to file: %v", err)
+	}
+
+	// CheckCompileErrors the code
+	return c.languageEnv.CheckCompileErrors(c.Filename, c.Language, c.cmdPrefix)
+}
+
+// appendSuffix appends the suffix to the filename if it is not already there depending on the language, panics if the language is not supported
+func appendSuffix(filename string, language Language) string {
+
+	suffix := ""
+	switch language {
+	case Go:
+		suffix = ".go"
+	case Rust:
+		suffix = ".rs"
+	default:
+		panic("Unsupported language")
+	}
+
+	// We check if the filename already has the suffix, if not we append it
+	if !strings.HasSuffix(filename, suffix) {
+		filename += suffix
+	}
+	return filename
+}
+
+// setupTempFolders creates the temp output directory for compiled files, panics if it fails
+func setupTempFolders(tempOutputDir string) {
+	// 0777 are the permissions for the directory, everyone can read, write and execute
+	err := os.MkdirAll(tempOutputDir, os.ModePerm)
+	if err != nil {
+		panic("Error creating temp output directory:\n\n" + err.Error())
+	}
+}
+
+// removeTempFolders removes the temp output directory for compiled files, panics if it fails
+func removeTempFolders(tempOutputDir string) {
+	err := os.RemoveAll(tempOutputDir)
+	if err != nil {
+		panic("Error removing temp output directory:\n\n" + err.Error())
+	}
+}
+
+func CompileStringToGo(code string, filename string, dependencies []string) (string, error) {
+
+	// Get the OS
+	os := runtime.GOOS
+
+	// SetupEnvironment
+	return InitCompiler(OS(os), Go, code, filename, dependencies).Compile()
+}
+
+func CompileStringToRust(code string, filename string, dependencies []string) (string, error) {
+
+	// Get the OS
+	os := runtime.GOOS
+
+	// SetupEnvironment
+	return InitCompiler(OS(os), Rust, code, filename, dependencies).Compile()
+}