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

   1 package main
   2
   3 import (
   4         "archive/tar"
   5         "context"
   6         "database/sql"
   7         "io"
   8         "os"
   9         "path/filepath"
  10         "strings"
  11         "sync"
  12         "testing"
  13         "time"
  14
  15         "github.com/stretchr/testify/require"
  16 )
  17
  18 func TestRunningDaemonBackupRestoresTheChatAndLinks(t *testing.T) {
  19         dir := t.TempDir()
  20         whatsappTestSession(t, dir, true)
  21
  22         files := map[string]string{
  23                 "auth.json":                  `{"access_token":"fake-secret"}`,
  24                 "telegram.json":              `{"token":"fake-token","owner":123}`,
  25                 "config.json":                `{"endpoint":"api"}`,
  26                 "agent/instructions.md":      "Keep the instructions.\n",
  27                 "agent/telegram-offset.json": "23\n",
  28         }
  29         for name, data := range files {
  30                 path := filepath.Join(dir, name)
  31                 require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
  32                 require.NoError(t, os.WriteFile(path, []byte(data), 0o600))
  33         }
  34
  35         memory, whatsapp, subagents := backupStartDaemon(t, dir)
  36         require.True(t, whatsapp.isLinked())
  37
  38         _, err := memory.Append("note", "A durable decision.")
  39         require.NoError(t, err)
  40
  41         _, err = memory.Append("note", "The next decision.")
  42         require.NoError(t, err)
  43         require.NoError(t, memory.Settle(t.Context()))
  44
  45         child, err := subagents.Create(t.Context())
  46         require.NoError(t, err)
  47
  48         _, err = child.Append("note", "subagent")
  49         require.NoError(t, err)
  50
  51         _, err = child.Append("user", "A saved task.")
  52         require.NoError(t, err)
  53
  54         childText, err := child.Show()
  55         require.NoError(t, err)
  56
  57         for _, excluded := range []string{"probe/skip", "agent-backup-example/skip"} {
  58                 path := filepath.Join(dir, excluded)
  59                 require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
  60                 require.NoError(t, os.WriteFile(path, []byte("excluded"), 0o600))
  61         }
  62
  63         archive := filepath.Join(t.TempDir(), "daily.tar.gz")
  64         t.Setenv("ONECHAT_DIR", dir)
  65
  66         var answer strings.Builder
  67
  68         require.NoError(t, backup(t.Context(), archive, &answer))
  69         require.Regexp(t, `^`+archive+` \([0-9]+ bytes\)\n$`, answer.String())
  70
  71         info, err := os.Stat(archive)
  72         require.NoError(t, err)
  73         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
  74
  75         restored := t.TempDir()
  76         require.NoError(t, backupRestore(t.Context(), archive, restored))
  77
  78         for name, data := range files {
  79                 bytes, err := os.ReadFile(filepath.Join(restored, name))
  80                 require.NoError(t, err)
  81                 require.Equal(t, data, string(bytes))
  82         }
  83
  84         for _, excluded := range []string{"probe", "agent-backup-example", socketFileName, "agent/lock"} {
  85                 _, err = os.Stat(filepath.Join(restored, excluded))
  86                 require.ErrorIs(t, err, os.ErrNotExist)
  87         }
  88
  89         restoredMemory, restoredWhatsApp, restoredSubagents := backupStartDaemon(t, restored)
  90         require.Equal(t, memory.Len(), restoredMemory.Len())
  91         require.Equal(t, memory.View(), restoredMemory.View())
  92
  93         for id := range memory.Len() {
  94                 expected, err := memory.Message(id)
  95                 require.NoError(t, err)
  96
  97                 actual, err := restoredMemory.Message(id)
  98                 require.NoError(t, err)
  99                 require.Equal(t, expected, actual)
 100         }
 101
 102         require.True(t, restoredWhatsApp.isLinked())
 103         require.Equal(t, whatsapp.link().own, restoredWhatsApp.link().own)
 104
 105         restoredChild, err := restoredSubagents.Get(t.Context(), child.name)
 106         require.NoError(t, err)
 107
 108         restoredText, err := restoredChild.Show()
 109         require.NoError(t, err)
 110         require.Equal(t, childText, restoredText)
 111 }
 112
 113 func TestRestoreRefusesNonemptyDataWithoutChangingIt(t *testing.T) {
 114         dir := t.TempDir()
 115         path := filepath.Join(dir, "keep")
 116         require.NoError(t, os.WriteFile(path, []byte("keep"), 0o600))
 117
 118         err := backupRestore(t.Context(), "missing.tar.gz", dir)
 119         require.ErrorContains(t, err, "empty data directory")
 120
 121         entries, err := os.ReadDir(dir)
 122         require.NoError(t, err)
 123         require.Len(t, entries, 1)
 124
 125         data, err := os.ReadFile(path)
 126         require.NoError(t, err)
 127         require.Equal(t, "keep", string(data))
 128 }
 129
 130 func TestRestoreRefusesTheDaemonTokenLock(t *testing.T) {
 131         dir := t.TempDir()
 132         tokens, err := lockTokens(dir)
 133         require.NoError(t, err)
 134
 135         defer func() {
 136                 require.NoError(t, tokens.Close())
 137         }()
 138
 139         err = backupRestore(t.Context(), "missing.tar.gz", dir)
 140         require.ErrorIs(t, err, errTokensInUse)
 141 }
 142
 143 func TestUnsafeRestoreInstallsNoArchiveFiles(t *testing.T) {
 144         archive := backupTestArchive(t,
 145                 tar.Header{
 146                         Name:     "safe",
 147                         Typeflag: tar.TypeReg,
 148                         Mode:     0o600,
 149                 },
 150                 tar.Header{
 151                         Name:     "../escape",
 152                         Typeflag: tar.TypeReg,
 153                         Mode:     0o600,
 154                 },
 155         )
 156         dir := t.TempDir()
 157         err := backupRestore(t.Context(), archive, dir)
 158         require.ErrorContains(t, err, "unsafe archive path")
 159
 160         entries, err := os.ReadDir(dir)
 161         require.NoError(t, err)
 162         require.Len(t, entries, 1)
 163         require.Equal(t, tokenLockFileName, entries[0].Name())
 164 }
 165
 166 func TestBackupNeedsARunningDaemon(t *testing.T) {
 167         t.Setenv("ONECHAT_DIR", t.TempDir())
 168
 169         err := backup(t.Context(), filepath.Join(t.TempDir(), "backup.tar.gz"), io.Discard)
 170         require.EqualError(t, err, "backup needs a running daemon")
 171 }
 172
 173 func TestWhatsAppSessionBackupIncludesCommittedWALWrites(t *testing.T) {
 174         dir := t.TempDir()
 175         whatsappTestSession(t, dir, true)
 176
 177         whatsapp, _, _ := whatsappStartTest(t, dir, newTestClock())
 178         address, err := whatsappSessionAddress(filepath.Join(dir, whatsappDirName, whatsappSessionFileName))
 179         require.NoError(t, err)
 180
 181         database, err := sql.Open("sqlite", address)
 182         require.NoError(t, err)
 183
 184         defer func() {
 185                 require.NoError(t, database.Close())
 186         }()
 187
 188         _, err = database.ExecContext(t.Context(), "PRAGMA journal_mode=WAL; CREATE TABLE backup_sample (value TEXT); INSERT INTO backup_sample VALUES ('saved');")
 189         require.NoError(t, err)
 190
 191         target := filepath.Join(t.TempDir(), "copy", whatsappSessionFileName)
 192         require.NoError(t, whatsapp.BackupSession(t.Context(), target))
 193
 194         address, err = whatsappSessionAddress(target)
 195         require.NoError(t, err)
 196
 197         restored, err := sql.Open("sqlite", address)
 198         require.NoError(t, err)
 199
 200         defer func() {
 201                 require.NoError(t, restored.Close())
 202         }()
 203
 204         var value string
 205
 206         err = restored.QueryRowContext(t.Context(), "SELECT value FROM backup_sample").Scan(&value)
 207         require.NoError(t, err)
 208         require.Equal(t, "saved", value)
 209
 210         info, err := os.Stat(target)
 211         require.NoError(t, err)
 212         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 213 }
 214
 215 func TestSessionBackupWaitsForRemovalToFinish(t *testing.T) {
 216         dir := t.TempDir()
 217         whatsappTestSession(t, dir, true)
 218
 219         whatsapp, fake, _ := whatsappStartTest(t, dir, newTestClock())
 220         disconnect := &BackupDisconnect{
 221                 WhatsAppClient: fake,
 222                 entered:        make(chan struct{}),
 223                 release:        make(chan struct{}),
 224         }
 225
 226         var release sync.Once
 227
 228         defer release.Do(func() {
 229                 close(disconnect.release)
 230         })
 231
 232         whatsapp.mutex.Lock()
 233
 234         whatsapp.client = disconnect
 235         whatsapp.mutex.Unlock()
 236
 237         removed := make(chan error, 1)
 238         go func() {
 239                 removed <- whatsapp.removeSession()
 240         }()
 241
 242         <-disconnect.entered
 243
 244         target := filepath.Join(t.TempDir(), "copy", whatsappSessionFileName)
 245         backedUp := make(chan error, 1)
 246         started := make(chan struct{})
 247         go func() {
 248                 close(started)
 249
 250                 backedUp <- whatsapp.BackupSession(t.Context(), target)
 251         }()
 252
 253         <-started
 254         require.Never(t, func() bool {
 255                 return len(backedUp) != 0
 256         }, 50*time.Millisecond, time.Millisecond)
 257         release.Do(func() {
 258                 close(disconnect.release)
 259         })
 260         require.NoError(t, <-removed)
 261         require.NoError(t, <-backedUp)
 262
 263         for _, path := range []string{whatsapp.sessionPath(), target} {
 264                 _, err := os.Stat(path)
 265                 require.ErrorIs(t, err, os.ErrNotExist)
 266         }
 267 }
 268
 269 // BackupDisconnect holds session removal in its unlocked client cleanup.
 270 type BackupDisconnect struct {
 271         WhatsAppClient
 272         entered chan struct{}
 273         release chan struct{}
 274 }
 275
 276 func (d *BackupDisconnect) Disconnect() {
 277         close(d.entered)
 278         <-d.release
 279 }
 280
 281 // backupStartDaemon serves fake model and WhatsApp clients on the socket of
 282 // dir. It keeps all state files under dir and holds the daemon token lock.
 283 func backupStartDaemon(t *testing.T, dir string) (*Memory, *WhatsApp, *Subagents) {
 284         t.Helper()
 285
 286         tokens, err := lockTokens(dir)
 287         require.NoError(t, err)
 288
 289         clock := newTestClock()
 290         agent, err := openAgent(filepath.Join(dir, daemonAgentDirName))
 291         require.NoError(t, err)
 292
 293         memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer(), discardLog())
 294         require.NoError(t, err)
 295
 296         whatsapp, _, _ := whatsappStartTest(t, dir, clock)
 297         subagents := newSubagents(filepath.Join(agent.dir, subagentDirName), clock.Now, subagentNameRandom)
 298         schedules, err := openSchedules(filepath.Join(agent.dir, scheduleFileName), clock.Now, discardLog())
 299         require.NoError(t, err)
 300
 301         server := newModelFake(t)
 302         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
 303         listener, err := daemonListenSocket(filepath.Join(dir, socketFileName))
 304         require.NoError(t, err)
 305
 306         ctx, cancel := context.WithCancel(t.Context())
 307         done := make(chan error, 1)
 308         go func() {
 309                 done <- daemonServe(ctx, listener, nil, whatsapp, schedules, func(notify func(Event)) (*TurnLoop, error) {
 310                         return newTurnLoop(agent, memory, subagents, whatsapp, schedules, client, clock.Now, discardLog(), notify)
 311                 })
 312         }()
 313
 314         t.Cleanup(func() {
 315                 cancel()
 316                 require.NoError(t, <-done)
 317                 require.NoError(t, subagents.Close())
 318                 require.NoError(t, memory.Close())
 319                 require.NoError(t, tokens.Close())
 320         })
 321
 322         return memory, whatsapp, subagents
 323 }