git.giacomolli.uk - onechat/blob - summarizer_helper_test.go

   1 package main
   2
   3 import (
   4         "context"
   5         "errors"
   6         "fmt"
   7         "slices"
   8         "strings"
   9         "sync"
  10         "unicode/utf8"
  11 )
  12
  13 // FakeSummarizer builds each node with a line of fakeLine. It fails each task
  14 // while failing is true, and each task waits for gate when gate is not nil.
  15 type FakeSummarizer struct {
  16         mutex   sync.Mutex
  17         gate    chan struct{}
  18         failing bool
  19         active  int
  20         peak    int
  21         tasks   []Task
  22 }
  23
  24 func (f *FakeSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
  25         f.mutex.Lock()
  26         f.tasks = append(f.tasks, task)
  27         f.active++
  28         f.peak = max(f.peak, f.active)
  29         fail := f.failing
  30         f.mutex.Unlock()
  31
  32         defer func() {
  33                 f.mutex.Lock()
  34                 f.active--
  35                 f.mutex.Unlock()
  36         }()
  37
  38         if f.gate != nil {
  39                 select {
  40                 case <-f.gate:
  41                 case <-ctx.Done():
  42                         return "", ctx.Err()
  43                 }
  44         }
  45
  46         if fail {
  47                 return "", errors.New("fake failure")
  48         }
  49
  50         return fakeLine(task), nil
  51 }
  52
  53 // SetFailing makes the next tasks fail or not.
  54 func (f *FakeSummarizer) SetFailing(failing bool) {
  55         f.mutex.Lock()
  56         defer f.mutex.Unlock()
  57
  58         f.failing = failing
  59 }
  60
  61 // Active returns the number of tasks that run now.
  62 func (f *FakeSummarizer) Active() int {
  63         f.mutex.Lock()
  64         defer f.mutex.Unlock()
  65
  66         return f.active
  67 }
  68
  69 // Tasks returns the tasks so far.
  70 func (f *FakeSummarizer) Tasks() []Task {
  71         f.mutex.Lock()
  72         defer f.mutex.Unlock()
  73
  74         return slices.Clone(f.tasks)
  75 }
  76
  77 // fakeLine returns a line of nodeBytes bytes or less that names the node and
  78 // starts with the input. A real line of this size makes each merge a model
  79 // call.
  80 func fakeLine(task Task) string {
  81         count := 1 << task.Level
  82         line := fmt.Sprintf("summary %d+%d: %s", task.Index*count, count, strings.Join(task.Input, " "))
  83         line = strings.ReplaceAll(line, "\n", " ")
  84         for len(line) > nodeBytes-12 {
  85                 _, size := utf8.DecodeLastRuneInString(line)
  86                 line = line[:len(line)-size]
  87         }
  88
  89         return line
  90 }
  91
  92 func newFakeSummarizer() *FakeSummarizer {
  93         return &FakeSummarizer{}
  94 }