add more tests for extract

llama/extraction/extract_test.go

@@ -43,19 +43,27 @@ import (
 // }
 
 // Inputs and Expected Outputs for the Test Cases
-// This can be considered a table-driven test.
+// This can be considered a table-driven test or equivalence partitioning
 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"},
+	{"Go Extraction 1 - Main", "```go\nfunc main() {}\n```", "\nfunc main() {}\n"},
+	{"Go Extraction 2 - Print", "```go\nfmt.Println('Hello World')\n```", "\nfmt.Println('Hello World')\n"},
+	{"Go Extraction 3 - Loop", "```go\nfor i := 0; i < 10; i++ {\nfmt.Println(i)\n}\n```", "\nfor i := 0; i < 10; i++ {\nfmt.Println(i)\n}\n"},
+	{"Go Extraction 4 - If Else", "```go\nif x > 10 {\nfmt.Println('Greater than 10')\n} else {\nfmt.Println('Less than or equal to 10')\n}\n```", "\nif x > 10 {\nfmt.Println('Greater than 10')\n} else {\nfmt.Println('Less than or equal to 10')\n}\n"},
+	{"Go Extraction 5 - Function with Parameters", "```go\nfunc add(a int, b int) int {\nreturn a + b\n}\n```", "\nfunc add(a int, b int) int {\nreturn a + b\n}\n"},
+	{"Go Extraction 6 - Nested Loops", "```go\nfor i := 0; i < 3; i++ {\nfor j := 0; j < 3; j++ {\nfmt.Printf('(%d, %d)', i, j)\n}\n}\n```", "\nfor i := 0; i < 3; i++ {\nfor j := 0; j < 3; j++ {\nfmt.Printf('(%d, %d)', i, j)\n}\n}\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"},
+	{"Rust Extraction 1 - Main", "```rust\nfn main() {}\n```", "\nfn main() {}\n"},
+	{"Rust Extraction 2 - Print", "```rust\nprintln!('Hello World')\n```", "\nprintln!('Hello World')\n"},
+	{"Rust Extraction 3 - Loop", "```rust\nfor i in 0..10 {\nprintf!(\"{}\", i);\n}\n```", "\nfor i in 0..10 {\nprintf!(\"{}\", i);\n}\n"},
+	{"Rust Extraction 4 - If Else", "```rust\nif x > 10 {\nprintln!(\"Greater than 10\");\n} else {\nprintln!(\"Less than or equal to 10\");\n}\n```", "\nif x > 10 {\nprintln!(\"Greater than 10\");\n} else {\nprintln!(\"Less than or equal to 10\");\n}\n"},
+	{"Rust Extraction 5 - Function with Parameters", "```rust\nfn add(a: i32, b: i32) -> i32 {\nreturn a + b;\n}\n```", "\nfn add(a: i32, b: i32) -> i32 {\nreturn a + b;\n}\n"},
+	{"Rust Extraction 6 - Nested Loops", "```rust\nfor i in 0..3 {\nfor j in 0..3 {\nprintf!(\"({},{})\", i, j);\n}\n}\n```", "\nfor i in 0..3 {\nfor j in 0..3 {\nprintf!(\"({},{})\", i, j);\n}\n}\n"},
 }
 
 // Refined Test Function using Table-Driven Approach

llama/ollama-implementation/ollama_test.go

@@ -55,7 +55,6 @@ func TestGetOllamaResponse(t *testing.T) {
 	context := []int{}
 	response, updatedContext, err := GetOllamaResponse(prompt, context)
 
-	// Verify the response and error handling
 	if err != nil {
 		t.Fatalf("Expected no error, got %v", err)
 	}
@@ -75,6 +74,8 @@ func TestGetOllamaResponse(t *testing.T) {
 	}
 }
 
+// Test for prompts.
+
 var promptTestCases = []struct {
 	name          string
 	prompt        string