notarius/internal/modules/dnd/extract/spells/schema_test.go

package spells

import (
	"bytes"
	"encoding/json"
	"strings"
	"testing"

	"github.com/santhosh-tekuri/jsonschema/v6"
)

func TestLoadResponseSchemaUsesExtractorOwnedLLMSchema(t *testing.T) {
	schema, err := loadResponseSchema()
	if err != nil {
		t.Fatalf("loadResponseSchema() error = %v, want nil", err)
	}
	if schema.Key != ResponseSchemaKey || schema.ID != ResponseSchemaID || schema.Version != SchemaVersion || schema.Name != ResponseSchemaName {
		t.Fatalf("schema identity = %#v, want maintained response identity", schema)
	}
	if !strings.HasPrefix(schema.SHA256, "sha256:") || !json.Valid(schema.JSONSchema) {
		t.Fatalf("schema metadata = %#v, want valid hashed JSON", schema)
	}
	var schemaDocument map[string]any
	if err := json.Unmarshal(schema.JSONSchema, &schemaDocument); err != nil {
		t.Fatalf("Unmarshal(schema.JSONSchema) error = %v, want nil", err)
	}
	if schemaDocument["$id"] != ResponseSchemaID {
		t.Fatalf("schema $id = %#v, want %q", schemaDocument["$id"], ResponseSchemaID)
	}

	valid := validSpellsResponse()
	validJSON, err := json.Marshal(valid)
	if err != nil {
		t.Fatalf("Marshal(validSpellsResponse()) error = %v, want nil", err)
	}
	if err := validateJSONSchema(validJSON, schema.JSONSchema); err != nil {
		t.Fatalf("valid private spells response rejected: %v", err)
	}
	for _, test := range []struct {
		name     string
		response map[string]any
		valid    bool
	}{
		{
			name: "semantic blanks and empty evidence",
			response: map[string]any{"spell_casts": []any{map[string]any{
				"caster": "", "spell": "", "source_refs": []any{},
			}}},
			valid: true,
		},
		{
			name: "nonpositive unit candidates",
			response: map[string]any{"spell_casts": []any{map[string]any{
				"caster": "Aria", "spell": "Cure Wounds",
				"source_refs": []any{map[string]any{"start_unit_id": 0, "end_unit_id": -1}},
			}}},
			valid: true,
		},
		{
			name: "missing required field",
			response: map[string]any{"spell_casts": []any{map[string]any{
				"spell": "Cure Wounds", "source_refs": []any{},
			}}},
		},
		{
			name: "unknown field",
			response: map[string]any{"spell_casts": []any{map[string]any{
				"caster": "Aria", "spell": "Cure Wounds", "source_refs": []any{}, "id": "assigned later",
			}}},
		},
		{
			name: "wrong field type",
			response: map[string]any{"spell_casts": []any{map[string]any{
				"caster": 7, "spell": "Cure Wounds", "source_refs": []any{},
			}}},
		},
		{
			name: "noninteger source identifier",
			response: map[string]any{"spell_casts": []any{map[string]any{
				"caster": "Aria", "spell": "Cure Wounds",
				"source_refs": []any{map[string]any{"start_unit_id": 1.5, "end_unit_id": 2}},
			}}},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			content, err := json.Marshal(test.response)
			if err != nil {
				t.Fatal(err)
			}
			err = validateJSONSchema(content, schema.JSONSchema)
			if (err == nil) != test.valid {
				t.Fatalf("validateJSONSchema() error = %v, want valid=%t", err, test.valid)
			}
		})
	}
	if err := validateJSONSchema([]byte(`{"spell_casts":`), schema.JSONSchema); err == nil {
		t.Fatal("validateJSONSchema() error = nil, want malformed JSON rejected")
	}

	withCanonicalSourceID := validSpellsResponse()
	withCanonicalSourceID["spell_casts"].([]any)[0].(map[string]any)["source_refs"].([]any)[0].(map[string]any)["source_id"] = "session-alpha"
	content, err := json.Marshal(withCanonicalSourceID)
	if err != nil {
		t.Fatalf("Marshal(response with source_id) error = %v, want nil", err)
	}
	if err := validateJSONSchema(content, schema.JSONSchema); err == nil {
		t.Fatal("validateJSONSchema() error = nil, want canonical source_id rejected")
	}
}

func TestResponseSchemaJSONIsMutationSafe(t *testing.T) {
	first, err := loadResponseSchema()
	if err != nil {
		t.Fatalf("loadResponseSchema() error = %v", err)
	}
	first.JSONSchema[0] = '['
	second, err := loadResponseSchema()
	if err != nil || !json.Valid(second.JSONSchema) || bytes.Equal(first.JSONSchema, second.JSONSchema) {
		t.Fatalf("second schema = %s, %v; want defensive valid copy", second.JSONSchema, err)
	}
}

func TestResponseSchemaDiagnosticsOmitRawSchema(t *testing.T) {
	schema, err := loadResponseSchema()
	if err != nil {
		t.Fatalf("loadResponseSchema() error = %v", err)
	}
	diagnostics := schema.DiagnosticsMap()
	if diagnostics["key"] != ResponseSchemaKey {
		t.Fatalf("diagnostics = %#v, want response key", diagnostics)
	}
	if _, ok := diagnostics["json_schema"]; ok {
		t.Fatalf("diagnostics include raw schema: %#v", diagnostics)
	}
}

func validSpellsResponse() map[string]any {
	return map[string]any{
		"spell_casts": []any{
			map[string]any{
				"caster": "Aria",
				"spell":  "Cure Wounds",
				"source_refs": []any{
					map[string]any{"start_unit_id": 1, "end_unit_id": 2},
				},
			},
		},
	}
}

func validateJSONSchema(instanceContent, schemaContent []byte) error {
	instance, err := jsonschema.UnmarshalJSON(bytes.NewReader(instanceContent))
	if err != nil {
		return err
	}
	schemaDocument, err := jsonschema.UnmarshalJSON(bytes.NewReader(schemaContent))
	if err != nil {
		return err
	}
	compiler := jsonschema.NewCompiler()
	if err := compiler.AddResource("schema.json", schemaDocument); err != nil {
		return err
	}
	schema, err := compiler.Compile("schema.json")
	if err != nil {
		return err
	}
	return schema.Validate(instance)
}