orderedLogDrainer_test.go

package main

import (
	"fmt"
	"sync"
	"testing"
)

func numLogsAndNumMsgs(numLogs int, numMsgs int) chan chan string {
	logs := make(chan chan string)

	go func() {
		log := make([]chan string, numLogs)
		var wg sync.WaitGroup

		for i := 0; i < numLogs; i++ {
			log[i] = make(chan string)
			logs <- log[i]
			wg.Add(1)
			go func(i int) {
				for j := 1; j <= numMsgs; j++ {
					log[i] <- testMsg(i+1, j)
				}
				close(log[i])
				wg.Done()
			}(i)
		}

		go func() {
			wg.Wait()
			close(logs)
		}()
	}()

	return logs
}

func testMsg(log int, msg int) string {
	return fmt.Sprintf("message %d in %d", msg, log)
}

func orderedNumLogsAndNumMsgs(numLogs int, numMsgs int) chan string {
	out := make(chan string)
	go func() {
		for i := 1; i <= numLogs; i++ {
			for j := 1; j <= numMsgs; j++ {
				out <- fmt.Sprintf("%d\t%s", i, testMsg(i, j))
			}
		}
		close(out)
	}()
	return out
}

func TestOrderedLogDrainer(t *testing.T) {
	cases := []struct {
		in   chan chan string
		want chan string
	}{
		{
			numLogsAndNumMsgs(1, 10),
			orderedNumLogsAndNumMsgs(1, 10),
		},
		{
			numLogsAndNumMsgs(10, 10),
			orderedNumLogsAndNumMsgs(10, 10),
		},
		{
			numLogsAndNumMsgs(10, 50),
			orderedNumLogsAndNumMsgs(10, 50),
		},
		{
			numLogsAndNumMsgs(100, 100),
			orderedNumLogsAndNumMsgs(100, 100),
		},
		{
			numLogsAndNumMsgs(1000, 100),
			orderedNumLogsAndNumMsgs(1000, 100),
		},
	}
	for _, c := range cases {
		result := OrderedLogDrainer(c.in)

		for wanted := range c.want {
			got := <-result
			if wanted != got {
				t.Errorf("Outputs doesn't match:\n wanted %s != got %s\n", wanted, got)
			}
		}
	}
}

func TestReadBeforeClose(t *testing.T) {
        logs := make(chan chan string)
        log := make(chan string)

        result := OrderedLogDrainer(logs)

        logs <- log
        log <- "message!"

        if res := <-result; "1\tmessage!" != res {
                t.Errorf("Outputs doesn't match:\n wanted 1\tmessage != got %s\n", res)
        }
        close(logs)
        close(log)
}