package activator

import (
	"context"
	"testing"
	"time"
)

func TestCleanupBudgetSharedAfterCancellation(t *testing.T) {
	t.Parallel()
	parent, stop := context.WithCancel(t.Context())
	ctx, finish := WithCleanupBudget(parent, 80*time.Millisecond)
	defer finish()
	stop()
	first, cancel := CleanupContext(ctx, time.Second)
	defer cancel()
	deadline, ok := first.Deadline()
	if !ok || time.Until(deadline) > 100*time.Millisecond {
		t.Fatal("cleanup escaped shutdown budget")
	}
	nested, done := CleanupContext(first, time.Hour)
	defer done()
	if next, _ := nested.Deadline(); next.After(deadline) {
		t.Fatal("nested cleanup extended deadline")
	}
	<-first.Done()
	later, end := CleanupContext(ctx, time.Second)
	defer end()
	if later.Err() == nil {
		t.Fatal("cleanup budget restarted")
	}
}

func TestCleanupBudgetDoesNotLimitNormalOperation(t *testing.T) {
	ctx, finish := WithCleanupBudget(t.Context(), time.Millisecond)
	defer finish()
	time.Sleep(3 * time.Millisecond)
	cleanup, cancel := CleanupContext(ctx, time.Second)
	defer cancel()
	if cleanup.Err() != nil || ctx.Err() != nil {
		t.Fatal("normal request exhausted shutdown budget")
	}
}

func TestCleanupBudgetNestedOwnerUsesShutdownDeadline(t *testing.T) {
	parent, stop := context.WithCancel(t.Context())
	parent, finish := WithCleanupBudget(parent, 60*time.Millisecond)
	defer finish()
	request, cancel := context.WithCancel(parent)
	defer cancel()
	request, end := WithCleanupBudget(request, time.Second)
	defer end()
	stop()
	owner, done := CleanupContext(parent, time.Second)
	defer done()
	handler, joined := CleanupContext(request, time.Second)
	defer joined()
	a, _ := owner.Deadline()
	b, _ := handler.Deadline()
	if !a.Equal(b) {
		t.Fatalf("handler extended shared deadline: %v %v", a, b)
	}
}