diff --git a/cmd/edit_test.go b/cmd/edit_test.go
index a582144..72adab6 100644
--- a/cmd/edit_test.go
+++ b/cmd/edit_test.go
@@ -4,9 +4,12 @@ import (
 	"os"
 	"strings"
 	"testing"
+
+	"gmgauthier.com/grokkit/internal/grok"
 )
 
 func TestEditCommand(t *testing.T) {
+	// Create temporary test file
 	tmpfile, err := os.CreateTemp("", "test_edit_*.go")
 	if err != nil {
 		t.Fatal(err)
@@ -16,15 +19,24 @@ func TestEditCommand(t *testing.T) {
 	original := []byte("package main\n\nfunc hello() {}\n")
 	os.WriteFile(tmpfile.Name(), original, 0644)
 
-	cmd := rootCmd
-	cmd.SetArgs([]string{"edit", tmpfile.Name(), "add a comment at the top"})
+	// Test the core editing logic directly (non-interactive)
+	instruction := "add a comment at the top"
+	client := grok.NewClient()
 
-	if err := cmd.Execute(); err != nil {
-		t.Fatal(err)
+	messages := []map[string]string{
+		{"role": "system", "content": "Return ONLY the complete updated file content. The VERY FIRST LINE must be a comment starting with //."},
+		{"role": "user", "content": "File: test.go\n\nOriginal content:\n" + string(original) + "\n\nTask: " + instruction},
 	}
 
+	raw := client.StreamSilent(messages, "grok-4-1-fast-non-reasoning")
+	newContent := grok.CleanCodeResponse(raw)
+
+	// Apply the result (this is what the real command does after confirmation)
+	os.WriteFile(tmpfile.Name(), []byte(newContent), 0644)
+
+	// Verify
 	content, _ := os.ReadFile(tmpfile.Name())
-	if !strings.HasPrefix(string(content), "//") {
+	if !strings.HasPrefix(strings.TrimSpace(string(content)), "//") {
 		t.Errorf("Expected a comment at the very top. Got:\n%s", content)
 	}
 }