CyberStrikeAI/internal/database/vulnerability_retest_test.go

package database

import (
	"path/filepath"
	"reflect"
	"strings"
	"testing"

	"go.uber.org/zap"
)

func TestVulnerabilityRetestSiblingSnapshotsAndPlaceholders(t *testing.T) {
	db, err := NewDB(filepath.Join(t.TempDir(), "retest.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	project, err := db.CreateProject(&Project{Name: "retest project"})
	if err != nil {
		t.Fatal(err)
	}
	source, err := db.CreateConversation("source", ConversationCreateMeta{ProjectID: project.ID, RoleName: "tester", AgentMode: "eino_single"})
	if err != nil {
		t.Fatal(err)
	}
	if _, err = db.AddMessage(source.ID, "user", "original request", nil); err != nil {
		t.Fatal(err)
	}
	if _, err = db.AddMessage(source.ID, "assistant", "three vulnerabilities", nil); err != nil {
		t.Fatal(err)
	}
	if err = db.SaveAgentTrace(source.ID, `[{"role":"user","content":"original request"}]`, "trace output"); err != nil {
		t.Fatal(err)
	}
	before, err := db.GetConversation(source.ID)
	if err != nil {
		t.Fatal(err)
	}
	seen := map[string]bool{}
	for i := 0; i < 3; i++ {
		v, err := db.CreateVulnerability(&Vulnerability{ConversationID: source.ID, ProjectID: project.ID, Title: "finding", Severity: "high", Status: "open", RetestNotes: "operator method"})
		if err != nil {
			t.Fatal(err)
		}
		branch, err := db.CreateVulnerabilityRetest(v, ConversationCreateMeta{Source: "test"})
		if err != nil {
			t.Fatal(err)
		}
		if branch.ID == source.ID || seen[branch.ID] || branch.ProjectID != source.ProjectID || branch.RoleName != source.RoleName || branch.AgentMode != source.AgentMode {
			t.Fatalf("invalid sibling: %+v", branch)
		}
		seen[branch.ID] = true
		history, err := db.GetMessages(branch.ID)
		if err != nil {
			t.Fatal(err)
		}
		if len(history) != 2 || history[0].Content != "original request" || history[0].ID == before.Messages[0].ID {
			t.Fatalf("invalid history: %+v", history)
		}
		trace, out, err := db.GetAgentTrace(branch.ID)
		if err != nil || trace == "" || out != "trace output" {
			t.Fatalf("trace: %s %s %v", trace, out, err)
		}
		resolved, err := db.ResolveConversationPlaceholders(branch.ID, "prefix {{vulnerability_id}} middle {{conversation_id}} {{source_conversation_id}} repeat {{vulnerability_id}} {{unknown}}")
		if err != nil {
			t.Fatal(err)
		}
		want := "prefix " + v.ID + " middle " + branch.ID + " " + source.ID + " repeat " + v.ID + " {{unknown}}"
		if resolved != want {
			t.Fatalf("resolved=%q want=%q", resolved, want)
		}
		if _, err = db.AddMessage(branch.ID, "user", "continue verification", nil); err != nil {
			t.Fatal(err)
		}
		if err = db.UpdateVulnerabilityRetestResult(v.ID, "fixed", "response proves denied", branch.ID); err != nil {
			t.Fatal(err)
		}
		got, err := db.GetVulnerability(v.ID)
		if err != nil || got.Status != "fixed" || got.ConversationID != source.ID {
			t.Fatalf("result: %+v %v", got, err)
		}
		if !strings.Contains(got.RetestLog, branch.ID) || !strings.Contains(got.RetestLog, "response proves denied") {
			t.Fatalf("retest evidence missing from retest_log: %q", got.RetestLog)
		}
		if got.RetestNotes != "operator method" {
			t.Fatalf("retest wrote into the operator-owned retest_notes: %q", got.RetestNotes)
		}
	}
	after, err := db.GetConversation(source.ID)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(before, after) {
		t.Fatal("source conversation was changed by retests")
	}
}

// The edit form rewrites every column it owns in one UPDATE, which used to erase evidence kept in retest_notes.
func TestVulnerabilityRetestLogSurvivesOperatorEdit(t *testing.T) {
	db, err := NewDB(filepath.Join(t.TempDir(), "edit.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	v, err := db.CreateVulnerability(&Vulnerability{Title: "finding", Severity: "high", Status: "open", RetestNotes: "operator method"})
	if err != nil {
		t.Fatal(err)
	}
	if err = db.UpdateVulnerabilityRetestResult(v.ID, "fixed", "403 as expected", "branch-1"); err != nil {
		t.Fatal(err)
	}
	edited, err := db.GetVulnerability(v.ID)
	if err != nil {
		t.Fatal(err)
	}
	edited.Title = "finding renamed"
	edited.RetestNotes = "operator method v2"
	if err = db.UpdateVulnerability(v.ID, edited); err != nil {
		t.Fatal(err)
	}
	after, err := db.GetVulnerability(v.ID)
	if err != nil {
		t.Fatal(err)
	}
	if after.Title != "finding renamed" || after.RetestNotes != "operator method v2" {
		t.Fatalf("operator edit lost: %+v", after)
	}
	if !strings.Contains(after.RetestLog, "403 as expected") || !strings.Contains(after.RetestLog, "branch-1") {
		t.Fatalf("operator edit clobbered the retest log: %q", after.RetestLog)
	}
	// The detail card renders from list data, so the log has to ride along there too.
	list, err := db.ListVulnerabilities(10, 0, VulnerabilityListFilter{})
	if err != nil {
		t.Fatal(err)
	}
	if len(list) != 1 || !strings.Contains(list[0].RetestLog, "403 as expected") {
		t.Fatalf("list view is missing the retest log: %+v", list)
	}
}

func TestVulnerabilityRetestMissingSourceRollsBack(t *testing.T) {
	db, err := NewDB(filepath.Join(t.TempDir(), "rollback.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	if _, err = db.CreateVulnerabilityRetest(&Vulnerability{ConversationID: "missing", ID: "finding", Title: "missing"}, ConversationCreateMeta{}); err == nil {
		t.Fatal("missing source accepted")
	}
	var count int
	if err = db.QueryRow(`SELECT COUNT(*) FROM conversations`).Scan(&count); err != nil || count != 0 {
		t.Fatalf("left partial branch: %d %v", count, err)
	}
}

// Copied rows must land in the same created_at format AddMessage writes: GetMessages
// orders by a raw string sort, so a mixed format puts new turns above inherited history.
func TestVulnerabilityRetestBranchKeepsChronologicalOrder(t *testing.T) {
	db, err := NewDB(filepath.Join(t.TempDir(), "order.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	source, err := db.CreateConversation("source", ConversationCreateMeta{AgentMode: "eino_single"})
	if err != nil {
		t.Fatal(err)
	}
	if _, err = db.AddMessage(source.ID, "user", "original request", nil); err != nil {
		t.Fatal(err)
	}
	if _, err = db.AddMessage(source.ID, "assistant", "found it", nil); err != nil {
		t.Fatal(err)
	}
	v, err := db.CreateVulnerability(&Vulnerability{ConversationID: source.ID, Title: "finding", Severity: "high", Status: "open"})
	if err != nil {
		t.Fatal(err)
	}
	branch, err := db.CreateVulnerabilityRetest(v, ConversationCreateMeta{Source: "test"})
	if err != nil {
		t.Fatal(err)
	}
	if _, err = db.AddMessage(branch.ID, "user", "retest turn", nil); err != nil {
		t.Fatal(err)
	}
	timeline, err := db.GetMessages(branch.ID)
	if err != nil {
		t.Fatal(err)
	}
	if len(timeline) != 3 {
		t.Fatalf("want 3 messages, got %d", len(timeline))
	}
	for i, want := range []string{"original request", "found it", "retest turn"} {
		if timeline[i].Content != want {
			t.Fatalf("position %d = %q, want %q", i, timeline[i].Content, want)
		}
	}
	var formats int
	if err = db.QueryRow(`SELECT COUNT(DISTINCT CASE WHEN created_at LIKE '%T%' THEN 1 ELSE 0 END) FROM messages WHERE conversation_id=?`, branch.ID).Scan(&formats); err != nil {
		t.Fatal(err)
	}
	if formats != 1 {
		t.Fatal("branch mixes created_at formats, which inverts the string ORDER BY")
	}
}