324 lines
10 KiB
Go
324 lines
10 KiB
Go
package cli
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"reflect"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
|
)
|
|
|
|
func TestRecomputeStepRecoversThroughFilesystemCheckpoints(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
invalidateOutput bool
|
|
wantCode int
|
|
}{
|
|
{name: "accepted producer is hydrated", wantCode: 0},
|
|
{name: "invalid producer stops dependents", invalidateOutput: true, wantCode: 1},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
roots := newRecomputeTestRoots(t)
|
|
harness := newRecomputeTestHarness()
|
|
fresh := runRecomputeCommand(roots, harness.options(), false)
|
|
if fresh.code != 0 {
|
|
t.Fatalf("fresh run code=%d stderr=%q", fresh.code, fresh.stderr)
|
|
}
|
|
removeCheckpointLaneStage(t, roots.checkpoints, "extract", "first", "producer")
|
|
removeCheckpointLaneStage(t, roots.checkpoints, "merge", "first", "producer")
|
|
if tt.invalidateOutput {
|
|
path := findCheckpointFile(t, roots.checkpoints, "normalize", "first", "producer", "output.json")
|
|
if err := os.WriteFile(path, []byte("{"), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
harness.resetCalls()
|
|
|
|
resumed := runRecomputeCommand(roots, harness.options(), true)
|
|
if resumed.code != tt.wantCode {
|
|
t.Fatalf("resumed code=%d stdout=%q stderr=%q", resumed.code, resumed.stdout, resumed.stderr)
|
|
}
|
|
events := readLatestCheckpointEvents(t, roots.debug)
|
|
if tt.invalidateOutput {
|
|
if harness.callsFor("test/extract/middle") != 0 || harness.callsFor("test/extract/dependent") != 0 {
|
|
t.Fatalf("dependent calls after invalid producer = %#v", harness.callsSnapshot())
|
|
}
|
|
if !strings.Contains(resumed.stderr, string(pipeline.CheckpointReasonDecodeFailed)) {
|
|
t.Fatalf("stderr=%q, want stable checkpoint reason", resumed.stderr)
|
|
}
|
|
assertNormalizeDecisionSequence(t, events, []checkpointDecisionExpectation{{"first", "producer", pipeline.CheckpointDecisionExecuted, pipeline.CheckpointReasonDecodeFailed}})
|
|
return
|
|
}
|
|
|
|
if got := harness.callsSnapshot(); !reflect.DeepEqual(got, map[string]int{"test/extract/dependent": 1, "test/extract/middle": 1}) {
|
|
t.Fatalf("resumed extractor calls = %#v", got)
|
|
}
|
|
outputPath := filepath.Join(latestChildDir(t, roots.output), "result.json")
|
|
data, err := os.ReadFile(outputPath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if string(data) != "[\"producer\",\"unrelated\",\"middle\",\"dependent\"]\n" {
|
|
t.Fatalf("ordered output = %q", data)
|
|
}
|
|
assertNormalizeDecisionSequence(t, events, []checkpointDecisionExpectation{
|
|
{"first", "producer", pipeline.CheckpointDecisionReused, pipeline.CheckpointReasonAcceptedArtifactReused},
|
|
{"first", "unrelated", pipeline.CheckpointDecisionReused, pipeline.CheckpointReasonReused},
|
|
{"second", "middle", pipeline.CheckpointDecisionForcedRecompute, pipeline.CheckpointReasonRecomputeStep},
|
|
{"third", "dependent", pipeline.CheckpointDecisionForcedRecompute, pipeline.CheckpointReasonRecomputeStep},
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|
|
type checkpointDecisionExpectation struct {
|
|
step, lane string
|
|
action pipeline.CheckpointDecisionCategory
|
|
reason pipeline.CheckpointReasonCode
|
|
}
|
|
|
|
func assertNormalizeDecisionSequence(t *testing.T, events []pipeline.CheckpointEvent, want []checkpointDecisionExpectation) {
|
|
t.Helper()
|
|
var got []checkpointDecisionExpectation
|
|
for _, event := range events {
|
|
if event.Stage == string(pipeline.StageNormalize) {
|
|
got = append(got, checkpointDecisionExpectation{event.StepID, event.LaneID, event.Action, event.ReasonCode})
|
|
}
|
|
}
|
|
if !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("normalize decisions = %#v, want %#v", got, want)
|
|
}
|
|
}
|
|
|
|
type recomputeTestHarness struct {
|
|
base *stateTestHarness
|
|
mu sync.Mutex
|
|
calls map[string]int
|
|
}
|
|
|
|
func newRecomputeTestHarness() *recomputeTestHarness {
|
|
return &recomputeTestHarness{base: newStateTestHarness(), calls: make(map[string]int)}
|
|
}
|
|
|
|
func (h *recomputeTestHarness) options() Options {
|
|
opts := h.base.options()
|
|
for _, key := range []string{"test/extract/producer", "test/extract/unrelated", "test/extract/middle", "test/extract/dependent"} {
|
|
moduleKey := key
|
|
spec := pipeline.ModuleSpec{
|
|
Key: moduleKey, Stage: pipeline.StageExtract, Requires: []string{"chunks"}, Provides: []string{"artifact"}, ArtifactKind: stateTestArtifactKind,
|
|
ReferenceSlots: []contracts.ReferenceSlot{{Name: "upstream", AcceptedMediaTypes: []string{"application/json"}, AcceptedArtifactKinds: []contracts.ArtifactKind{stateTestArtifactKind}}},
|
|
}
|
|
if err := pipeline.RegisterExtractor(opts.Registries.Extractors, spec, func() (contracts.Extractor[stateTestArtifact], error) {
|
|
return recomputeTestExtractor{key: moduleKey, harness: h}, nil
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
if err := opts.Registries.Outputs.RegisterWithSpec(pipeline.ModuleSpec{Key: "test/recompute-output", Stage: pipeline.StageOutput, Requires: []string{"normalized"}, Provides: []string{"output"}}, func() (contracts.OutputEncoder, error) {
|
|
return recomputeTestOutput{}, nil
|
|
}); err != nil {
|
|
panic(err)
|
|
}
|
|
opts.Catalog = catalogFromRegistries(opts.Registries)
|
|
return opts
|
|
}
|
|
|
|
func (h *recomputeTestHarness) record(key string) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
h.calls[key]++
|
|
}
|
|
|
|
func (h *recomputeTestHarness) resetCalls() {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
h.calls = make(map[string]int)
|
|
}
|
|
|
|
func (h *recomputeTestHarness) callsFor(key string) int {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
return h.calls[key]
|
|
}
|
|
|
|
func (h *recomputeTestHarness) callsSnapshot() map[string]int {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
result := make(map[string]int, len(h.calls))
|
|
for key, value := range h.calls {
|
|
result[key] = value
|
|
}
|
|
return result
|
|
}
|
|
|
|
type recomputeTestExtractor struct {
|
|
key string
|
|
harness *recomputeTestHarness
|
|
}
|
|
|
|
func (e recomputeTestExtractor) Key() string { return e.key }
|
|
func (e recomputeTestExtractor) ReferenceSlots() []contracts.ReferenceSlot {
|
|
return []contracts.ReferenceSlot{{Name: "upstream", AcceptedMediaTypes: []string{"application/json"}, AcceptedArtifactKinds: []contracts.ArtifactKind{stateTestArtifactKind}}}
|
|
}
|
|
func (e recomputeTestExtractor) Extract(context.Context, contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[stateTestArtifact], error) {
|
|
e.harness.record(e.key)
|
|
return contracts.TypedExtractionResult[stateTestArtifact]{Value: stateTestArtifact{Value: e.key}}, nil
|
|
}
|
|
|
|
type recomputeTestOutput struct{}
|
|
|
|
func (recomputeTestOutput) Key() string { return "test/recompute-output" }
|
|
func (recomputeTestOutput) Encode(_ context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
|
|
lanes := make([]string, 0, len(req.NormalizeOutputs))
|
|
for _, output := range req.NormalizeOutputs {
|
|
lanes = append(lanes, output.LaneID)
|
|
}
|
|
data, err := json.Marshal(lanes)
|
|
if err != nil {
|
|
return contracts.OutputResult{}, err
|
|
}
|
|
return contracts.OutputResult{Files: []contracts.OutputFile{{Name: "result.json", Bytes: append(data, '\n')}}}, nil
|
|
}
|
|
|
|
func newRecomputeTestRoots(t *testing.T) stateTestRoots {
|
|
t.Helper()
|
|
roots := newStateTestRoots(t)
|
|
config := fmt.Sprintf(`version: 3
|
|
output:
|
|
directory: %q
|
|
cache:
|
|
chunk_plans:
|
|
directory: %q
|
|
mode: bypass
|
|
checkpoints:
|
|
enabled: true
|
|
directory: %q
|
|
debug:
|
|
directory: %q
|
|
pipelines:
|
|
sample:
|
|
input: test/input
|
|
chunk: test/chunk
|
|
steps:
|
|
- id: first
|
|
artifacts:
|
|
producer:
|
|
extract: test/extract/producer
|
|
merge: test/merge
|
|
normalize: test/normalize
|
|
unrelated:
|
|
extract: test/extract/unrelated
|
|
merge: test/merge
|
|
normalize: test/normalize
|
|
- id: second
|
|
references:
|
|
upstream:
|
|
artifact:
|
|
step: first
|
|
lane: producer
|
|
artifacts:
|
|
middle:
|
|
extract: test/extract/middle
|
|
merge: test/merge
|
|
normalize: test/normalize
|
|
- id: third
|
|
references:
|
|
upstream:
|
|
artifact:
|
|
step: second
|
|
lane: middle
|
|
artifacts:
|
|
dependent:
|
|
extract: test/extract/dependent
|
|
merge: test/merge
|
|
normalize: test/normalize
|
|
output: test/recompute-output
|
|
`, roots.output, roots.plans, roots.checkpoints, roots.debug)
|
|
if err := os.WriteFile(roots.config, []byte(config), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return roots
|
|
}
|
|
|
|
func runRecomputeCommand(roots stateTestRoots, opts Options, recompute bool) stateTestResult {
|
|
args := []string{"run", "sample", "--config", roots.config, "--input", roots.input, "--chunk_cache", "bypass"}
|
|
if recompute {
|
|
args = append(args, "--resume", "--recompute-step", "second", "--debug")
|
|
}
|
|
var stdout, stderr bytes.Buffer
|
|
return stateTestResult{code: RunWithOptions(args, &stdout, &stderr, opts), stdout: stdout.String(), stderr: stderr.String()}
|
|
}
|
|
|
|
func removeCheckpointLaneStage(t *testing.T, root, stage, step, lane string) {
|
|
t.Helper()
|
|
dir := filepath.Dir(findCheckpointFile(t, root, stage, step, lane, "manifest.json"))
|
|
if err := os.RemoveAll(dir); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func findCheckpointFile(t *testing.T, root, stage, step, lane, name string) string {
|
|
t.Helper()
|
|
want := filepath.Join(stage, step, lane, name)
|
|
var matches []string
|
|
err := filepath.WalkDir(root, func(path string, entry os.DirEntry, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !entry.IsDir() && strings.HasSuffix(path, want) {
|
|
matches = append(matches, path)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(matches) != 1 {
|
|
t.Fatalf("checkpoint files ending in %q = %v", want, matches)
|
|
}
|
|
return matches[0]
|
|
}
|
|
|
|
func latestChildDir(t *testing.T, root string) string {
|
|
t.Helper()
|
|
entries, err := os.ReadDir(root)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var dirs []string
|
|
for _, entry := range entries {
|
|
if entry.IsDir() {
|
|
dirs = append(dirs, filepath.Join(root, entry.Name()))
|
|
}
|
|
}
|
|
if len(dirs) == 0 {
|
|
t.Fatal("no child directory")
|
|
}
|
|
sort.Strings(dirs)
|
|
return dirs[len(dirs)-1]
|
|
}
|
|
|
|
func readLatestCheckpointEvents(t *testing.T, root string) []pipeline.CheckpointEvent {
|
|
t.Helper()
|
|
var events []pipeline.CheckpointEvent
|
|
data, err := os.ReadFile(filepath.Join(latestChildDir(t, root), "summary", "checkpoint-events.json"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := json.Unmarshal(data, &events); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return events
|
|
}
|