doormile_backend/internal/assignment/balance_pg_test.go

package assignment

import (
	"errors"
	"fmt"
	"sync"
	"testing"
	"time"

	"doormile/constants"
	"doormile/models"

	"github.com/redis/go-redis/v9"
	"gorm.io/gorm"
)

// Balanced auto-assignment (docs/balanced-assignment-plan.md, Phase A)
// against a real Postgres. Skipped unless REGISTRY_TEST_DSN is set — see
// eligibility_pg_test.go.

// onDuty adds a rider with live GPS and an open duty session.
func onDuty(t *testing.T, gdb *gorm.DB, id int, sessionStart time.Time) {
	t.Helper()
	rider(t, gdb, id, constants.MilerAvailable, time.Minute)
	mustDo(t, gdb.Create(&models.MilerDutyLog{Userid: id, Loginat: sessionStart}).Error)
}

// assignOne runs one attempt the way tryAssign does (minus the Redis GEO
// lookup, which `nearby` stands in for) and returns the rider it went to.
func assignOne(t *testing.T, gdb *gorm.DB, bookingID int, nearby []redis.GeoLocation) (int, error) {
	t.Helper()
	var b models.PickupBooking
	if err := gdb.First(&b, bookingID).Error; err != nil {
		return 0, err
	}
	cand, pool, dec, found := selectMilerWithAI(&b, nearby)
	if !found {
		return 0, nil
	}
	final, err := commitAssignment(&b, cand, pool, dec)
	if err != nil {
		return 0, err
	}
	return final.profile.Userid, nil
}

// near puts every rider at a different distance from the pickup, so a
// distance-first rule would pile everything onto rider ids[0].
func near(ids ...int) []redis.GeoLocation {
	out := make([]redis.GeoLocation, len(ids))
	for i, id := range ids {
		out[i] = redis.GeoLocation{Name: fmt.Sprint(id), Dist: 0.5 + float64(i)*1.5}
	}
	return out
}

func perRider(t *testing.T, gdb *gorm.DB) map[int]int64 {
	t.Helper()
	type row struct {
		Mileruserid int
		N           int64
	}
	var rows []row
	mustDo(t, gdb.Table("bookingassignments").Select("mileruserid, COUNT(*) AS n").
		Where("assignmentstatus IN ?", []string{constants.AssignmentAssigned, constants.AssignmentAccepted}).
		Group("mileruserid").Scan(&rows).Error)
	out := map[int]int64{}
	for _, r := range rows {
		out[r.Mileruserid] = r.N
	}
	return out
}

func spread(m map[int]int64, ids ...int) (lo, hi int64) {
	lo, hi = 1<<62, -1
	for _, id := range ids {
		n := m[id]
		if n < lo {
			lo = n
		}
		if n > hi {
			hi = n
		}
	}
	return lo, hi
}

func balanceSetup(t *testing.T) *gorm.DB {
	gdb := eligibilityDB(t)
	t.Setenv("AI_LAYER_BASE_URL", "http://127.0.0.1:1") // unreachable: the fallback chooses
	t.Setenv("MILER_MAX_ACTIVE_BOOKINGS", "20")
	return gdb
}

// 5 riders, 50 orders one after another: 10 each, although rider 1 is the
// nearest to every pickup.
func TestBalancedSplitSequential(t *testing.T) {
	gdb := balanceSetup(t)
	ids := []int{1, 2, 3, 4, 5}
	for _, id := range ids {
		onDuty(t, gdb, id, time.Now().Add(-time.Hour))
	}
	for i := 0; i < 50; i++ {
		b := order(t, gdb, 0, constants.BookingPendingPickup, "", time.Time{})
		if _, err := assignOne(t, gdb, b, near(ids...)); err != nil {
			t.Fatal(err)
		}
	}
	got := perRider(t, gdb)
	for _, id := range ids {
		if got[id] != 10 {
			t.Fatalf("per rider = %v, want 10 each", got)
		}
	}
}

// Any count splits within one: 23 orders over 5 riders -> 5,5,5,4,4.
func TestBalancedSplitUneven(t *testing.T) {
	gdb := balanceSetup(t)
	ids := []int{1, 2, 3, 4, 5}
	for _, id := range ids {
		onDuty(t, gdb, id, time.Now().Add(-time.Hour))
	}
	for i := 0; i < 23; i++ {
		b := order(t, gdb, 0, constants.BookingPendingPickup, "", time.Time{})
		if _, err := assignOne(t, gdb, b, near(ids...)); err != nil {
			t.Fatal(err)
		}
	}
	if lo, hi := spread(perRider(t, gdb), ids...); lo != 4 || hi != 5 {
		t.Fatalf("spread %d..%d, want 4..5", lo, hi)
	}
}

// A bulk upload: 50 orders attempted at the same moment. Without the final
// re-check under the lock they all saw the same counts and piled onto one
// rider. Still balanced, nobody over the ceiling, one rider per order.
func TestBalancedSplitConcurrentBurst(t *testing.T) {
	gdb := balanceSetup(t)
	ids := []int{1, 2, 3, 4, 5}
	for _, id := range ids {
		onDuty(t, gdb, id, time.Now().Add(-time.Hour))
	}
	bookings := make([]int, 50)
	for i := range bookings {
		bookings[i] = order(t, gdb, 0, constants.BookingPendingPickup, "", time.Time{})
	}
	var wg sync.WaitGroup
	errs := make(chan error, len(bookings))
	for _, b := range bookings {
		wg.Add(1)
		go func(b int) {
			defer wg.Done()
			if _, err := assignOne(t, gdb, b, near(ids...)); err != nil {
				errs <- err
			}
		}(b)
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		t.Fatal(err)
	}
	got := perRider(t, gdb)
	if lo, hi := spread(got, ids...); hi-lo > 1 {
		t.Fatalf("burst split %v: spread %d..%d, want within 1", got, lo, hi)
	}
	var total int64
	for _, n := range got {
		total += n
	}
	var perOrder int64
	gdb.Raw("SELECT COALESCE(MAX(n),0) FROM (SELECT COUNT(*) n FROM bookingassignments GROUP BY bookingid) x").Scan(&perOrder)
	if total != 50 || perOrder != 1 {
		t.Fatalf("assigned %d (want 50), max riders on one order %d (want 1)", total, perOrder)
	}
}

// Riders come and go: 3 riders hold 4 each, 2 log in fresh. The next 8 orders
// go to the newcomers (4 each), then everyone shares.
func TestNewRidersCatchUpThenShare(t *testing.T) {
	gdb := balanceSetup(t)
	busy := []int{1, 2, 3}
	fresh := []int{4, 5}
	for _, id := range busy {
		onDuty(t, gdb, id, time.Now().Add(-3*time.Hour))
		for i := 0; i < 4; i++ {
			order(t, gdb, id, constants.BookingMilerAssigned, constants.AssignmentAccepted, time.Now().Add(-time.Hour))
		}
	}
	for _, id := range fresh {
		onDuty(t, gdb, id, time.Now())
	}
	all := append(append([]int{}, busy...), fresh...)
	for i := 0; i < 8; i++ {
		b := order(t, gdb, 0, constants.BookingPendingPickup, "", time.Time{})
		id, err := assignOne(t, gdb, b, near(all...))
		if err != nil {
			t.Fatal(err)
		}
		if id != 4 && id != 5 {
			t.Fatalf("order %d went to busy rider %d while a newcomer had fewer", i, id)
		}
	}
	if lo, hi := spread(perRider(t, gdb), all...); lo != 4 || hi != 4 {
		t.Fatalf("after catch-up spread %d..%d, want all 4", lo, hi)
	}
	// From here on everyone shares: 5 more orders -> one each.
	for i := 0; i < 5; i++ {
		b := order(t, gdb, 0, constants.BookingPendingPickup, "", time.Time{})
		if _, err := assignOne(t, gdb, b, near(all...)); err != nil {
			t.Fatal(err)
		}
	}
	if lo, hi := spread(perRider(t, gdb), all...); lo != 5 || hi != 5 {
		t.Fatalf("after sharing spread %d..%d, want all 5", lo, hi)
	}
}

// Tie on orders in hand: fewer orders this duty session wins, then distance.
func TestTieBreakSessionThenDistance(t *testing.T) {
	gdb := balanceSetup(t)
	// Rider 1: nearest, but has had 2 orders this session (both now closed).
	onDuty(t, gdb, 1, time.Now().Add(-2*time.Hour))
	for i := 0; i < 2; i++ {
		order(t, gdb, 1, constants.BookingConvertedConsignment, constants.AssignmentCompleted, time.Now().Add(-time.Hour))
	}
	// Riders 2 and 3: no orders this session; 2 is nearer than 3.
	onDuty(t, gdb, 2, time.Now().Add(-2*time.Hour))
	onDuty(t, gdb, 3, time.Now().Add(-2*time.Hour))

	b := order(t, gdb, 0, constants.BookingPendingPickup, "", time.Time{})
	id, err := assignOne(t, gdb, b, near(1, 2, 3))
	if err != nil {
		t.Fatal(err)
	}
	if id != 2 {
		t.Fatalf("went to %d, want 2 (fewest this session, then nearest)", id)
	}
}

// Everyone at the ceiling: the order waits (no assignment, no error that
// would stop the retries).
func TestAllRidersAtCeilingWaits(t *testing.T) {
	gdb := balanceSetup(t)
	t.Setenv("MILER_MAX_ACTIVE_BOOKINGS", "2")
	for _, id := range []int{1, 2} {
		onDuty(t, gdb, id, time.Now().Add(-time.Hour))
		for i := 0; i < 2; i++ {
			order(t, gdb, id, constants.BookingMilerAssigned, constants.AssignmentAccepted, time.Now())
		}
	}
	b := order(t, gdb, 0, constants.BookingPendingPickup, "", time.Time{})
	id, err := assignOne(t, gdb, b, near(1, 2))
	if err != nil && !errors.Is(err, errNoRiderCapacity) {
		t.Fatal(err)
	}
	if id != 0 {
		t.Fatalf("assigned to %d although everyone is at the ceiling", id)
	}
	var bk models.PickupBooking
	mustDo(t, gdb.First(&bk, b).Error)
	if bk.Status != constants.BookingPendingPickup || bk.Assignedmileruserid != nil {
		t.Fatalf("booking should still be pending: %+v", bk.Status)
	}
}