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