brain/kvbrain: implement forget by message id

brain/kvbrain/forget.go

@@ -3,6 +3,7 @@ package kvbrain
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"slices"
 	"sync"
 	"time"
@@ -102,7 +103,24 @@ func (br *Brain) Forget(ctx context.Context, tag string, tuples []brain.Tuple) e
 // ForgetMessage forgets everything learned from a single given message.
 // If nothing has been learned from the message, it should be ignored.
 func (br *Brain) ForgetMessage(ctx context.Context, tag string, msg uuid.UUID) error {
-	panic("not implemented") // TODO: Implement
+	past, _ := br.past.Load(tag)
+	if past == nil {
+		return nil
+	}
+	keys := past.findID(msg)
+	batch := br.knowledge.NewWriteBatch()
+	defer batch.Cancel()
+	for _, key := range keys {
+		err := batch.Delete(key)
+		if err != nil {
+			return err
+		}
+	}
+	err := batch.Flush()
+	if err != nil {
+		return fmt.Errorf("couldn't commit deleting message %v: %w", msg, err)
+	}
+	return nil
 }
 
 // ForgetDuring forgets all messages learned in the given time span.

brain/kvbrain/forget_test.go

@@ -2,11 +2,15 @@ package kvbrain
 
 import (
 	"bytes"
+	"context"
 	"slices"
 	"testing"
+	"time"
 
+	"github.com/dgraph-io/badger/v4"
 	"github.com/google/uuid"
 
+	"github.com/zephyrtronium/robot/brain"
 	"github.com/zephyrtronium/robot/userhash"
 )
 
@@ -204,3 +208,115 @@ func BenchmarkPastFindUser(b *testing.B) {
 
 //go:noinline
 func use(x [][]byte) {}
+
+func TestForgetMessage(t *testing.T) {
+	type message struct {
+		msg  brain.MessageMeta
+		tups []brain.Tuple
+	}
+	cases := []struct {
+		name string
+		msgs []message
+		uu   uuid.UUID
+		want map[string]string
+	}{
+		{
+			name: "single",
+			msgs: []message{
+				{
+					msg: brain.MessageMeta{
+						ID:   uuid.UUID{1},
+						User: userhash.Hash{2},
+						Tag:  "kessoku",
+						Time: time.Unix(0, 0),
+					},
+					tups: []brain.Tuple{
+						{Prefix: []string{"bocchi"}, Suffix: "ryou"},
+					},
+				},
+			},
+			uu:   uuid.UUID{1},
+			want: map[string]string{},
+		},
+		{
+			name: "several",
+			msgs: []message{
+				{
+					msg: brain.MessageMeta{
+						ID:   uuid.UUID{1},
+						User: userhash.Hash{2},
+						Tag:  "kessoku",
+						Time: time.Unix(0, 0),
+					},
+					tups: []brain.Tuple{
+						{Prefix: []string{"bocchi"}, Suffix: "ryou"},
+						{Prefix: []string{"nijika"}, Suffix: "kita"},
+					},
+				},
+			},
+			uu:   uuid.UUID{1},
+			want: map[string]string{},
+		},
+		{
+			name: "tagged",
+			msgs: []message{
+				{
+					msg: brain.MessageMeta{
+						ID:   uuid.UUID{1},
+						User: userhash.Hash{2},
+						Tag:  "sickhack",
+						Time: time.Unix(0, 0),
+					},
+					tups: []brain.Tuple{
+						{Prefix: []string{"bocchi"}, Suffix: "ryou"},
+					},
+				},
+			},
+			uu: uuid.UUID{1},
+			want: map[string]string{
+				mkey("sickhack", "bocchi\xff\xff", uuid.UUID{1}): "ryou",
+			},
+		},
+		{
+			name: "unseen",
+			msgs: []message{
+				{
+					msg: brain.MessageMeta{
+						ID:   uuid.UUID{1},
+						User: userhash.Hash{2},
+						Tag:  "kessoku",
+						Time: time.Unix(0, 0),
+					},
+					tups: []brain.Tuple{
+						{Prefix: []string{"bocchi"}, Suffix: "ryou"},
+					},
+				},
+			},
+			uu: uuid.UUID{2},
+			want: map[string]string{
+				mkey("kessoku", "bocchi\xff\xff", uuid.UUID{1}): "ryou",
+			},
+		},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			t.Parallel()
+			ctx := context.Background()
+			db, err := badger.Open(badger.DefaultOptions("").WithInMemory(true).WithLogger(nil))
+			if err != nil {
+				t.Fatal(err)
+			}
+			br := New(db)
+			for _, msg := range c.msgs {
+				err := br.Learn(ctx, &msg.msg, msg.tups)
+				if err != nil {
+					t.Errorf("failed to learn: %v", err)
+				}
+			}
+			if err := br.ForgetMessage(ctx, "kessoku", c.uu); err != nil {
+				t.Errorf("couldn't forget: %v", err)
+			}
+			dbcheck(t, db, c.want)
+		})
+	}
+}

brain/kvbrain/learn_test.go

@@ -13,14 +13,44 @@ import (
 	"github.com/zephyrtronium/robot/userhash"
 )
 
-func TestLearn(t *testing.T) {
-	mkey := func(tag, toks string, id uuid.UUID) string {
-		b := make([]byte, 8, 8+len(toks)+len(id))
-		binary.LittleEndian.PutUint64(b, hashTag(tag))
-		b = append(b, toks...)
-		b = append(b, id[:]...)
-		return string(b)
+func mkey(tag, toks string, id uuid.UUID) string {
+	b := make([]byte, 8, 8+len(toks)+len(id))
+	binary.LittleEndian.PutUint64(b, hashTag(tag))
+	b = append(b, toks...)
+	b = append(b, id[:]...)
+	return string(b)
+}
+
+func dbcheck(t *testing.T, db *badger.DB, want map[string]string) {
+	t.Helper()
+	seen := 0
+	err := db.View(func(txn *badger.Txn) error {
+		opts := badger.IteratorOptions{}
+		it := txn.NewIterator(opts)
+		defer it.Close()
+		for it.Rewind(); it.Valid(); it.Next() {
+			item := it.Item()
+			k := string(item.Key())
+			v, err := item.ValueCopy(nil)
+			if err != nil {
+				t.Errorf("couldn't get value for key %q: %v", k, err)
+			}
+			if got := string(v); want[k] != got {
+				t.Errorf("wrong value for key %q: want %q, got %q", k, want[k], got)
+			}
+			seen++
+		}
+		return nil
+	})
+	if err != nil {
+		t.Errorf("view failed: %v", err)
+	}
+	if seen != len(want) {
+		t.Errorf("saw wrong number of items: want %d, got %d", len(want), seen)
 	}
+}
+
+func TestLearn(t *testing.T) {
 	uu := uuid.UUID{':', ')', ':', ')', ':', ')', ':', ')', ':', ')', ':', ')', ':', ')', ':', ')'}
 	h := userhash.Hash{2}
 	cases := []struct {
@@ -103,31 +133,7 @@ func TestLearn(t *testing.T) {
 			if err := br.Learn(ctx, &c.msg, c.tups); err != nil {
 				t.Errorf("failed to learn: %v", err)
 			}
-			seen := 0
-			err = db.View(func(txn *badger.Txn) error {
-				opts := badger.IteratorOptions{}
-				it := txn.NewIterator(opts)
-				defer it.Close()
-				for it.Rewind(); it.Valid(); it.Next() {
-					item := it.Item()
-					k := string(item.Key())
-					v, err := item.ValueCopy(nil)
-					if err != nil {
-						t.Errorf("couldn't get value for key %q: %v", k, err)
-					}
-					if got := string(v); c.want[k] != got {
-						t.Errorf("wrong value for key %q: want %q, got %q", k, c.want[k], got)
-					}
-					seen++
-				}
-				return nil
-			})
-			if err != nil {
-				t.Errorf("view failed: %v", err)
-			}
-			if seen != len(c.want) {
-				t.Errorf("saw wrong number of items: want %d, got %d", len(c.want), seen)
-			}
+			dbcheck(t, db, c.want)
 		})
 	}
 }