优化

internal/data/data_scope_audit_test.go

@@ -0,0 +1,72 @@
+package data
+
+import (
+	"io"
+	"log/slog"
+	"testing"
+	"time"
+
+	"github.com/glebarez/sqlite"
+	"gorm.io/gorm"
+)
+
+func openDataScopeAuditTestDB(t *testing.T, name string) *gorm.DB {
+	t.Helper()
+	db, err := gorm.Open(sqlite.Open("file:"+name+"?mode=memory&cache=shared"), &gorm.Config{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = db.AutoMigrate(&dataAccessLogPO{}); err != nil {
+		t.Fatal(err)
+	}
+	return db
+}
+
+func waitForDataScopeAuditCount(t *testing.T, db *gorm.DB, want int64) {
+	t.Helper()
+	deadline := time.Now().Add(time.Second)
+	for time.Now().Before(deadline) {
+		var count int64
+		if err := db.Model(&dataAccessLogPO{}).Count(&count).Error; err != nil {
+			t.Fatal(err)
+		}
+		if count == want {
+			return
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+	t.Fatalf("data access log count did not reach %d", want)
+}
+
+func TestDataScopeAuditWriterFlushesBatch(t *testing.T) {
+	db := openDataScopeAuditTestDB(t, "data-scope-audit-batch")
+	d := &Data{gormDB: newReloadableDB(db, nil)}
+	w := newDataScopeAuditWriterWithOptions(d, slog.New(slog.NewTextHandler(io.Discard, nil)), 8, 2, time.Hour)
+	d.auditLog = w
+	t.Cleanup(w.Close)
+
+	d.enqueueDataScopeAudit(dataAccessLogPO{EventType: "no_identity", TargetTable: "example"})
+	d.enqueueDataScopeAudit(dataAccessLogPO{EventType: "blocked_write", TargetTable: "example"})
+	waitForDataScopeAuditCount(t, db, 2)
+}
+
+func TestDataScopeAuditWriterUsesReloadedDatabase(t *testing.T) {
+	first := openDataScopeAuditTestDB(t, "data-scope-audit-first")
+	second := openDataScopeAuditTestDB(t, "data-scope-audit-second")
+	d := &Data{gormDB: newReloadableDB(first, nil)}
+	w := newDataScopeAuditWriterWithOptions(d, slog.New(slog.NewTextHandler(io.Discard, nil)), 8, 100, 10*time.Millisecond)
+	d.auditLog = w
+	t.Cleanup(w.Close)
+
+	d.enqueueDataScopeAudit(dataAccessLogPO{EventType: "no_identity", TargetTable: "example"})
+	d.gormDB.replace(second, nil)
+	waitForDataScopeAuditCount(t, second, 1)
+
+	var firstCount int64
+	if err := first.Model(&dataAccessLogPO{}).Count(&firstCount).Error; err != nil {
+		t.Fatal(err)
+	}
+	if firstCount != 0 {
+		t.Fatalf("audit log was written to retired database: %d", firstCount)
+	}
+}