snowpick/classify_test.go

package main

import (
	"testing"
)

func TestClassifyTablePIIInferred(t *testing.T) {
	hints := ClassifyTable("sys_user", "name", nil)
	if len(hints) == 0 {
		t.Fatal("expected hint for sys_user")
	}
	if hints[0].Category != ImpactPII {
		t.Fatalf("category = %q", hints[0].Category)
	}
	if hints[0].Basis != "inferred" {
		t.Fatalf("basis = %q, expected inferred", hints[0].Basis)
	}
	if hints[0].Reason != reasonTablePII {
		t.Fatalf("reason = %q", hints[0].Reason)
	}
}

func TestClassifyTableOperationsInferred(t *testing.T) {
	hints := ClassifyTable("incident", "number", nil)
	if len(hints) == 0 || hints[0].Category != ImpactOperations {
		t.Fatalf("expected operations hint for incident, got %#v", hints)
	}
	if hints[0].Basis != "inferred" {
		t.Fatalf("basis = %q", hints[0].Basis)
	}
}

func TestClassifyTableCredentialsInferred(t *testing.T) {
	hints := ClassifyTable("oauth_entity", "name", nil)
	if len(hints) == 0 || hints[0].Category != ImpactCredentials {
		t.Fatalf("expected credentials hint for oauth_entity, got %#v", hints)
	}
	if hints[0].Basis != "inferred" {
		t.Fatalf("basis = %q", hints[0].Basis)
	}
}

func TestClassifyTableInfrastructureInferred(t *testing.T) {
	hints := ClassifyTable("cmdb_model", "name", nil)
	if len(hints) == 0 || hints[0].Category != ImpactInfrastructure {
		t.Fatalf("expected infrastructure hint for cmdb_model, got %#v", hints)
	}
}

func TestClassifyFieldNameInferred(t *testing.T) {
	hints := ClassifyTable("some_table", "email_address", nil)
	if len(hints) == 0 || hints[0].Category != ImpactPII {
		t.Fatalf("expected PII hint for email_address field, got %#v", hints)
	}
	if hints[0].Basis != "inferred" {
		t.Fatalf("basis = %q", hints[0].Basis)
	}
}

func TestClassifyTableByContentEmailObserved(t *testing.T) {
	samples := []TableSample{{DisplayValue: "[email protected]"}}
	hints := ClassifyTable("unknown_table", "unknown_field", samples)
	if len(hints) == 0 || hints[0].Category != ImpactPII {
		t.Fatalf("expected PII hint for email content, got %#v", hints)
	}
	if hints[0].Basis != "observed" {
		t.Fatalf("basis = %q, expected observed", hints[0].Basis)
	}
	if hints[0].Reason != reasonContentEmail {
		t.Fatalf("reason = %q", hints[0].Reason)
	}
}

func TestBareIPDoesNotMatchObserved(t *testing.T) {
	samples := []TableSample{{DisplayValue: "Server 10.0.1.55 rack 3"}}
	hints := ClassifyTable("unknown_table", "unknown_field", samples)
	for _, h := range hints {
		if h.Category == ImpactInfrastructure && h.Basis == "observed" {
			t.Fatalf("bare IP should not produce observed infrastructure hint, got %#v", h)
		}
	}
}

func TestVersionStringDoesNotMatchObserved(t *testing.T) {
	samples := []TableSample{{DisplayValue: "Chrome 120.0.6099.71"}}
	hints := ClassifyTable("unknown_table", "unknown_field", samples)
	for _, h := range hints {
		if h.Category == ImpactInfrastructure && h.Basis == "observed" {
			t.Fatalf("version string should not produce observed infrastructure hint, got %#v", h)
		}
	}
}

func TestCorpDomainObserved(t *testing.T) {
	samples := []TableSample{{DisplayValue: "host app01.corp"}}
	hints := ClassifyTable("unknown_table", "unknown_field", samples)
	if len(hints) == 0 || hints[0].Category != ImpactInfrastructure || hints[0].Basis != "observed" {
		t.Fatalf("expected observed infrastructure hint for .corp domain, got %#v", hints)
	}
}

func TestInternalDomainObserved(t *testing.T) {
	samples := []TableSample{{DisplayValue: "vpn.example.internal"}}
	hints := ClassifyTable("unknown_table", "unknown_field", samples)
	if len(hints) == 0 || hints[0].Category != ImpactInfrastructure || hints[0].Basis != "observed" {
		t.Fatalf("expected observed infrastructure hint for .internal domain, got %#v", hints)
	}
}

func TestLocalDomainObserved(t *testing.T) {
	samples := []TableSample{{DisplayValue: "asset.local"}}
	hints := ClassifyTable("unknown_table", "unknown_field", samples)
	if len(hints) == 0 || hints[0].Category != ImpactInfrastructure || hints[0].Basis != "observed" {
		t.Fatalf("expected observed infrastructure hint for .local domain, got %#v", hints)
	}
}

func TestInfrastructureTableStillInferred(t *testing.T) {
	hints := ClassifyTable("cmdb_model", "name", nil)
	if len(hints) == 0 || hints[0].Category != ImpactInfrastructure || hints[0].Basis != "inferred" {
		t.Fatalf("expected inferred infrastructure hint for cmdb_model, got %#v", hints)
	}
}

func hasObserved(hints []Hint, cat ImpactCategory) bool {
	for _, h := range hints {
		if h.Category == cat && h.Basis == "observed" {
			return true
		}
	}
	return false
}

func TestCredentialKeywordWithoutValueNotObserved(t *testing.T) {
	// Prose mentions of credential words must not produce observed credentials.
	for _, s := range []string{"the bearer of bad news", "checkbearertest", "Password Reset Procedure", "reset your api_key today"} {
		hints := ClassifyWidget("widget", []RecordSample{{Snippet: s}})
		if hasObserved(hints, ImpactCredentials) {
			t.Fatalf("bare credential keyword %q should not be observed credentials, got %#v", s, hints)
		}
	}
}

func TestCredentialAssignmentAndTokenObserved(t *testing.T) {
	for _, s := range []string{"api_key = AKIAIOSFODNN7", "token: abc123def456", "Authorization: Bearer eyJhbGciOiJIUzI1NiJ9"} {
		hints := ClassifyWidget("widget", []RecordSample{{Snippet: s}})
		if !hasObserved(hints, ImpactCredentials) {
			t.Fatalf("credential value %q should be observed credentials, got %#v", s, hints)
		}
	}
}

func TestInfraSuffixInFilenameNotObserved(t *testing.T) {
	// A TLD-like suffix mid-token (filename) must not produce observed infrastructure.
	for _, s := range []string{"see notes.local.txt", "report.corp.pdf attached", "data.internal.csv export"} {
		hints := ClassifyTable("unknown_table", "unknown_field", []TableSample{{DisplayValue: s}})
		if hasObserved(hints, ImpactInfrastructure) {
			t.Fatalf("filename %q should not be observed infrastructure, got %#v", s, hints)
		}
	}
}

func TestSysPropertiesAndCmdbCiClassified(t *testing.T) {
	if h := ClassifyTable("sys_properties", "", nil); len(h) == 0 || h[0].Category != ImpactCredentials {
		t.Fatalf("expected credentials hint for sys_properties, got %#v", h)
	}
	if h := ClassifyTable("cmdb_ci", "", nil); len(h) == 0 || h[0].Category != ImpactInfrastructure {
		t.Fatalf("expected infrastructure hint for cmdb_ci, got %#v", h)
	}
}

func TestMixedBasisResolvesToObserved(t *testing.T) {
	// sys_user (inferred PII) + email in content (observed PII) -> observed wins
	samples := []TableSample{{DisplayValue: "[email protected]"}}
	hints := ClassifyTable("sys_user", "name", samples)

	var piiHint *Hint
	for i := range hints {
		if hints[i].Category == ImpactPII {
			piiHint = &hints[i]
			break
		}
	}
	if piiHint == nil {
		t.Fatal("expected PII hint")
	}
	if piiHint.Basis != "observed" {
		t.Fatalf("basis = %q, expected observed (observed should win over inferred)", piiHint.Basis)
	}
	// Should have both table and content signals
	hasTable := false
	hasContent := false
	for _, s := range piiHint.Signals {
		if s == "table:sys_user" {
			hasTable = true
		}
		if s == "content:[email protected]" {
			hasContent = true
		}
	}
	if !hasTable || !hasContent {
		t.Fatalf("expected both table and content signals, got %v", piiHint.Signals)
	}
}

func TestClassifyWidgetKBInferred(t *testing.T) {
	hints := ClassifyWidget("kb-article-page", nil)
	if len(hints) == 0 || hints[0].Category != ImpactOperations {
		t.Fatalf("expected operations hint for kb widget, got %#v", hints)
	}
	if hints[0].Basis != "inferred" {
		t.Fatalf("basis = %q", hints[0].Basis)
	}
}

func TestClassifyWidgetWithSamplesObserved(t *testing.T) {
	samples := []RecordSample{
		{Title: "Password Reset", Snippet: "Contact [email protected] for details"},
	}
	hints := ClassifyWidget("some-widget", samples)
	found := map[ImpactCategory]string{}
	for _, h := range hints {
		found[h.Category] = h.Basis
	}
	if found[ImpactPII] != "observed" {
		t.Fatal("expected observed PII hint for email in snippet")
	}
	// "Password Reset" should NOT trigger credentials:observed (bare keyword)
	if found[ImpactCredentials] == "observed" {
		t.Fatal("bare 'password' in title should not produce observed credentials hint")
	}
}

func TestClassifyCredentialAssignmentObserved(t *testing.T) {
	samples := []RecordSample{
		{Snippet: "config: api_key = AKIAIOSFODNN7EXAMPLE"},
	}
	hints := ClassifyWidget("some-widget", samples)
	found := map[ImpactCategory]string{}
	for _, h := range hints {
		found[h.Category] = h.Basis
	}
	if found[ImpactCredentials] != "observed" {
		t.Fatal("expected observed credentials hint for api_key assignment pattern")
	}
}

func TestClassifyBearerTokenObserved(t *testing.T) {
	samples := []RecordSample{
		{Snippet: "Authorization: Bearer eyJhbGciOiJIUzI1NiJ9"},
	}
	hints := ClassifyWidget("some-widget", samples)
	found := map[ImpactCategory]string{}
	for _, h := range hints {
		found[h.Category] = h.Basis
	}
	if found[ImpactCredentials] != "observed" {
		t.Fatal("expected observed credentials hint for bearer token")
	}
}

func TestClassifyUnknownTableNoSamples(t *testing.T) {
	hints := ClassifyTable("unknown_table", "unknown_field", nil)
	if len(hints) != 0 {
		t.Fatalf("expected no hints for unknown table, got %#v", hints)
	}
}

func TestReviewRequiredHint(t *testing.T) {
	hints := ReviewRequiredHint(500)
	if len(hints) != 1 {
		t.Fatalf("expected 1 hint, got %d", len(hints))
	}
	if hints[0].Category != ImpactReviewRequired {
		t.Fatalf("category = %q", hints[0].Category)
	}
	if hints[0].Basis != "none" {
		t.Fatalf("basis = %q", hints[0].Basis)
	}
	if hints[0].Reason != reasonReviewRequired {
		t.Fatalf("reason = %q", hints[0].Reason)
	}
	if len(hints[0].Signals) != 1 || hints[0].Signals[0] != "records_returned:500" {
		t.Fatalf("signals = %v", hints[0].Signals)
	}
}

func TestHintTags(t *testing.T) {
	hints := []Hint{
		{Category: ImpactPII, Basis: "observed", Reason: reasonContentEmail, Signals: []string{"content:[email protected]"}},
		{Category: ImpactCredentials, Basis: "inferred", Reason: reasonFieldCredentials, Signals: []string{"field:password"}},
	}
	tags := hintTags(hints)
	// Colors are not active in tests, so tags are plain text.
	if tags != " [pii:observed] [credentials:inferred]" {
		t.Fatalf("hintTags() = %q", tags)
	}
}

func TestHintTagsReviewRequired(t *testing.T) {
	hints := ReviewRequiredHint(100)
	tags := hintTags(hints)
	if tags != " [review_required]" {
		t.Fatalf("hintTags() = %q", tags)
	}
}

func TestHintTagsEmpty(t *testing.T) {
	tags := hintTags(nil)
	if tags != "" {
		t.Fatalf("expected empty string, got %q", tags)
	}
}