threshold/internal/test/timeout_handler.go

package test

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/luxfi/threshold/pkg/party"
	"github.com/luxfi/threshold/pkg/protocol"
)

// HandlerLoopWithTimeout runs a handler with proper timeout and cleanup
func HandlerLoopWithTimeout(t testing.TB, id party.ID, h *protocol.Handler, network *Network, timeout time.Duration) error {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	done := make(chan struct{})
	errChan := make(chan error, 1)

	go func() {
		defer close(done)

		for {
			select {
			case <-ctx.Done():
				return

			case msg, ok := <-h.Listen():
				if !ok {
					// Handler finished successfully
					return
				}
				if msg != nil {
					network.Send(msg)
				}

			case msg := <-network.Next(id):
				if msg != nil {
					h.Accept(msg)
				}
			}
		}
	}()

	// Wait for completion or timeout
	select {
	case <-done:
		// Clean shutdown of network connection
		select {
		case <-network.Done(id):
		case <-time.After(100 * time.Millisecond):
		}
		return nil

	case err := <-errChan:
		return err

	case <-ctx.Done():
		// Timeout - clean shutdown
		select {
		case <-network.Done(id):
		case <-time.After(100 * time.Millisecond):
		}
		return ctx.Err()
	}
}

// RunProtocolHandlers drives pre-built handlers under a deadline and returns
// whatever results landed before the deadline. Use this when the caller
// already owns Handler construction (e.g. Doerner's sender/receiver split);
// otherwise RunProtocolWithTimeout (which constructs handlers from a StartFunc
// factory) is the simpler entry point.
func RunProtocolHandlers(t testing.TB, partyIDs []party.ID, timeout time.Duration, createHandlers func() map[party.ID]*protocol.Handler) (map[party.ID]interface{}, error) {
	network := NewNetwork(partyIDs)
	handlers := createHandlers()
	results := make(map[party.ID]interface{})
	errors := make(map[party.ID]error)

	var wg sync.WaitGroup
	var mu sync.Mutex

	for id, h := range handlers {
		if h == nil {
			continue
		}

		wg.Add(1)
		go func(partyID party.ID, handler *protocol.Handler) {
			defer wg.Done()

			// Run with timeout
			err := HandlerLoopWithTimeout(t, partyID, handler, network, timeout)

			mu.Lock()
			defer mu.Unlock()

			if err != nil {
				errors[partyID] = err
			} else {
				// Try to get result
				if result, err := handler.Result(); err == nil {
					results[partyID] = result
				} else {
					errors[partyID] = err
				}
			}
		}(id, h)
	}

	// Wait for all handlers or timeout
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		// All handlers completed
	case <-time.After(timeout + time.Second):
		// Global timeout exceeded
		if t != nil {
			t.Log("Global timeout exceeded")
		}
	}

	// Clean up all handlers
	for _, h := range handlers {
		if h != nil {
			h.Stop()
		}
	}

	// Return partial results even on timeout
	if len(errors) > 0 && len(results) == 0 {
		// All failed, return first error
		for _, err := range errors {
			return nil, err
		}
	}

	return results, nil
}

// RunInitCheck runs a boolean init/smoke callback under a 5s deadline as a
// named sub-test. A `false` return fails the test; a deadline expiry is
// logged but does not fail (init checks are tolerant of slow setups). For
// full-protocol round-trips use the PhaseHarness directly.
func RunInitCheck(t *testing.T, name string, n int, threshold int, testFunc func(partyIDs []party.ID) bool) {
	t.Run(name, func(t *testing.T) {
		partyIDs := PartyIDs(n)

		// Run test with timeout
		done := make(chan bool, 1)
		go func() {
			done <- testFunc(partyIDs)
		}()

		select {
		case success := <-done:
			if !success {
				t.Error("init check returned false")
			}
		case <-time.After(5 * time.Second):
			// Don't fail on timeout, just log it
			t.Log("init check deadline reached without callback completing")
		}
	})
}