fn-serverless/api/common/wait_utils_test.go

package common

import (
	"testing"
	"time"
)

func isClosed(ch chan struct{}) bool {
	select {
	case <-ch:
		return true
	default:
	}
	return false
}

func TestWaitGroupEmpty(t *testing.T) {

	wg := NewWaitGroup()

	if !wg.AddSession(0) {
		t.Fatalf("Add 0 should not fail")
	}

	if isClosed(wg.Closer()) {
		t.Fatalf("Should not be closed yet")
	}

	done := wg.CloseGroupNB()

	// gate-on close
	wg.CloseGroup()

	if !isClosed(wg.Closer()) {
		t.Fatalf("Should be closing state")
	}

	if isClosed(done) {
		t.Fatalf("NB Chan I should be closed")
	}

	done = wg.CloseGroupNB()
	if isClosed(done) {
		t.Fatalf("NB Chan II should be closed")
	}
}

func TestWaitGroupBlast(t *testing.T) {
	wg := NewWaitGroup()
	wg.AddSession(1)

	blastRadius := 500
	for i := 0; i < blastRadius; i++ {
		go func(i int) {
			if !wg.AddSession(1) {
				// this is OK, we are creating a race
				// and some cannot make it after main
				// go-routine issues a CloseGroupNB()
				return
			}

			if isClosed(wg.Closer()) {
				t.Fatalf("Should not be closing state")
			}

			wg.DoneSession()
		}(i)
	}

	if isClosed(wg.Closer()) {
		t.Fatalf("Should not be closing state")
	}

	done := wg.CloseGroupNB()

	if !isClosed(wg.Closer()) {
		t.Fatalf("Should be closing state")
	}

	select {
	case <-done:
		t.Fatalf("NB Chan should not be closed yet, since sum is 1")
	case <-time.After(time.Duration(1) * time.Second):
		wg.DoneSession()
	}

	select {
	case <-done:
	case <-time.After(time.Duration(1) * time.Second):
		t.Fatalf("NB Chan should be closed by now, since sum is 0")
	}
}

func TestWaitGroupSingle(t *testing.T) {

	wg := NewWaitGroup()

	if isClosed(wg.Closer()) {
		t.Fatalf("Should not be closing state yet")
	}

	if !wg.AddSession(1) {
		t.Fatalf("Add 1 should not fail")
	}

	if isClosed(wg.Closer()) {
		t.Fatalf("Should not be closing state yet")
	}

	wg.DoneSession()
	// sum should be zero now.

	if !wg.AddSession(2) {
		t.Fatalf("Add 2 should not fail")
	}

	// sum is 2 now
	// initiate shutdown
	done := wg.CloseGroupNB()

	if isClosed(done) {
		t.Fatalf("NB Chan should not be closed yet, since sum is 2")
	}

	wg.DoneSession()

	if wg.AddSession(1) {
		t.Fatalf("Add 1 should fail (we are shutting down)")
	}
	if !isClosed(wg.Closer()) {
		t.Fatalf("Should be closing state")
	}

	// sum is 1 now

	if isClosed(done) {
		t.Fatalf("NB Chan should not be closed yet, since sum is 1")
	}

	if wg.AddSession(0) {
		t.Fatalf("Add 0 should fail (considered positive number and we are closing)")
	}

	if wg.AddSession(100) {
		t.Fatalf("Add 100 should fail (we are shutting down)")
	}

	if !isClosed(wg.Closer()) {
		t.Fatalf("Should be closing state")
	}

	wg.DoneSession()

	// sum is 0 now
	<-done

	if !isClosed(done) {
		t.Fatalf("NB Chan should be closed, since sum is 0")
	}

	if !isClosed(wg.Closer()) {
		t.Fatalf("Should be closing state")
	}
}