66 lines
2.0 KiB
Go
66 lines
2.0 KiB
Go
package store
|
|
|
|
import (
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func openLocalDiscoveryTestStore(t *testing.T) *Store {
|
|
t.Helper()
|
|
st, err := Open(filepath.Join(t.TempDir(), "dispatcher.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() { _ = st.Close() })
|
|
return st
|
|
}
|
|
|
|
func seedLocalCommandAssignments(t *testing.T, st *Store, records ...LocalCommandRecord) {
|
|
t.Helper()
|
|
if len(records) == 0 {
|
|
return
|
|
}
|
|
dispatcherID := records[0].DispatcherID
|
|
tasks := make([]LocalDiscoveredTask, 0, len(records))
|
|
seen := make(map[string]struct{}, len(records))
|
|
for _, record := range records {
|
|
if record.DispatcherID != dispatcherID {
|
|
t.Fatalf("test commands use different Dispatcher IDs: %q and %q", dispatcherID, record.DispatcherID)
|
|
}
|
|
if _, exists := seen[record.TaskID]; exists {
|
|
continue
|
|
}
|
|
seen[record.TaskID] = struct{}{}
|
|
revision := int64(1)
|
|
if record.Admission != nil && record.Admission.Task.TaskRevision > 0 {
|
|
revision = record.Admission.Task.TaskRevision
|
|
}
|
|
tasks = append(tasks, localDiscoveredTask(dispatcherID, record.TaskID, record.TenantID, record.TenantKey, revision, "running"))
|
|
}
|
|
seedLocalSnapshotTasks(t, st, dispatcherID, tasks...)
|
|
}
|
|
|
|
func seedLocalDiscovery(t *testing.T, st *Store) {
|
|
t.Helper()
|
|
seedLocalSnapshotTasks(t, st, "d-1",
|
|
localDiscoveredTask("d-1", "task-a", "tenant-a", "tenant-key-a", 1, "running"),
|
|
localDiscoveredTask("d-1", "task-b", "tenant-a", "tenant-key-a", 1, "running"),
|
|
)
|
|
}
|
|
|
|
func seedLocalSnapshotTasks(t *testing.T, st *Store, dispatcherID string, tasks ...LocalDiscoveredTask) {
|
|
t.Helper()
|
|
now := time.Now().UTC()
|
|
if err := st.ApplyLocalTaskSnapshot(dispatcherID, tasks, now); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := st.MarkLocalTaskDiscoveryReadyV04(dispatcherID, now); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func localDiscoveredTask(_ string, taskID, tenantID, tenantKey string, revision int64, status string) LocalDiscoveredTask {
|
|
return LocalDiscoveredTask{TaskID: taskID, TenantID: tenantID, TenantKey: tenantKey, TaskRevision: revision, Status: status}
|
|
}
|