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