Skip to content
Snippets Groups Projects
Commit bafb34c3 authored by Ammar Ahmed's avatar Ammar Ahmed :speech_balloon:
Browse files

fixed return values of extract and handled error in test cases

parent 043900de
No related branches found
No related tags found
No related merge requests found
package extraction package extraction
import ( import (
"fmt"
"strings" "strings"
) )
...@@ -20,17 +21,17 @@ var RustPrompt = "The code should be in the Rust programming language. There sho ...@@ -20,17 +21,17 @@ var RustPrompt = "The code should be in the Rust programming language. There sho
// } // }
// Extract extracts the code snippet between ``` and removes the language identifier. // Extract extracts the code snippet between ``` and removes the language identifier.
func Extract(output string) string { func Extract(output string) (string, error) {
parts := strings.Split(output, "```") parts := strings.Split(output, "```")
if len(parts) < 2 { if len(parts) < 2 {
return "" // Handle the case if format is incorrect: Return empty string return "", fmt.Errorf("the string wasn't in a proper format") // Handle the case if format is incorrect: Return empty string
} }
// Trim the language identifier like `go` or `rust` from the code // Trim the language identifier like `go` or `rust` from the code
code := parts[1] code := parts[1]
lines := strings.SplitN(code, "\n", 2) lines := strings.SplitN(code, "\n", 2)
if len(lines) > 1 { if len(lines) > 1 {
return "\n" + lines[1] // Return the code without the first line (language identifier) return "\n" + lines[1], nil // Return the code without the first line (language identifier)
} }
return "" return "", fmt.Errorf("the string doesn't contain any lines")
} }
...@@ -71,9 +71,11 @@ var testCases = []struct { ...@@ -71,9 +71,11 @@ var testCases = []struct {
func TestExtraction(t *testing.T) { func TestExtraction(t *testing.T) {
for _, tc := range testCases { for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) { t.Run(tc.name, func(t *testing.T) {
output := Extract(tc.input) output, err := Extract(tc.input)
if output != tc.expected { if output != tc.expected {
t.Errorf("Test %s failed: Expected %q, got %q", tc.name, tc.expected, output) t.Errorf("Test %s failed: Expected %q, got %q", tc.name, tc.expected, output)
t.Log(err.Error())
} }
}) })
} }
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment