package main

import (
	"context"
	"errors"
	"slices"
	"testing"
)

func TestSquares(t *testing.T) {
	got, err := Squares(context.Background(), []int{1, 2, 3, 4}, 2)
	if err != nil || !slices.Equal(got, []int{1, 4, 9, 16}) {
		t.Fatal(got, err)
	}
}
func TestAlreadyCanceled(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := Squares(ctx, make([]int, 100), 2)
	if !errors.Is(err, context.Canceled) {
		t.Fatal(err)
	}
}
func TestEmptyAndInvalid(t *testing.T) {
	if out, err := Squares(context.Background(), nil, 3); err != nil || len(out) != 0 {
		t.Fatal(out, err)
	}
	if _, err := Squares(context.Background(), nil, 0); err == nil {
		t.Fatal("expected error")
	}
}

func TestPublishDropPolicy(t *testing.T) {
	a, b := make(chan int, 1), make(chan int, 1)
	if n := Publish([]chan int{a, b}, 7); n != 0 {
		t.Fatal(n)
	}
	<-a
	if n := Publish([]chan int{a, b}, 9); n != 1 {
		t.Fatal(n)
	}
	if x, y := <-a, <-b; x != 9 || y != 7 {
		t.Fatal(x, y)
	}
	close(a)
	close(b)
}
func TestCancelPeers(t *testing.T) {
	if e := CancelPeers(); e == nil || e.Error() != "stage failed" {
		t.Fatal(e)
	}
}