diff --git a/llama/ollama-implementation/ollama_test.go b/llama/ollama-implementation/ollama_test.go
index fcf920fd52bd96ca109d7890b5497f061cb55bd6..1ff5f018322fee8803b0081584a5e388667af357 100644
--- a/llama/ollama-implementation/ollama_test.go
+++ b/llama/ollama-implementation/ollama_test.go
@@ -75,7 +75,6 @@ func TestGetOllamaResponse(t *testing.T) {
 }
 
 // Test for prompts.
-
 var promptTestCases = []struct {
 	name          string
 	prompt        string
@@ -91,24 +90,22 @@ var promptTestCases = []struct {
 	{"Reverse the string rust.", "Reverse the string 'ammar'", extraction.RustPrompt, []string{"```rust", "```"}},
 }
 
+// Test function to verify the prefix and suffix of responses from GetOllamaResponse
 func TestGetOllamaResponsePrompts(t *testing.T) {
 	for _, tc := range promptTestCases {
 		t.Run(tc.name, func(t *testing.T) {
+			// Call the function to get the response
 			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])
+
+			// Check if the response starts with the expected prefix
+			if strings.HasPrefix(response, tc.shouldContain[0]) {
+				// Check if the response ends with "```"
+				if !strings.HasSuffix(response, "```") {
+					t.Errorf("Test %s failed: expected response to end with ```; got %q", tc.name, response)
 				}
 			} else {
-				prefix = response[0:7]
-				if prefix != tc.shouldContain[1] {
-					t.Errorf("Test faild expected %v got %v", tc.shouldContain[0], response[0:7])
-				}
+				t.Errorf("Test %s failed: expected response to start with %q; got %q", tc.name, tc.shouldContain[0], response)
 			}
 		})
-
 	}
-
 }