<commit_before>package transloadit

import (
	"fmt"
	"os"
	"strings"
	"testing"
)

var assemblyId string
var assemblyUrl string

func TestAssembly(t *testing.T) {

	client := setup(t)

	assembly := client.CreateAssembly()

	file, err := os.Open("./fixtures/lol_cat.jpg")
	if err != nil {
		t.Fatal(err)
	}

	file2, err := os.Open("./fixtures/mona_lisa.jpg")
	if err != nil {
		t.Fatal(err)
	}

	assembly.AddReader("image", "lol_cat.jpg", file)
	assembly.AddReader("image2", "mona_lisa.jpg", file2)

	assembly.AddStep("resize", map[string]interface{}{
		"robot":           "/image/resize",
		"width":           75,
		"height":          75,
		"resize_strategy": "pad",
		"background":      "#000000",
	})

	assembly.NotifyUrl = "http://requestb.in/1kwp6lx1"

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.AssemblyId == "" {
		t.Fatal("response doesn't contain assembly_id")
	}

	if info.NotifyUrl != "http://requestb.in/1kwp6lx1" {
		t.Fatal("wrong notify url")
	}

	if len(info.Uploads) != 2 {
		t.Fatal("wrong number of uploads")
	}

	if info.Uploads[0].Name == "lol_cat.jpg" {
		if info.Uploads[0].Field != "image" {
			t.Fatal("wrong field name")
		}
	} else if info.Uploads[1].Name == "lol_cat.jpg" {
		if info.Uploads[1].Field != "image" {
			t.Fatal("wrong field name")
		}
	} else {
		t.Fatal("lol_cat.jpg not found in uploads")
	}

	if info.ClientAgent != "Transloadit Go SDK v1" {
		t.Fatal("wrong user agent")
	}

	assemblyId = info.AssemblyId
	assemblyUrl = info.AssemblyUrl
}

func TestAssemblyFail(t *testing.T) {

	config := DefaultConfig
	config.AuthKey = "does not exist"
	config.AuthSecret = "does not matter"

	client, err := NewClient(&config)
	if err != nil {
		t.Fatal(err)
	}

	assembly := client.CreateAssembly()

	file, err := os.Open("./fixtures/lol_cat.jpg")
	if err != nil {
		t.Fatal(err)
	}

	assembly.AddReader("image", "lol_cat.jpg", file)

	assembly.AddStep("resize", map[string]interface{}{
		"robot":           "/image/resize",
		"width":           75,
		"height":          75,
		"resize_strategy": "pad",
		"background":      "#000000",
	})

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.Error != "GET_ACCOUNT_UNKNOWN_AUTH_KEY" {
		t.Fatal("reponse doesn't contain error message")
	}

}

func TestAssemblyBlocking(t *testing.T) {

	client := setup(t)

	assembly := client.CreateAssembly()

	file, err := os.Open("./fixtures/lol_cat.jpg")
	if err != nil {
		t.Fatal(err)
	}

	assembly.AddReader("image", "lol_cat.jpg", file)

	assembly.AddStep("resize", map[string]interface{}{
		"robot":           "/image/resize",
		"width":           75,
		"height":          75,
		"resize_strategy": "pad",
		"background":      "#000000",
	})

	assembly.Blocking = true

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.Ok != "ASSEMBLY_COMPLETED" {
		t.Fatal("wrong assembly status")
	}

	if info.AssemblyId == "" {
		t.Fatal("response doesn't contain assembly_id")
	}

	if len(info.Uploads) != 1 {
		t.Fatal("wrong number of uploads")
	}
}

func TestGetAssembly(t *testing.T) {

	client := setup(t)

	assembly, err := client.GetAssembly(assemblyUrl)
	if err != nil {
		t.Fatal(err)
	}

	if assembly.AssemblyId == "" {
		t.Fatal("assembly id not contained")
	}

	if assembly.AssemblyUrl != assemblyUrl {
		t.Fatal("assembly urls don't match")
	}

}

func TestReplayAssembly(t *testing.T) {

	client := setup(t)

	assembly := client.ReplayAssembly(assemblyId)

	assembly.NotifyUrl = "http://requestb.in/1kwp6lx1"
	assembly.ReparseTemplate()

	info, err := assembly.Start()
	if err != nil {
		t.Fatal(err)
	}

	if info.Ok != "ASSEMBLY_REPLAYING" {
		t.Fatal("wrong status code returned")
	}

	if info.NotifyUrl != "http://requestb.in/1kwp6lx1" {
		t.Fatal("wrong notify url")
	}

}

func TestReplayAssemblyBlocking(t *testing.T) {

	client := setup(t)

	assembly := client.ReplayAssembly(assemblyId)

	assembly.Blocking = true

	info, err := assembly.Start()
	if err != nil {
		t.Fatal(err)
	}

	if info.Ok != "ASSEMBLY_COMPLETED" {
		t.Fatal("wrong status code returned")
	}

}

func TestAssemblyUsingTemplate(t *testing.T) {

	client := setup(t)

	assembly := client.CreateAssembly()

	assembly.TemplateId = "64c11b20308811e4b5548d4f316c150f"

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.AssemblyId == "" {
		t.Fatal(fmt.Sprintf("response doesn't contain assembly_id. %s", info.Error))
	}

	if !strings.Contains(info.Params, "64c11b20308811e4b5548d4f316c150f") {
		t.Fatal("template id not as parameter submitted")
	}
}

func TestCancelAssembly(t *testing.T) {

	client := setup(t)

	assembly := client.CreateAssembly()

	assembly.AddStep("import", map[string]interface{}{
		"robot": "/http/import",
		"url":   "http://mirror.nl.leaseweb.net/speedtest/10000mb.bin",
	})

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.AssemblyUrl == "" {
		t.Fatal("response doesn't contain assembly_url")
	}

	_, err = client.CancelAssembly(info.AssemblyUrl)
	if err != nil {
		t.Fatal(err)
	}

}

func TestListAssemblies(t *testing.T) {

	client := setup(t)

	assemblies, err := client.ListAssemblies(&ListOptions{
		PageSize: 3,
	})
	if err != nil {
		t.Fatal(err)
	}

	if len(assemblies.Assemblies) < 3 {
		t.Fatal("wrong number of assemblies")
	}

	if assemblies.Count == 0 {
		t.Fatal("wrong count")
	}

	if assemblies.Assemblies[0].AssemblyId == "" {
		t.Fatal("wrong template name")
	}

}
<commit_msg>Test returned error in TestAssemblyFail<commit_after>package transloadit

import (
	"fmt"
	"os"
	"strings"
	"testing"
)

var assemblyId string
var assemblyUrl string

func TestAssembly(t *testing.T) {

	client := setup(t)

	assembly := client.CreateAssembly()

	file, err := os.Open("./fixtures/lol_cat.jpg")
	if err != nil {
		t.Fatal(err)
	}

	file2, err := os.Open("./fixtures/mona_lisa.jpg")
	if err != nil {
		t.Fatal(err)
	}

	assembly.AddReader("image", "lol_cat.jpg", file)
	assembly.AddReader("image2", "mona_lisa.jpg", file2)

	assembly.AddStep("resize", map[string]interface{}{
		"robot":           "/image/resize",
		"width":           75,
		"height":          75,
		"resize_strategy": "pad",
		"background":      "#000000",
	})

	assembly.NotifyUrl = "http://requestb.in/1kwp6lx1"

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.AssemblyId == "" {
		t.Fatal("response doesn't contain assembly_id")
	}

	if info.NotifyUrl != "http://requestb.in/1kwp6lx1" {
		t.Fatal("wrong notify url")
	}

	if len(info.Uploads) != 2 {
		t.Fatal("wrong number of uploads")
	}

	if info.Uploads[0].Name == "lol_cat.jpg" {
		if info.Uploads[0].Field != "image" {
			t.Fatal("wrong field name")
		}
	} else if info.Uploads[1].Name == "lol_cat.jpg" {
		if info.Uploads[1].Field != "image" {
			t.Fatal("wrong field name")
		}
	} else {
		t.Fatal("lol_cat.jpg not found in uploads")
	}

	if info.ClientAgent != "Transloadit Go SDK v1" {
		t.Fatal("wrong user agent")
	}

	assemblyId = info.AssemblyId
	assemblyUrl = info.AssemblyUrl
}

func TestAssemblyFail(t *testing.T) {

	config := DefaultConfig
	config.AuthKey = "does not exist"
	config.AuthSecret = "does not matter"

	client, err := NewClient(&config)
	if err != nil {
		t.Fatal(err)
	}

	assembly := client.CreateAssembly()

	file, err := os.Open("./fixtures/lol_cat.jpg")
	if err != nil {
		t.Fatal(err)
	}

	assembly.AddReader("image", "lol_cat.jpg", file)

	assembly.AddStep("resize", map[string]interface{}{
		"robot":           "/image/resize",
		"width":           75,
		"height":          75,
		"resize_strategy": "pad",
		"background":      "#000000",
	})

	info, err := assembly.Upload()
	if err == nil {
		t.Fatal("no error returned")
	}

	if info.Error != "GET_ACCOUNT_UNKNOWN_AUTH_KEY" {
		t.Fatal("reponse doesn't contain error message")
	}

}

func TestAssemblyBlocking(t *testing.T) {

	client := setup(t)

	assembly := client.CreateAssembly()

	file, err := os.Open("./fixtures/lol_cat.jpg")
	if err != nil {
		t.Fatal(err)
	}

	assembly.AddReader("image", "lol_cat.jpg", file)

	assembly.AddStep("resize", map[string]interface{}{
		"robot":           "/image/resize",
		"width":           75,
		"height":          75,
		"resize_strategy": "pad",
		"background":      "#000000",
	})

	assembly.Blocking = true

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.Ok != "ASSEMBLY_COMPLETED" {
		t.Fatal("wrong assembly status")
	}

	if info.AssemblyId == "" {
		t.Fatal("response doesn't contain assembly_id")
	}

	if len(info.Uploads) != 1 {
		t.Fatal("wrong number of uploads")
	}
}

func TestGetAssembly(t *testing.T) {

	client := setup(t)

	assembly, err := client.GetAssembly(assemblyUrl)
	if err != nil {
		t.Fatal(err)
	}

	if assembly.AssemblyId == "" {
		t.Fatal("assembly id not contained")
	}

	if assembly.AssemblyUrl != assemblyUrl {
		t.Fatal("assembly urls don't match")
	}

}

func TestReplayAssembly(t *testing.T) {

	client := setup(t)

	assembly := client.ReplayAssembly(assemblyId)

	assembly.NotifyUrl = "http://requestb.in/1kwp6lx1"
	assembly.ReparseTemplate()

	info, err := assembly.Start()
	if err != nil {
		t.Fatal(err)
	}

	if info.Ok != "ASSEMBLY_REPLAYING" {
		t.Fatal("wrong status code returned")
	}

	if info.NotifyUrl != "http://requestb.in/1kwp6lx1" {
		t.Fatal("wrong notify url")
	}

}

func TestReplayAssemblyBlocking(t *testing.T) {

	client := setup(t)

	assembly := client.ReplayAssembly(assemblyId)

	assembly.Blocking = true

	info, err := assembly.Start()
	if err != nil {
		t.Fatal(err)
	}

	if info.Ok != "ASSEMBLY_COMPLETED" {
		t.Fatal("wrong status code returned")
	}

}

func TestAssemblyUsingTemplate(t *testing.T) {

	client := setup(t)

	assembly := client.CreateAssembly()

	assembly.TemplateId = "64c11b20308811e4b5548d4f316c150f"

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.AssemblyId == "" {
		t.Fatal(fmt.Sprintf("response doesn't contain assembly_id. %s", info.Error))
	}

	if !strings.Contains(info.Params, "64c11b20308811e4b5548d4f316c150f") {
		t.Fatal("template id not as parameter submitted")
	}
}

func TestCancelAssembly(t *testing.T) {

	client := setup(t)

	assembly := client.CreateAssembly()

	assembly.AddStep("import", map[string]interface{}{
		"robot": "/http/import",
		"url":   "http://mirror.nl.leaseweb.net/speedtest/10000mb.bin",
	})

	info, err := assembly.Upload()
	if err != nil {
		t.Fatal(err)
	}

	if info.AssemblyUrl == "" {
		t.Fatal("response doesn't contain assembly_url")
	}

	_, err = client.CancelAssembly(info.AssemblyUrl)
	if err != nil {
		t.Fatal(err)
	}

}

func TestListAssemblies(t *testing.T) {

	client := setup(t)

	assemblies, err := client.ListAssemblies(&ListOptions{
		PageSize: 3,
	})
	if err != nil {
		t.Fatal(err)
	}

	if len(assemblies.Assemblies) < 3 {
		t.Fatal("wrong number of assemblies")
	}

	if assemblies.Count == 0 {
		t.Fatal("wrong count")
	}

	if assemblies.Assemblies[0].AssemblyId == "" {
		t.Fatal("wrong template name")
	}

}
<|endoftext|>


<commit_before>/*Package assert provides assertions for comparing expected values to actual
values. When an assertion fails a helpful error message is printed.

Assert and Check

Assert() and Check() both accept a Comparison, and fail the test when the
comparison fails. The one difference is that Assert() will end the test execution
immediately (using t.FailNow()) whereas Check() will fail the test (using t.Fail()),
return the value of the comparison, then proceed with the rest of the test case.

Example Usage

The example below shows assert used with some common types.

	import (
	    "testing"

	    "github.com/gotestyourself/gotestyourself/assert"
	    is "github.com/gotestyourself/gotestyourself/assert/cmp"
	)

	func TestEverything(t *testing.T) {
	    // booleans
	    assert.Assert(t, ok)
	    assert.Assert(t, !missing)

	    // primitives
	    assert.Equal(t, count, 1)
	    assert.Equal(t, msg, "the message")
	    assert.Assert(t, total != 10) // NotEqual

	    // errors
	    assert.NilError(t, closer.Close())
	    assert.Assert(t, is.Error(err, "the exact error message"))
	    assert.Assert(t, is.ErrorContains(err, "includes this"))
	    assert.Assert(t, os.IsNotExist(err), "got %+v", err)

	    // complex types
	    assert.Assert(t, is.Len(items, 3))
	    assert.Assert(t, len(sequence) != 0) // NotEmpty
	    assert.Assert(t, is.Contains(mapping, "key"))
	    assert.Assert(t, is.DeepEqual(result, myStruct{Name: "title"}))

	    // pointers and interface
	    assert.Assert(t, is.Nil(ref))
	    assert.Assert(t, ref != nil) // NotNil
	}

Comparisons

https://godoc.org/github.com/gotestyourself/gotestyourself/assert/cmp provides
many common comparisons. Additional comparisons can be written to compare
values in other ways.

Below is an example of a custom comparison using a regex pattern:

	func RegexP(value string, pattern string) func() (bool, string) {
	    return func() (bool, string) {
	        re := regexp.MustCompile(pattern)
	        msg := fmt.Sprintf("%q did not match pattern %q", value, pattern)
	        return re.MatchString(value), msg
	    }
	}

*/
package assert

import (
	"fmt"
	"go/ast"
	"go/token"

	"github.com/gotestyourself/gotestyourself/assert/cmp"
	"github.com/gotestyourself/gotestyourself/internal/format"
	"github.com/gotestyourself/gotestyourself/internal/source"
)

// BoolOrComparison can be a bool, or cmp.Comparison. Other types will panic.
type BoolOrComparison interface{}

// TestingT is the subset of testing.T used by the assert package.
type TestingT interface {
	FailNow()
	Fail()
	Log(args ...interface{})
}

type helperT interface {
	Helper()
}

const failureMessage = "assertion failed: "

func assert(
	t TestingT,
	failer func(),
	argsFilter astExprListFilter,
	comparison BoolOrComparison,
	msgAndArgs ...interface{},
) bool {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	var success bool
	switch check := comparison.(type) {
	case bool:
		if check {
			return true
		}
		logFailureFromBool(t, msgAndArgs...)

	// Undocumented legacy comparison without Result type
	case func() (success bool, message string):
		success = runCompareFunc(t, check, msgAndArgs...)

	case cmp.Comparison:
		success = runComparison(t, argsFilter, check, msgAndArgs...)

	case func() cmp.Result:
		success = runComparison(t, argsFilter, check, msgAndArgs...)

	default:
		panic(fmt.Sprintf("comparison arg must be bool or Comparison, not %T", comparison))
	}

	if success {
		return true
	}
	failer()
	return false
}

func runCompareFunc(
	t TestingT,
	f func() (success bool, message string),
	msgAndArgs ...interface{},
) bool {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	if success, message := f(); !success {
		t.Log(format.WithCustomMessage(failureMessage+message, msgAndArgs...))
		return false
	}
	return true
}

func logFailureFromBool(t TestingT, msgAndArgs ...interface{}) {
	const stackIndex = 3 // Assert()/Check(), assert(), formatFailureFromBool()
	const comparisonArgPos = 1
	args, err := source.CallExprArgs(stackIndex)
	if err != nil {
		t.Log(err.Error())
		return
	}

	msg, err := boolFailureMessage(args[comparisonArgPos])
	if err != nil {
		t.Log(err.Error())
		msg = "expression is false"
	}

	t.Log(format.WithCustomMessage(failureMessage+msg, msgAndArgs...))
}

func boolFailureMessage(expr ast.Expr) (string, error) {
	if binaryExpr, ok := expr.(*ast.BinaryExpr); ok && binaryExpr.Op == token.NEQ {
		x, err := source.FormatNode(binaryExpr.X)
		if err != nil {
			return "", err
		}
		y, err := source.FormatNode(binaryExpr.Y)
		if err != nil {
			return "", err
		}
		return x + " is " + y, nil
	}

	if unaryExpr, ok := expr.(*ast.UnaryExpr); ok && unaryExpr.Op == token.NOT {
		x, err := source.FormatNode(unaryExpr.X)
		if err != nil {
			return "", err
		}
		return x + " is true", nil
	}

	formatted, err := source.FormatNode(expr)
	if err != nil {
		return "", err
	}
	return "expression is false: " + formatted, nil
}

// Assert performs a comparison. If the comparison fails the test is marked as
// failed, a failure message is logged, and execution is stopped immediately.
//
// The comparison argument may be one of two types: bool or cmp.Comparison.
// When called with a bool the failure message will contain the literal source
// code of the expression.
// When called with a cmp.Comparison the comparison is responsible for producing
// a helpful failure message.
func Assert(t TestingT, comparison BoolOrComparison, msgAndArgs ...interface{}) {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	assert(t, t.FailNow, filterExprArgsFromComparison, comparison, msgAndArgs...)
}

// Check performs a comparison. If the comparison fails the test is marked as
// failed, a failure message is logged, and Check returns false. Otherwise returns
// true.
//
// See Assert for details about the comparison arg and failure messages.
func Check(t TestingT, comparison BoolOrComparison, msgAndArgs ...interface{}) bool {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	return assert(t, t.Fail, filterExprArgsFromComparison, comparison, msgAndArgs...)
}

// NilError fails the test immediately if the last arg is a non-nil error.
// This is equivalent to Assert(t, cmp.NilError(err)).
func NilError(t TestingT, err error, msgAndArgs ...interface{}) {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	assert(t, t.FailNow, filterExprExcludeFirst, cmp.NilError(err), msgAndArgs...)
}

// Equal uses the == operator to assert two values are equal and fails the test
// if they are not equal. This is equivalent to Assert(t, cmp.Equal(x, y)).
func Equal(t TestingT, x, y interface{}, msgAndArgs ...interface{}) {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	assert(t, t.FailNow, filterExprExcludeFirst, cmp.Equal(x, y), msgAndArgs...)
}
<commit_msg>Support error as a value for Comparison<commit_after>/*Package assert provides assertions for comparing expected values to actual
values. When an assertion fails a helpful error message is printed.

Assert and Check

Assert() and Check() both accept a Comparison, and fail the test when the
comparison fails. The one difference is that Assert() will end the test execution
immediately (using t.FailNow()) whereas Check() will fail the test (using t.Fail()),
return the value of the comparison, then proceed with the rest of the test case.

Example Usage

The example below shows assert used with some common types.

	import (
	    "testing"

	    "github.com/gotestyourself/gotestyourself/assert"
	    is "github.com/gotestyourself/gotestyourself/assert/cmp"
	)

	func TestEverything(t *testing.T) {
	    // booleans
	    assert.Assert(t, ok)
	    assert.Assert(t, !missing)

	    // primitives
	    assert.Equal(t, count, 1)
	    assert.Equal(t, msg, "the message")
	    assert.Assert(t, total != 10) // NotEqual

	    // errors
	    assert.NilError(t, closer.Close())
	    assert.Assert(t, is.Error(err, "the exact error message"))
	    assert.Assert(t, is.ErrorContains(err, "includes this"))
	    assert.Assert(t, os.IsNotExist(err), "got %+v", err)

	    // complex types
	    assert.Assert(t, is.Len(items, 3))
	    assert.Assert(t, len(sequence) != 0) // NotEmpty
	    assert.Assert(t, is.Contains(mapping, "key"))
	    assert.Assert(t, is.DeepEqual(result, myStruct{Name: "title"}))

	    // pointers and interface
	    assert.Assert(t, is.Nil(ref))
	    assert.Assert(t, ref != nil) // NotNil
	}

Comparisons

https://godoc.org/github.com/gotestyourself/gotestyourself/assert/cmp provides
many common comparisons. Additional comparisons can be written to compare
values in other ways.

Below is an example of a custom comparison using a regex pattern:

	func RegexP(value string, pattern string) func() (bool, string) {
	    return func() (bool, string) {
	        re := regexp.MustCompile(pattern)
	        msg := fmt.Sprintf("%q did not match pattern %q", value, pattern)
	        return re.MatchString(value), msg
	    }
	}

*/
package assert

import (
	"fmt"
	"go/ast"
	"go/token"

	"github.com/gotestyourself/gotestyourself/assert/cmp"
	"github.com/gotestyourself/gotestyourself/internal/format"
	"github.com/gotestyourself/gotestyourself/internal/source"
)

// BoolOrComparison can be a bool, or cmp.Comparison. Other types will panic.
type BoolOrComparison interface{}

// TestingT is the subset of testing.T used by the assert package.
type TestingT interface {
	FailNow()
	Fail()
	Log(args ...interface{})
}

type helperT interface {
	Helper()
}

const failureMessage = "assertion failed: "

// nolint: gocyclo
func assert(
	t TestingT,
	failer func(),
	argsFilter astExprListFilter,
	comparison BoolOrComparison,
	msgAndArgs ...interface{},
) bool {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	var success bool
	switch check := comparison.(type) {
	case bool:
		if check {
			return true
		}
		logFailureFromBool(t, msgAndArgs...)

	// Undocumented legacy comparison without Result type
	case func() (success bool, message string):
		success = runCompareFunc(t, check, msgAndArgs...)

	case nil:
		return true

	case error:
		msg := "error is not nil: "
		t.Log(format.WithCustomMessage(failureMessage+msg+check.Error(), msgAndArgs...))

	case cmp.Comparison:
		success = runComparison(t, argsFilter, check, msgAndArgs...)

	case func() cmp.Result:
		success = runComparison(t, argsFilter, check, msgAndArgs...)

	default:
		t.Log(fmt.Sprintf("invalid Comparison: %v (%T)", check, check))
	}

	if success {
		return true
	}
	failer()
	return false
}

func runCompareFunc(
	t TestingT,
	f func() (success bool, message string),
	msgAndArgs ...interface{},
) bool {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	if success, message := f(); !success {
		t.Log(format.WithCustomMessage(failureMessage+message, msgAndArgs...))
		return false
	}
	return true
}

func logFailureFromBool(t TestingT, msgAndArgs ...interface{}) {
	const stackIndex = 3 // Assert()/Check(), assert(), formatFailureFromBool()
	const comparisonArgPos = 1
	args, err := source.CallExprArgs(stackIndex)
	if err != nil {
		t.Log(err.Error())
		return
	}

	msg, err := boolFailureMessage(args[comparisonArgPos])
	if err != nil {
		t.Log(err.Error())
		msg = "expression is false"
	}

	t.Log(format.WithCustomMessage(failureMessage+msg, msgAndArgs...))
}

func boolFailureMessage(expr ast.Expr) (string, error) {
	if binaryExpr, ok := expr.(*ast.BinaryExpr); ok && binaryExpr.Op == token.NEQ {
		x, err := source.FormatNode(binaryExpr.X)
		if err != nil {
			return "", err
		}
		y, err := source.FormatNode(binaryExpr.Y)
		if err != nil {
			return "", err
		}
		return x + " is " + y, nil
	}

	if unaryExpr, ok := expr.(*ast.UnaryExpr); ok && unaryExpr.Op == token.NOT {
		x, err := source.FormatNode(unaryExpr.X)
		if err != nil {
			return "", err
		}
		return x + " is true", nil
	}

	formatted, err := source.FormatNode(expr)
	if err != nil {
		return "", err
	}
	return "expression is false: " + formatted, nil
}

// Assert performs a comparison. If the comparison fails the test is marked as
// failed, a failure message is logged, and execution is stopped immediately.
//
// The comparison argument may be one of two types: bool or cmp.Comparison.
// When called with a bool the failure message will contain the literal source
// code of the expression.
// When called with a cmp.Comparison the comparison is responsible for producing
// a helpful failure message.
func Assert(t TestingT, comparison BoolOrComparison, msgAndArgs ...interface{}) {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	assert(t, t.FailNow, filterExprArgsFromComparison, comparison, msgAndArgs...)
}

// Check performs a comparison. If the comparison fails the test is marked as
// failed, a failure message is logged, and Check returns false. Otherwise returns
// true.
//
// See Assert for details about the comparison arg and failure messages.
func Check(t TestingT, comparison BoolOrComparison, msgAndArgs ...interface{}) bool {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	return assert(t, t.Fail, filterExprArgsFromComparison, comparison, msgAndArgs...)
}

// NilError fails the test immediately if the last arg is a non-nil error.
// This is equivalent to Assert(t, cmp.NilError(err)).
func NilError(t TestingT, err error, msgAndArgs ...interface{}) {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	assert(t, t.FailNow, filterExprExcludeFirst, cmp.NilError(err), msgAndArgs...)
}

// Equal uses the == operator to assert two values are equal and fails the test
// if they are not equal. This is equivalent to Assert(t, cmp.Equal(x, y)).
func Equal(t TestingT, x, y interface{}, msgAndArgs ...interface{}) {
	if ht, ok := t.(helperT); ok {
		ht.Helper()
	}
	assert(t, t.FailNow, filterExprExcludeFirst, cmp.Equal(x, y), msgAndArgs...)
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"
	"strings"

	"github.com/DimShadoWWW/power-pg/common"
	"github.com/op/go-logging"
)

var (
	connid = uint64(0)
)

// Pkg PostgreSQL package structure
type Pkg struct {
	Type    byte
	Content []byte
}

// Start function
func Start(localHost, remoteHost *string, remotePort *string, msgBytes chan []byte, msgCh chan Pkg, recreate bool, log *logging.Logger) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost, remotePort)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
			log:    log,
		}
		go p.start(msgBytes, msgCh, recreate)
	}
}

func getResolvedAddresses(localHost, remoteHost, remotePort *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", *remoteHost, *remotePort))
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
	result        *[]string
	log           *logging.Logger
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(msgBytes chan []byte, msgCh chan Pkg, recreate bool) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, msgBytes, msgCh, recreate, p.log)
	go p.pipe(p.rconn, p.lconn, nil, nil, recreate, p.log)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, msgBytes chan []byte, msgCh chan Pkg, recreate bool, log *logging.Logger) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(ReadBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	// spaces := regexp.MustCompile("[\n\t ]+")
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r ReadBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			if msgBytes != nil {
				log.Debug("Readed bytes: %d\n", n)
			}
			b := buff[:n]
			msgBytes <- b
			//write out result
			if !recreate {
				n, err = dst.Write(b)
				if err != nil {
					p.err("Write failed '%s'\n", err)
					return
				}
			}

			r = b
			log.Debug("PostgreSQL full message: %s\n", string(r))
			// if msgCh != nil {
			// 						msgCh <- 	fmt.Sprintf("%#v", string(buff[:n]))}
			if msgBytes != nil {
				log.Debug("Remaining bytes: %d\n", remainingBytes)
			}
			if msgBytes != nil {
				log.Debug("newPacket : %v\n", newPacket)
			}
			if msgBytes != nil {
				log.Debug("len(r) : %v\n", len(r))
			}
			fmt.Println("3")
			// NewP:
			fmt.Println("4")
			if newPacket || (len(msg) > 4 && len(r) > 4 && remainingBytes == 0) {
				fmt.Println("5")
				// remainingBytes = 0
				newPacket = false
				if msgBytes != nil && msg != "" {
					log.Debug("2 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
				}
				var msg []byte
				t := r.Byte()
				n = n - 1
				fmt.Println("t: ", string(t))
				switch t {
				case 'Q', 'B', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'P', 'p', 'S', 'X':
					// case 'B', 'P':
					// c.rxReadyForQuery(r)
					log.Debug("PostgreSQL pkg type: %s\n", string(t))
					remainingBytes = r.Int32()
					if remainingBytes < 4 {
						fmt.Println("ERROR: remainingBytes can't be less than 4 bytes if int32")
					} else {
						remainingBytes = remainingBytes - 4
						if remainingBytes > 0 {
							if remainingBytes > n {
								remainingBytes = n
							}
							newPacket = true
							msg = append(msg, r.Next(remainingBytes)[:]...)
							// msg = spaces.ReplaceAll(msg, []byte{' '})
							remainingBytes = n - remainingBytes
							if msgCh != nil {
								log.Debug("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							}
							if msgBytes != nil {
								log.Debug("3 Remaining bytes: %d \tmsg: %v\n", remainingBytes, msg)
							}

							if msgCh != nil {
								msgCh <- Pkg{
									Type:    t,
									Content: msg,
								}
							}
							remainingBytes = 0
							// goto NewP
							// } else {
							// 	newPacket = false
							// 	msg = append(msg, r.Next(remainingBytes)[:]...)
							// 	// msg = bytes.Replace(msg, []byte("\n\t"), []byte(" "), -1)
							// 	msg = spaces.ReplaceAll(msg, []byte{' '})
							// 	// msg = []byte(stripchars(string(msg),
							// 	// 	"\n\t"))
							// 	remainingBytes = remainingBytes - n
							// 	if msgBytes != nil {
							// 		log.Debug("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							// 	}
						}
					}
					// case :
					// 	fmt.Println("TODO")
					// 	// c.rxReadyForQuery(r)
					// 	remainingBytes = r.int32()
					// 	remainingBytes = remainingBytes - 4
					// 	if remainingBytes > 0 {
					// 		if remainingBytes <= n {
					// 			newPacket = true
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = n - remainingBytes
					// 			if msgCh != nil {
					// msgCh <- 	fmt.Sprintf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))}
					// 			// fmt.Println(msg)
					// 			goto NewP
					// 		} else {
					// 			newPacket = false
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = remainingBytes - n
					// 			if msgCh != nil {
					// msgCh <- 	fmt.Sprintf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))}
					// 		}
					// 	}
					// case rowDescription:
					// case dataRow:
					// case bindComplete:
					// case commandComplete:
					// 	commandTag = CommandTag(r.readCString())
				// case 'Q', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'p', 'S', 'X':
				default:
					fmt.Println("6")
					remainingBytes = 0
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
				remainingBytes = 0
			} else {
				fmt.Println("7")
				remainingBytes = 0
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// if msgCh != nil {
			// msgCh <- 	fmt.Sprintf("t: %#v\n", t)}

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}

func stripchars(str, chr string) string {
	return strings.Map(func(r rune) rune {
		if strings.IndexRune(chr, r) < 0 {
			return r
		}
		return -1
	}, str)
}
<commit_msg>Update<commit_after>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"
	"strings"

	"github.com/DimShadoWWW/power-pg/common"
	"github.com/op/go-logging"
)

var (
	connid = uint64(0)
)

// Pkg PostgreSQL package structure
type Pkg struct {
	Type    byte
	Content []byte
}

// Start function
func Start(localHost, remoteHost *string, remotePort *string, msgBytes chan []byte, msgCh chan Pkg, recreate bool, log *logging.Logger) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost, remotePort)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
			log:    log,
		}
		go p.start(msgBytes, msgCh, recreate)
	}
}

func getResolvedAddresses(localHost, remoteHost, remotePort *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", *remoteHost, *remotePort))
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
	result        *[]string
	log           *logging.Logger
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(msgBytes chan []byte, msgCh chan Pkg, recreate bool) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, msgBytes, msgCh, recreate, p.log)
	go p.pipe(p.rconn, p.lconn, nil, nil, recreate, p.log)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, msgBytes chan []byte, msgCh chan Pkg, recreate bool, log *logging.Logger) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(ReadBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	// spaces := regexp.MustCompile("[\n\t ]+")
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r ReadBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			if msgBytes != nil {
				log.Debug("Readed bytes: %d\n", n)
			}
			b := buff[:n]
			log.Info("Readed: %v\n", b)
			msgBytes <- b
			//write out result
			if !recreate {
				n, err = dst.Write(b)
				if err != nil {
					p.err("Write failed '%s'\n", err)
					return
				}
			}
			r = buff[:n]
			log.Debug("PostgreSQL full message: %s\n", string(r))
			// if msgCh != nil {
			// 						msgCh <- 	fmt.Sprintf("%#v", string(buff[:n]))}
			if msgBytes != nil {
				log.Debug("Remaining bytes: %d\n", remainingBytes)
			}
			if msgBytes != nil {
				log.Debug("newPacket : %v\n", newPacket)
			}
			if msgBytes != nil {
				log.Debug("len(r) : %v\n", len(r))
			}
			fmt.Println("3")
			// NewP:
			fmt.Println("4")
			if newPacket || (len(msg) > 4 && len(r) > 4 && remainingBytes == 0) {
				fmt.Println("5")
				// remainingBytes = 0
				newPacket = false
				if msgBytes != nil && msg != "" {
					log.Debug("2 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
				}
				var msg []byte
				t := r.Byte()
				n = n - 1
				fmt.Println("t: ", string(t))
				switch t {
				case 'Q', 'B', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'P', 'p', 'S', 'X':
					// case 'B', 'P':
					// c.rxReadyForQuery(r)
					log.Debug("PostgreSQL pkg type: %s\n", string(t))
					remainingBytes = r.Int32()
					if remainingBytes < 4 {
						fmt.Println("ERROR: remainingBytes can't be less than 4 bytes if int32")
					} else {
						remainingBytes = remainingBytes - 4
						if remainingBytes > 0 {
							if remainingBytes > n {
								remainingBytes = n
							}
							newPacket = true
							msg = append(msg, r.Next(remainingBytes)[:]...)
							// msg = spaces.ReplaceAll(msg, []byte{' '})
							remainingBytes = n - remainingBytes
							if msgCh != nil {
								log.Debug("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							}
							if msgBytes != nil {
								log.Debug("3 Remaining bytes: %d \tmsg: %v\n", remainingBytes, msg)
							}

							if msgCh != nil {
								msgCh <- Pkg{
									Type:    t,
									Content: msg,
								}
							}
							remainingBytes = 0
							// goto NewP
							// } else {
							// 	newPacket = false
							// 	msg = append(msg, r.Next(remainingBytes)[:]...)
							// 	// msg = bytes.Replace(msg, []byte("\n\t"), []byte(" "), -1)
							// 	msg = spaces.ReplaceAll(msg, []byte{' '})
							// 	// msg = []byte(stripchars(string(msg),
							// 	// 	"\n\t"))
							// 	remainingBytes = remainingBytes - n
							// 	if msgBytes != nil {
							// 		log.Debug("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							// 	}
						}
					}
					// case :
					// 	fmt.Println("TODO")
					// 	// c.rxReadyForQuery(r)
					// 	remainingBytes = r.int32()
					// 	remainingBytes = remainingBytes - 4
					// 	if remainingBytes > 0 {
					// 		if remainingBytes <= n {
					// 			newPacket = true
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = n - remainingBytes
					// 			if msgCh != nil {
					// msgCh <- 	fmt.Sprintf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))}
					// 			// fmt.Println(msg)
					// 			goto NewP
					// 		} else {
					// 			newPacket = false
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = remainingBytes - n
					// 			if msgCh != nil {
					// msgCh <- 	fmt.Sprintf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))}
					// 		}
					// 	}
					// case rowDescription:
					// case dataRow:
					// case bindComplete:
					// case commandComplete:
					// 	commandTag = CommandTag(r.readCString())
				// case 'Q', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'p', 'S', 'X':
				default:
					fmt.Println("6")
					remainingBytes = 0
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
				remainingBytes = 0
			} else {
				fmt.Println("7")
				remainingBytes = 0
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// if msgCh != nil {
			// msgCh <- 	fmt.Sprintf("t: %#v\n", t)}

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}

func stripchars(str, chr string) string {
	return strings.Map(func(r rune) rune {
		if strings.IndexRune(chr, r) < 0 {
			return r
		}
		return -1
	}, str)
}
<|endoftext|>


<commit_before>package rabbit

import (
	"fmt"
	"os"
	"sync"
	"time"

	log "github.com/cihub/seelog"
	uuid "github.com/nu7hatch/gouuid"
	"github.com/streadway/amqp"
	"gopkg.in/tomb.v2"

	"github.com/obeattie/typhon/errors"
	"github.com/obeattie/typhon/message"
	"github.com/obeattie/typhon/transport"
)

const (
	connectTimeout  = 30 * time.Second
	chanSendTimeout = 10 * time.Second
	respondTimeout  = 10 * time.Second
)

var (
	ErrCouldntConnect   = errors.InternalService("Could not connect to RabbitMQ")
	ErrDeliveriesClosed = errors.InternalService("Delivery channel closed")
	ErrNoReplyTo        = errors.BadRequest("Request does not have appropriate X-Rabbit-ReplyTo header")
)

type rabbitTransport struct {
	tomb          *tomb.Tomb
	connM         sync.RWMutex                       // protects conn + connReady
	conn          *RabbitConnection                  // underlying connection
	connReady     chan struct{}                      // swapped along with conn (reconnecting)
	replyQueue    string                             // message reply queue name
	inflightReqs  map[string]chan<- message.Response // correlation id: response chan
	inflightReqsM sync.Mutex                         // protects inflightReqs
	listeners     map[string]*tomb.Tomb              // service name: tomb
	listenersM    sync.RWMutex                       // protects listeners
}

// run starts the asynchronous run-loop connecting to RabbitMQ
func (t *rabbitTransport) run() {
	initConn := func() *RabbitConnection {
		conn := NewRabbitConnection()
		t.connM.Lock()
		defer t.connM.Unlock()
		t.conn = conn
		select {
		case <-t.connReady:
			// Only swap connReady if it's already closed
			t.connReady = make(chan struct{})
		default:
		}
		return conn
	}
	conn := initConn()

	t.tomb.Go(func() error {
		defer func() {
			t.killListeners()
			conn.Close()
			log.Info("[Typhon:RabbitTransport] Dead; connection closed")
		}()

	runLoop:
		for {
			log.Info("[Typhon:RabbitTransport] Run loop connecting…")
			select {
			case <-t.tomb.Dying():
				return nil

			case <-conn.Init():
				log.Info("[Typhon:RabbitTransport] Run loop connected")
				t.listenReplies()

				select {
				case <-t.tomb.Dying():
					// Do not loop again
					return nil
				default:
					conn.Close()
					conn = initConn()
					continue runLoop
				}

			case <-time.After(connectTimeout):
				log.Criticalf("[Typhon:RabbitTransport] Run loop timed out after %s waiting to connect",
					connectTimeout.String())
				return ErrCouldntConnect
			}
		}
	})
}

// deliveryChan returns the name of a delivery channel for a given service
func (t *rabbitTransport) deliveryChan(serviceName string) string {
	return serviceName
}

func (t *rabbitTransport) Tomb() *tomb.Tomb {
	return t.tomb
}

func (t *rabbitTransport) connection() *RabbitConnection {
	t.connM.RLock()
	defer t.connM.RUnlock()
	return t.conn
}

func (t *rabbitTransport) Ready() <-chan struct{} {
	t.connM.RLock()
	defer t.connM.RUnlock()
	return t.connReady
}

func (t *rabbitTransport) killListeners() {
	t.listenersM.RLock()
	ts := make([]*tomb.Tomb, 0, len(t.listeners))
	for _, t := range t.listeners {
		t.Killf("Listeners killed")
		ts = append(ts, t)
	}
	t.listenersM.RUnlock()
	for _, t := range ts {
		t.Wait()
	}
}

func (t *rabbitTransport) StopListening(serviceName string) bool {
	t.listenersM.RLock()
	tm, ok := t.listeners[t.deliveryChan(serviceName)]
	if ok {
		tm.Killf("Stopped listening")
	}
	t.listenersM.RUnlock()
	if ok {
		tm.Wait()
		return true
	}
	return false
}

func (t *rabbitTransport) Listen(serviceName string, rc chan<- message.Request) error {
	tm := &tomb.Tomb{}
	cn := t.deliveryChan(serviceName)
	t.listenersM.Lock()
	if _, ok := t.listeners[cn]; ok {
		t.listenersM.Unlock()
		return transport.ErrAlreadyListening
	} else {
		t.listeners[cn] = tm
		t.listenersM.Unlock()
	}

	// Used to convey a connection error to the caller (we block until listening has begun)
	errChan := make(chan error)

	tm.Go(func() error {
		timeout := time.NewTimer(connectTimeout)
		defer func() {
			timeout.Stop()
			t.listenersM.Lock()
			defer t.listenersM.Unlock()
			delete(t.listeners, cn)
			close(rc)
			close(errChan)
			log.Debugf("[Typhon:RabbitTransport] Listener %s stopped", cn)
		}()

		select {
		case <-t.tomb.Dying():
			return nil
		case <-tm.Dying():
			return nil
		case <-timeout.C:
			errChan <- transport.ErrTimeout
			return nil
		case <-t.Ready():
		}

		deliveryChan, rabbitChannel, err := t.connection().Consume(cn)
		if err != nil {
			log.Warnf("[Typhon:RabbitTransport] Failed to consume from %s: %v", cn, err)
			errChan <- err
			return nil
		}
		defer rabbitChannel.Close()
		errChan <- nil
		log.Infof("[Typhon:RabbitTransport] Listening on %s…", cn)

		for {
			select {
			case <-t.tomb.Dying():
				return nil

			case <-tm.Dying():
				return nil

			case delivery, ok := <-deliveryChan:
				if !ok {
					log.Warnf("[Typhon:RabbitTransport] Delivery channel closed; stopping listener %s", cn)
					return nil
				}
				go t.handleReqDelivery(delivery, rc)
			}
		}
	})
	return <-errChan
}

func (t *rabbitTransport) logId(delivery amqp.Delivery) string {
	return fmt.Sprintf("%s[%s]", delivery.RoutingKey, delivery.CorrelationId)
}

func (t *rabbitTransport) Respond(req message.Request, rsp message.Response) error {
	headers := rsp.Headers()
	headers["Content-Encoding"] = "response"
	headers["Service"] = rsp.Service()
	headers["Endpoint"] = rsp.Endpoint()

	timeout := time.NewTimer(respondTimeout)
	defer timeout.Stop()
	select {
	case <-t.Ready():
		timeout.Stop()
	case <-t.tomb.Dying():
		return tomb.ErrDying
	case <-timeout.C:
		return transport.ErrTimeout
	}

	return t.connection().Publish("", req.Headers()["X-Rabbit-ReplyTo"], amqp.Publishing{
		CorrelationId: rsp.Id(),
		Timestamp:     time.Now().UTC(),
		Body:          rsp.Payload(),
		Headers:       headersToTable(headers),
	})
}

func (t *rabbitTransport) Send(req message.Request, _timeout time.Duration) (message.Response, error) {
	id := req.Id()
	if id == "" {
		_uuid, err := uuid.NewV4()
		if err != nil {
			log.Errorf("[Typhon:RabbitTransport] Failed to generate request uuid: %v", err)
			return nil, err
		}
		req.SetId(_uuid.String())
	}

	rspQueue := req.Id()
	defer func() {
		t.inflightReqsM.Lock()
		delete(t.inflightReqs, rspQueue)
		t.inflightReqsM.Unlock()
	}()
	rspChan := make(chan message.Response, 1)
	t.inflightReqsM.Lock()
	t.inflightReqs[rspQueue] = rspChan
	t.inflightReqsM.Unlock()

	timeout := time.NewTimer(_timeout)
	defer timeout.Stop()

	headers := req.Headers()
	headers["Content-Encoding"] = "request"
	headers["Service"] = req.Service()
	headers["Endpoint"] = req.Endpoint()

	select {
	case <-t.Ready():
	case <-timeout.C:
		log.Warnf("[Typhon:RabbitTransport] Timed out after %s waiting for ready", _timeout.String())
		return nil, transport.ErrTimeout
	}

	log.Tracef("[Typhon:RabbitTransport] Sending request %s…", req.Id())
	if err := t.connection().Publish(Exchange, req.Service(), amqp.Publishing{
		CorrelationId: req.Id(),
		Timestamp:     time.Now().UTC(),
		Body:          req.Payload(),
		ReplyTo:       t.replyQueue,
		Headers:       headersToTable(headers),
	}); err != nil {
		log.Errorf("[Typhon:RabbitTransport] Failed to publish: %v", err)
		return nil, err
	}

	select {
	case rsp := <-rspChan:
		return rsp, nil
	case <-timeout.C:
		log.Warnf("[Typhon:RabbitTransport] Timed out after %s waiting for response to %s", _timeout.String(),
			req.Id())
		return nil, transport.ErrTimeout
	}
}

func (t *rabbitTransport) listenReplies() error {
	if err := t.connection().Channel.DeclareReplyQueue(t.replyQueue); err != nil {
		log.Criticalf("[Typhon:RabbitTransport] Failed to declare reply queue %s: %v", t.replyQueue, err)
		os.Exit(1)
		return err
	}

	deliveries, err := t.connection().Channel.ConsumeQueue(t.replyQueue)
	if err != nil {
		log.Criticalf("[Typhon:RabbitTransport] Failed to consume from reply queue %s: %v", t.replyQueue, err)
		os.Exit(1)
		return err
	}

	log.Debugf("[Typhon:RabbitTransport] Listening for replies on %s…", t.replyQueue)
	t.connM.RLock()
	readyC := t.connReady
	t.connM.RUnlock()
	select {
	case <-readyC:
		// Make sure not to close the channel if it's already closed
	default:
		close(readyC)
	}

	for {
		select {
		case delivery, ok := <-deliveries:
			if !ok {
				log.Warnf("[Typhon:RabbitTransport] Delivery channel %s closed", t.replyQueue)
				return ErrDeliveriesClosed
			}
			go t.handleRspDelivery(delivery)

		case <-t.tomb.Dying():
			log.Info("[Typhon:RabbitTransport] Reply listener terminating (tomb death)")
			return tomb.ErrDying
		}
	}
}

func (t *rabbitTransport) deliveryToMessage(delivery amqp.Delivery, msg message.Message) {
	msg.SetId(delivery.CorrelationId)
	msg.SetHeaders(tableToHeaders(delivery.Headers))
	msg.SetHeader("X-Rabbit-ReplyTo", delivery.ReplyTo)
	msg.SetPayload(delivery.Body)
	switch service := delivery.Headers["Service"].(type) {
	case string:
		msg.SetService(service)
	}
	switch endpoint := delivery.Headers["Endpoint"].(type) {
	case string:
		msg.SetEndpoint(endpoint)
	}
}

func (t *rabbitTransport) handleReqDelivery(delivery amqp.Delivery, reqChan chan<- message.Request) {
	logId := t.logId(delivery)
	enc := delivery.Headers["Content-Encoding"].(string)
	switch enc {
	case "request":
		req := message.NewRequest()
		t.deliveryToMessage(delivery, req)

		timeout := time.NewTimer(chanSendTimeout)
		defer timeout.Stop()
		select {
		case reqChan <- req:
		case <-timeout.C:
			log.Errorf("[Typhon:RabbitTransport] Could not deliver request %s after %s: receiving channel is full",
				logId, chanSendTimeout.String())
		}

	default:
		log.Debugf("[Typhon:RabbitTransport] Cannot handle Content-Encoding \"%s\" as request for %s", enc, logId)
	}
}

func (t *rabbitTransport) handleRspDelivery(delivery amqp.Delivery) {
	logId := t.logId(delivery)

	enc := delivery.Headers["Content-Encoding"].(string)
	switch enc {
	case "response":
		rsp := message.NewResponse()
		t.deliveryToMessage(delivery, rsp)

		t.inflightReqsM.Lock()
		rspChan, ok := t.inflightReqs[rsp.Id()]
		delete(t.inflightReqs, rsp.Id())
		t.inflightReqsM.Unlock()
		if !ok {
			log.Warnf("[Typhon:RabbitTransport] Could not match response %s to channel", logId)
			return
		}

		timeout := time.NewTimer(chanSendTimeout)
		defer timeout.Stop()
		select {
		case rspChan <- rsp:
		case <-timeout.C:
			log.Errorf("[Typhon:RabbitTransport] Could not deliver responde %s after %s: receiving channel is full",
				logId, chanSendTimeout.String())
		}

	default:
		log.Errorf("[Typhon:RabbitTransport] Cannot handle Content-Encoding \"%s\" as response for %s", enc, logId)
	}
}

func NewTransport() transport.Transport {
	rt := &rabbitTransport{
		tomb:         new(tomb.Tomb),
		inflightReqs: make(map[string]chan<- message.Response),
		listeners:    make(map[string]*tomb.Tomb),
		connReady:    make(chan struct{}),
	}

	if uid, err := uuid.NewV4(); err != nil {
		log.Criticalf("[Typhon:RabbitTransport] Failed to create UUID for reply queue: %v", err)
		os.Exit(1) // TODO: Is this really necessary?
	} else {
		rt.replyQueue = fmt.Sprintf("replyTo-%s", uid.String())
	}
	rt.run()
	return rt
}
<commit_msg>Remove annoying logging<commit_after>package rabbit

import (
	"fmt"
	"os"
	"sync"
	"time"

	log "github.com/cihub/seelog"
	uuid "github.com/nu7hatch/gouuid"
	"github.com/streadway/amqp"
	"gopkg.in/tomb.v2"

	"github.com/obeattie/typhon/errors"
	"github.com/obeattie/typhon/message"
	"github.com/obeattie/typhon/transport"
)

const (
	connectTimeout  = 30 * time.Second
	chanSendTimeout = 10 * time.Second
	respondTimeout  = 10 * time.Second
)

var (
	ErrCouldntConnect   = errors.InternalService("Could not connect to RabbitMQ")
	ErrDeliveriesClosed = errors.InternalService("Delivery channel closed")
	ErrNoReplyTo        = errors.BadRequest("Request does not have appropriate X-Rabbit-ReplyTo header")
)

type rabbitTransport struct {
	tomb          *tomb.Tomb
	connM         sync.RWMutex                       // protects conn + connReady
	conn          *RabbitConnection                  // underlying connection
	connReady     chan struct{}                      // swapped along with conn (reconnecting)
	replyQueue    string                             // message reply queue name
	inflightReqs  map[string]chan<- message.Response // correlation id: response chan
	inflightReqsM sync.Mutex                         // protects inflightReqs
	listeners     map[string]*tomb.Tomb              // service name: tomb
	listenersM    sync.RWMutex                       // protects listeners
}

// run starts the asynchronous run-loop connecting to RabbitMQ
func (t *rabbitTransport) run() {
	initConn := func() *RabbitConnection {
		conn := NewRabbitConnection()
		t.connM.Lock()
		defer t.connM.Unlock()
		t.conn = conn
		select {
		case <-t.connReady:
			// Only swap connReady if it's already closed
			t.connReady = make(chan struct{})
		default:
		}
		return conn
	}
	conn := initConn()

	t.tomb.Go(func() error {
		defer func() {
			t.killListeners()
			conn.Close()
			log.Info("[Typhon:RabbitTransport] Dead; connection closed")
		}()

	runLoop:
		for {
			log.Info("[Typhon:RabbitTransport] Run loop connecting…")
			select {
			case <-t.tomb.Dying():
				return nil

			case <-conn.Init():
				log.Info("[Typhon:RabbitTransport] Run loop connected")
				t.listenReplies()

				select {
				case <-t.tomb.Dying():
					// Do not loop again
					return nil
				default:
					conn.Close()
					conn = initConn()
					continue runLoop
				}

			case <-time.After(connectTimeout):
				log.Criticalf("[Typhon:RabbitTransport] Run loop timed out after %s waiting to connect",
					connectTimeout.String())
				return ErrCouldntConnect
			}
		}
	})
}

// deliveryChan returns the name of a delivery channel for a given service
func (t *rabbitTransport) deliveryChan(serviceName string) string {
	return serviceName
}

func (t *rabbitTransport) Tomb() *tomb.Tomb {
	return t.tomb
}

func (t *rabbitTransport) connection() *RabbitConnection {
	t.connM.RLock()
	defer t.connM.RUnlock()
	return t.conn
}

func (t *rabbitTransport) Ready() <-chan struct{} {
	t.connM.RLock()
	defer t.connM.RUnlock()
	return t.connReady
}

func (t *rabbitTransport) killListeners() {
	t.listenersM.RLock()
	ts := make([]*tomb.Tomb, 0, len(t.listeners))
	for _, t := range t.listeners {
		t.Killf("Listeners killed")
		ts = append(ts, t)
	}
	t.listenersM.RUnlock()
	for _, t := range ts {
		t.Wait()
	}
}

func (t *rabbitTransport) StopListening(serviceName string) bool {
	t.listenersM.RLock()
	tm, ok := t.listeners[t.deliveryChan(serviceName)]
	if ok {
		tm.Killf("Stopped listening")
	}
	t.listenersM.RUnlock()
	if ok {
		tm.Wait()
		return true
	}
	return false
}

func (t *rabbitTransport) Listen(serviceName string, rc chan<- message.Request) error {
	tm := &tomb.Tomb{}
	cn := t.deliveryChan(serviceName)
	t.listenersM.Lock()
	if _, ok := t.listeners[cn]; ok {
		t.listenersM.Unlock()
		return transport.ErrAlreadyListening
	} else {
		t.listeners[cn] = tm
		t.listenersM.Unlock()
	}

	// Used to convey a connection error to the caller (we block until listening has begun)
	errChan := make(chan error)

	tm.Go(func() error {
		timeout := time.NewTimer(connectTimeout)
		defer func() {
			timeout.Stop()
			t.listenersM.Lock()
			defer t.listenersM.Unlock()
			delete(t.listeners, cn)
			close(rc)
			close(errChan)
			log.Debugf("[Typhon:RabbitTransport] Listener %s stopped", cn)
		}()

		select {
		case <-t.tomb.Dying():
			return nil
		case <-tm.Dying():
			return nil
		case <-timeout.C:
			errChan <- transport.ErrTimeout
			return nil
		case <-t.Ready():
		}

		deliveryChan, rabbitChannel, err := t.connection().Consume(cn)
		if err != nil {
			log.Warnf("[Typhon:RabbitTransport] Failed to consume from %s: %v", cn, err)
			errChan <- err
			return nil
		}
		defer rabbitChannel.Close()
		errChan <- nil
		log.Infof("[Typhon:RabbitTransport] Listening on %s…", cn)

		for {
			select {
			case <-t.tomb.Dying():
				return nil

			case <-tm.Dying():
				return nil

			case delivery, ok := <-deliveryChan:
				if !ok {
					log.Warnf("[Typhon:RabbitTransport] Delivery channel closed; stopping listener %s", cn)
					return nil
				}
				go t.handleReqDelivery(delivery, rc)
			}
		}
	})
	return <-errChan
}

func (t *rabbitTransport) logId(delivery amqp.Delivery) string {
	return fmt.Sprintf("%s[%s]", delivery.RoutingKey, delivery.CorrelationId)
}

func (t *rabbitTransport) Respond(req message.Request, rsp message.Response) error {
	headers := rsp.Headers()
	headers["Content-Encoding"] = "response"
	headers["Service"] = rsp.Service()
	headers["Endpoint"] = rsp.Endpoint()

	timeout := time.NewTimer(respondTimeout)
	defer timeout.Stop()
	select {
	case <-t.Ready():
		timeout.Stop()
	case <-t.tomb.Dying():
		return tomb.ErrDying
	case <-timeout.C:
		return transport.ErrTimeout
	}

	return t.connection().Publish("", req.Headers()["X-Rabbit-ReplyTo"], amqp.Publishing{
		CorrelationId: rsp.Id(),
		Timestamp:     time.Now().UTC(),
		Body:          rsp.Payload(),
		Headers:       headersToTable(headers),
	})
}

func (t *rabbitTransport) Send(req message.Request, _timeout time.Duration) (message.Response, error) {
	id := req.Id()
	if id == "" {
		_uuid, err := uuid.NewV4()
		if err != nil {
			log.Errorf("[Typhon:RabbitTransport] Failed to generate request uuid: %v", err)
			return nil, err
		}
		req.SetId(_uuid.String())
	}

	rspQueue := req.Id()
	defer func() {
		t.inflightReqsM.Lock()
		delete(t.inflightReqs, rspQueue)
		t.inflightReqsM.Unlock()
	}()
	rspChan := make(chan message.Response, 1)
	t.inflightReqsM.Lock()
	t.inflightReqs[rspQueue] = rspChan
	t.inflightReqsM.Unlock()

	timeout := time.NewTimer(_timeout)
	defer timeout.Stop()

	headers := req.Headers()
	headers["Content-Encoding"] = "request"
	headers["Service"] = req.Service()
	headers["Endpoint"] = req.Endpoint()

	select {
	case <-t.Ready():
	case <-timeout.C:
		log.Warnf("[Typhon:RabbitTransport] Timed out after %s waiting for ready", _timeout.String())
		return nil, transport.ErrTimeout
	}

	if err := t.connection().Publish(Exchange, req.Service(), amqp.Publishing{
		CorrelationId: req.Id(),
		Timestamp:     time.Now().UTC(),
		Body:          req.Payload(),
		ReplyTo:       t.replyQueue,
		Headers:       headersToTable(headers),
	}); err != nil {
		log.Errorf("[Typhon:RabbitTransport] Failed to publish: %v", err)
		return nil, err
	}

	select {
	case rsp := <-rspChan:
		return rsp, nil
	case <-timeout.C:
		log.Warnf("[Typhon:RabbitTransport] Timed out after %s waiting for response to %s", _timeout.String(),
			req.Id())
		return nil, transport.ErrTimeout
	}
}

func (t *rabbitTransport) listenReplies() error {
	if err := t.connection().Channel.DeclareReplyQueue(t.replyQueue); err != nil {
		log.Criticalf("[Typhon:RabbitTransport] Failed to declare reply queue %s: %v", t.replyQueue, err)
		os.Exit(1)
		return err
	}

	deliveries, err := t.connection().Channel.ConsumeQueue(t.replyQueue)
	if err != nil {
		log.Criticalf("[Typhon:RabbitTransport] Failed to consume from reply queue %s: %v", t.replyQueue, err)
		os.Exit(1)
		return err
	}

	log.Debugf("[Typhon:RabbitTransport] Listening for replies on %s…", t.replyQueue)
	t.connM.RLock()
	readyC := t.connReady
	t.connM.RUnlock()
	select {
	case <-readyC:
		// Make sure not to close the channel if it's already closed
	default:
		close(readyC)
	}

	for {
		select {
		case delivery, ok := <-deliveries:
			if !ok {
				log.Warnf("[Typhon:RabbitTransport] Delivery channel %s closed", t.replyQueue)
				return ErrDeliveriesClosed
			}
			go t.handleRspDelivery(delivery)

		case <-t.tomb.Dying():
			log.Info("[Typhon:RabbitTransport] Reply listener terminating (tomb death)")
			return tomb.ErrDying
		}
	}
}

func (t *rabbitTransport) deliveryToMessage(delivery amqp.Delivery, msg message.Message) {
	msg.SetId(delivery.CorrelationId)
	msg.SetHeaders(tableToHeaders(delivery.Headers))
	msg.SetHeader("X-Rabbit-ReplyTo", delivery.ReplyTo)
	msg.SetPayload(delivery.Body)
	switch service := delivery.Headers["Service"].(type) {
	case string:
		msg.SetService(service)
	}
	switch endpoint := delivery.Headers["Endpoint"].(type) {
	case string:
		msg.SetEndpoint(endpoint)
	}
}

func (t *rabbitTransport) handleReqDelivery(delivery amqp.Delivery, reqChan chan<- message.Request) {
	logId := t.logId(delivery)
	enc := delivery.Headers["Content-Encoding"].(string)
	switch enc {
	case "request":
		req := message.NewRequest()
		t.deliveryToMessage(delivery, req)

		timeout := time.NewTimer(chanSendTimeout)
		defer timeout.Stop()
		select {
		case reqChan <- req:
		case <-timeout.C:
			log.Errorf("[Typhon:RabbitTransport] Could not deliver request %s after %s: receiving channel is full",
				logId, chanSendTimeout.String())
		}

	default:
		log.Debugf("[Typhon:RabbitTransport] Cannot handle Content-Encoding \"%s\" as request for %s", enc, logId)
	}
}

func (t *rabbitTransport) handleRspDelivery(delivery amqp.Delivery) {
	logId := t.logId(delivery)

	enc := delivery.Headers["Content-Encoding"].(string)
	switch enc {
	case "response":
		rsp := message.NewResponse()
		t.deliveryToMessage(delivery, rsp)

		t.inflightReqsM.Lock()
		rspChan, ok := t.inflightReqs[rsp.Id()]
		delete(t.inflightReqs, rsp.Id())
		t.inflightReqsM.Unlock()
		if !ok {
			log.Warnf("[Typhon:RabbitTransport] Could not match response %s to channel", logId)
			return
		}

		timeout := time.NewTimer(chanSendTimeout)
		defer timeout.Stop()
		select {
		case rspChan <- rsp:
		case <-timeout.C:
			log.Errorf("[Typhon:RabbitTransport] Could not deliver responde %s after %s: receiving channel is full",
				logId, chanSendTimeout.String())
		}

	default:
		log.Errorf("[Typhon:RabbitTransport] Cannot handle Content-Encoding \"%s\" as response for %s", enc, logId)
	}
}

func NewTransport() transport.Transport {
	rt := &rabbitTransport{
		tomb:         new(tomb.Tomb),
		inflightReqs: make(map[string]chan<- message.Response),
		listeners:    make(map[string]*tomb.Tomb),
		connReady:    make(chan struct{}),
	}

	if uid, err := uuid.NewV4(); err != nil {
		log.Criticalf("[Typhon:RabbitTransport] Failed to create UUID for reply queue: %v", err)
		os.Exit(1) // TODO: Is this really necessary?
	} else {
		rt.replyQueue = fmt.Sprintf("replyTo-%s", uid.String())
	}
	rt.run()
	return rt
}
<|endoftext|>


<commit_before>package main

import (
	//"github.com/jmoiron/sqlx"
	//_ "github.com/lib/pq"
	"time"
)

func classifyGoods() {
	defer time.Sleep(time.Second)

	type GroupId int
	type Hash string
	type GoodId int64

	// load existing group infos
	type GroupInfo struct {
		GroupId GroupId     `db:"group_id"`
		Hashes  StringArray `db:"hashes"`
		Title   string      `db:"title"`
	}
	var groupRows []GroupInfo
	ce(db.Select(&groupRows, `SELECT
		group_id, hashes, title
		FROM groups`), "select group infos")
	groupIdToHashSet := make(map[GroupId]map[Hash]struct{})
	hashToGroupIdSet := make(map[Hash]map[GroupId]struct{})
	for _, info := range groupRows {
		for _, hash := range info.Hashes {
			hash := Hash(hash)
			if _, ok := groupIdToHashSet[info.GroupId]; !ok {
				groupIdToHashSet[info.GroupId] = make(map[Hash]struct{})
			}
			groupIdToHashSet[info.GroupId][hash] = struct{}{}
			if _, ok := hashToGroupIdSet[hash]; !ok {
				hashToGroupIdSet[hash] = make(map[GroupId]struct{})
			}
			hashToGroupIdSet[hash][info.GroupId] = struct{}{}
		}
	}
	pt("group infos loaded\n")

select_goods:
	tx := db.MustBegin()

	var goodIds Int64Array
	err := tx.Select(&goodIds, `SELECT good_id
			FROM goods
			WHERE added_at >= $1
			--AND status > 0
			AND group_id IS NULL
			ORDER BY good_id DESC
			LIMIT 256
			`,
		time.Now().Add(-time.Hour*24*120).Format("2006-01-02"),
	)
	ce(err, "select goods")
	pt("select %d goods\n", len(goodIds))
	if len(goodIds) == 0 {
		return
	}

	//XXX debug
	//var err error
	//goodIds := Int64Array{
	//	2398924,
	//}

	//XXX from vvic db
	//lalaDB, err := sqlx.Connect("postgres", "host=reus.mobi user=reus dbname=lala connect_timeout=60")
	//ce(err, "connect to lala db")
	//var goodIdsFromLala Int64Array
	//ce(lalaDB.Select(&goodIdsFromLala, `SELECT vvic_id FROM items`), "select good ids from lala db")
	//var goodIds Int64Array
	//ce(db.Select(&goodIds, `SELECT good_id FROM goods
	//	WHERE good_id = ANY($1)
	//	AND group_id IS NULL
	//	`,
	//	goodIdsFromLala,
	//), "select good ids")
	//pt("%d good ids\n", len(goodIds))
	//if len(goodIds) == 0 {
	//	return
	//}

	var infos []struct {
		GoodId GoodId `db:"good_id"`
		Hash   string `db:"hash"`
	}
	err = tx.Select(&infos, `SELECT i.good_id, encode(sha512_16k, 'hex') AS hash
		FROM images i
		LEFT JOIN urls u ON u.url_id = i.url_id
		WHERE i.good_id = ANY($1)
		AND sha512_16k IS NOT NULL
		`,
		goodIds,
	)
	ce(err, "select hashes")
	goodIdToHashSet := make(map[GoodId]map[Hash]struct{})
	for _, info := range infos {
		if _, ok := goodIdToHashSet[info.GoodId]; !ok {
			goodIdToHashSet[info.GoodId] = make(map[Hash]struct{})
		}
		goodIdToHashSet[info.GoodId][Hash(info.Hash)] = struct{}{}
	}
	pt("select %d rows of infos\n", len(infos))

loop_goods:
	for _, goodId := range goodIds {
		goodId := GoodId(goodId)
		candidateGroupIdSet := make(map[GroupId]struct{})
		for hash := range goodIdToHashSet[goodId] {
			if groupIdSet, ok := hashToGroupIdSet[hash]; ok {
				for groupId := range groupIdSet {
					candidateGroupIdSet[groupId] = struct{}{}
				}
			}
		}

		for groupId := range candidateGroupIdSet {
			count := 0
			total := 0
			for hash := range groupIdToHashSet[groupId] {
				if _, ok := goodIdToHashSet[goodId][hash]; ok {
					count++
				}
				total++
			}
			if total-count <= 3 && count > 5 {
				var similarity float64
				ce(tx.Get(&similarity, `SELECT similarity(
					(SELECT title FROM groups WHERE group_id = $1),
					(SELECT title FROM goods WHERE good_id = $2)
				)`,
					groupId,
					goodId,
				), "get similarity")
				if similarity < 0.3 && count < 15 {
					// 有15个图相同的话，就不管标题了
					// 如果发现有商品的小图用量大于15,那就提高
					continue
				}
				//same group
				_, err := tx.Exec(`UPDATE goods
							SET group_id = $1
							WHERE good_id = $2
							`,
					groupId,
					goodId,
				)
				ce(err, "update goods")
				continue loop_goods
			}
		}

		// new group
		var groupHashes StringArray
		for hash := range goodIdToHashSet[goodId] {
			groupHashes = append(groupHashes, string(hash))
		}
		row := tx.QueryRow(`INSERT INTO groups
				(hashes, title)
				VALUES ($1, (SELECT title FROM goods WHERE good_id = $2))
				RETURNING group_id
				`,
			groupHashes,
			goodId,
		)
		ce(err, "insert new group")
		var newGroupId GroupId
		ce(row.Scan(&newGroupId), "scan")
		for hash := range goodIdToHashSet[goodId] {
			if _, ok := groupIdToHashSet[newGroupId]; !ok {
				groupIdToHashSet[newGroupId] = make(map[Hash]struct{})
			}
			groupIdToHashSet[newGroupId][hash] = struct{}{}
			if _, ok := hashToGroupIdSet[hash]; !ok {
				hashToGroupIdSet[hash] = make(map[GroupId]struct{})
			}
			hashToGroupIdSet[hash][newGroupId] = struct{}{}
		}
		_, err := tx.Exec(`UPDATE goods
				SET group_id = $1
				WHERE good_id = $2
				`,
			newGroupId,
			goodId,
		)
		ce(err, "update goods")

	}

	ce(tx.Commit(), "commit")

	goto select_goods

}
<commit_msg>use base64 to reduce memory consumption<commit_after>package main

import (
	//"github.com/jmoiron/sqlx"
	//_ "github.com/lib/pq"
	"time"
)

func classifyGoods() {
	defer time.Sleep(time.Second)

	type GroupId int
	type Hash string
	type GoodId int64

	// load existing group infos
	type GroupInfo struct {
		GroupId GroupId     `db:"group_id"`
		Hashes  StringArray `db:"hashes"`
		Title   string      `db:"title"`
	}
	var groupRows []GroupInfo
	ce(db.Select(&groupRows, `SELECT
		group_id, hashes, title
		FROM groups`), "select group infos")
	groupIdToHashSet := make(map[GroupId]map[Hash]struct{})
	hashToGroupIdSet := make(map[Hash]map[GroupId]struct{})
	for _, info := range groupRows {
		for _, hash := range info.Hashes {
			hash := Hash(hash)
			if _, ok := groupIdToHashSet[info.GroupId]; !ok {
				groupIdToHashSet[info.GroupId] = make(map[Hash]struct{})
			}
			groupIdToHashSet[info.GroupId][hash] = struct{}{}
			if _, ok := hashToGroupIdSet[hash]; !ok {
				hashToGroupIdSet[hash] = make(map[GroupId]struct{})
			}
			hashToGroupIdSet[hash][info.GroupId] = struct{}{}
		}
	}
	pt("group infos loaded\n")

select_goods:
	tx := db.MustBegin()

	var goodIds Int64Array
	err := tx.Select(&goodIds, `SELECT good_id
			FROM goods
			WHERE added_at >= $1
			--AND status > 0
			AND group_id IS NULL
			ORDER BY good_id DESC
			LIMIT 256
			`,
		time.Now().Add(-time.Hour*24*120).Format("2006-01-02"),
	)
	ce(err, "select goods")
	pt("select %d goods\n", len(goodIds))
	if len(goodIds) == 0 {
		return
	}

	//XXX debug
	//var err error
	//goodIds := Int64Array{
	//	2398924,
	//}

	//XXX from vvic db
	//lalaDB, err := sqlx.Connect("postgres", "host=reus.mobi user=reus dbname=lala connect_timeout=60")
	//ce(err, "connect to lala db")
	//var goodIdsFromLala Int64Array
	//ce(lalaDB.Select(&goodIdsFromLala, `SELECT vvic_id FROM items`), "select good ids from lala db")
	//var goodIds Int64Array
	//ce(db.Select(&goodIds, `SELECT good_id FROM goods
	//	WHERE good_id = ANY($1)
	//	AND group_id IS NULL
	//	`,
	//	goodIdsFromLala,
	//), "select good ids")
	//pt("%d good ids\n", len(goodIds))
	//if len(goodIds) == 0 {
	//	return
	//}

	var infos []struct {
		GoodId GoodId `db:"good_id"`
		Hash   string `db:"hash"`
	}
	err = tx.Select(&infos, `SELECT i.good_id, encode(sha512_16k, 'base64') AS hash
		FROM images i
		LEFT JOIN urls u ON u.url_id = i.url_id
		WHERE i.good_id = ANY($1)
		AND sha512_16k IS NOT NULL
		`,
		goodIds,
	)
	ce(err, "select hashes")
	goodIdToHashSet := make(map[GoodId]map[Hash]struct{})
	for _, info := range infos {
		if _, ok := goodIdToHashSet[info.GoodId]; !ok {
			goodIdToHashSet[info.GoodId] = make(map[Hash]struct{})
		}
		goodIdToHashSet[info.GoodId][Hash(info.Hash)] = struct{}{}
	}
	pt("select %d rows of infos\n", len(infos))

loop_goods:
	for _, goodId := range goodIds {
		goodId := GoodId(goodId)
		candidateGroupIdSet := make(map[GroupId]struct{})
		for hash := range goodIdToHashSet[goodId] {
			if groupIdSet, ok := hashToGroupIdSet[hash]; ok {
				for groupId := range groupIdSet {
					candidateGroupIdSet[groupId] = struct{}{}
				}
			}
		}

		for groupId := range candidateGroupIdSet {
			count := 0
			total := 0
			for hash := range groupIdToHashSet[groupId] {
				if _, ok := goodIdToHashSet[goodId][hash]; ok {
					count++
				}
				total++
			}
			if total-count <= 3 && count > 5 {
				var similarity float64
				ce(tx.Get(&similarity, `SELECT similarity(
					(SELECT title FROM groups WHERE group_id = $1),
					(SELECT title FROM goods WHERE good_id = $2)
				)`,
					groupId,
					goodId,
				), "get similarity")
				if similarity < 0.3 && count < 15 {
					// 有15个图相同的话，就不管标题了
					// 如果发现有商品的小图用量大于15,那就提高
					continue
				}
				//same group
				_, err := tx.Exec(`UPDATE goods
							SET group_id = $1
							WHERE good_id = $2
							`,
					groupId,
					goodId,
				)
				ce(err, "update goods")
				continue loop_goods
			}
		}

		// new group
		var groupHashes StringArray
		for hash := range goodIdToHashSet[goodId] {
			groupHashes = append(groupHashes, string(hash))
		}
		row := tx.QueryRow(`INSERT INTO groups
				(hashes, title)
				VALUES ($1, (SELECT title FROM goods WHERE good_id = $2))
				RETURNING group_id
				`,
			groupHashes,
			goodId,
		)
		ce(err, "insert new group")
		var newGroupId GroupId
		ce(row.Scan(&newGroupId), "scan")
		for hash := range goodIdToHashSet[goodId] {
			if _, ok := groupIdToHashSet[newGroupId]; !ok {
				groupIdToHashSet[newGroupId] = make(map[Hash]struct{})
			}
			groupIdToHashSet[newGroupId][hash] = struct{}{}
			if _, ok := hashToGroupIdSet[hash]; !ok {
				hashToGroupIdSet[hash] = make(map[GroupId]struct{})
			}
			hashToGroupIdSet[hash][newGroupId] = struct{}{}
		}
		_, err := tx.Exec(`UPDATE goods
				SET group_id = $1
				WHERE good_id = $2
				`,
			newGroupId,
			goodId,
		)
		ce(err, "update goods")

	}

	ce(tx.Commit(), "commit")

	goto select_goods

}
<|endoftext|>


<commit_before>package pt

import (
	"fmt"
	"image"
	"math"
	"math/rand"
	"runtime"
	"strings"
	"sync"
	"time"
)

type Renderer struct {
	Scene              *Scene
	Camera             *Camera
	Sampler            Sampler
	Buffer             *Buffer
	SamplesPerPixel    int
	StratifiedSampling bool
	AdaptiveSamples    int
	FireflySamples     int
	FireflyThreshold   float64
	NumCPU             int
	Verbose            bool
}

func NewRenderer(scene *Scene, camera *Camera, sampler Sampler, w, h int) *Renderer {
	r := Renderer{}
	r.Scene = scene
	r.Camera = camera
	r.Sampler = sampler
	r.Buffer = NewBuffer(w, h)
	r.SamplesPerPixel = 1
	r.StratifiedSampling = false
	r.AdaptiveSamples = 0
	r.FireflySamples = 0
	r.FireflyThreshold = 1
	r.NumCPU = runtime.NumCPU()
	r.Verbose = true
	return &r
}

func (r *Renderer) run() {
	scene := r.Scene
	camera := r.Camera
	sampler := r.Sampler
	buf := r.Buffer
	w, h := buf.W, buf.H
	spp := r.SamplesPerPixel
	sppRoot := int(math.Sqrt(float64(r.SamplesPerPixel)))
	ncpu := r.NumCPU

	runtime.GOMAXPROCS(ncpu)
	scene.Compile()
	ch := make(chan int, h)
	r.printf("%d x %d pixels, %d spp, %d cores\n", w, h, spp, ncpu)
	start := time.Now()
	scene.rays = 0
	for i := 0; i < ncpu; i++ {
		go func(i int) {
			rnd := rand.New(rand.NewSource(time.Now().UnixNano()))
			for y := i; y < h; y += ncpu {
				for x := 0; x < w; x++ {
					if r.StratifiedSampling {
						// stratified subsampling
						for u := 0; u < sppRoot; u++ {
							for v := 0; v < sppRoot; v++ {
								fu := (float64(u) + 0.5) / float64(sppRoot)
								fv := (float64(v) + 0.5) / float64(sppRoot)
								ray := camera.CastRay(x, y, w, h, fu, fv, rnd)
								sample := sampler.Sample(scene, ray, rnd)
								buf.AddSample(x, y, sample)
							}
						}
					} else {
						// random subsampling
						for i := 0; i < spp; i++ {
							fu := rnd.Float64()
							fv := rnd.Float64()
							ray := camera.CastRay(x, y, w, h, fu, fv, rnd)
							sample := sampler.Sample(scene, ray, rnd)
							buf.AddSample(x, y, sample)
						}
					}
					// adaptive sampling
					if r.AdaptiveSamples > 0 {
						v := Clamp(buf.StandardDeviation(x, y).MaxComponent(), 0, 1)
						// v = math.Pow(v, 2)
						samples := int(v * float64(r.AdaptiveSamples))
						for i := 0; i < samples; i++ {
							fu := rnd.Float64()
							fv := rnd.Float64()
							ray := camera.CastRay(x, y, w, h, fu, fv, rnd)
							sample := sampler.Sample(scene, ray, rnd)
							buf.AddSample(x, y, sample)
						}
					}
					// firefly reduction
					if r.FireflySamples > 0 {
						if buf.StandardDeviation(x, y).MaxComponent() > r.FireflyThreshold {
							for i := 0; i < r.FireflySamples; i++ {
								fu := rnd.Float64()
								fv := rnd.Float64()
								ray := camera.CastRay(x, y, w, h, fu, fv, rnd)
								sample := sampler.Sample(scene, ray, rnd)
								buf.AddSample(x, y, sample)
							}
						}
					}
				}
				ch <- 1
			}
		}(i)
	}
	r.showProgress(start, scene.RayCount(), 0, h)
	for i := 0; i < h; i++ {
		<-ch
		r.showProgress(start, scene.RayCount(), i+1, h)
	}
	r.printf("\n")
}

func (r *Renderer) printf(format string, a ...interface{}) {
	if !r.Verbose {
		return
	}
	fmt.Printf(format, a...)
}

func (r *Renderer) showProgress(start time.Time, rays uint64, i, h int) {
	if !r.Verbose {
		return
	}
	pct := int(100 * float64(i) / float64(h))
	elapsed := time.Since(start)
	rps := float64(rays) / elapsed.Seconds()
	fmt.Printf("\r%4d / %d (%3d%%) [", i, h, pct)
	for p := 0; p < 100; p += 3 {
		if pct > p {
			fmt.Print("=")
		} else {
			fmt.Print(" ")
		}
	}
	fmt.Printf("] %s %s ", DurationString(elapsed), NumberString(rps))
}

func (r *Renderer) writeImage(path string, buf *Buffer, channel Channel, wg *sync.WaitGroup) {
	defer wg.Done()
	im := buf.Image(channel)
	if err := SavePNG(path, im); err != nil {
		panic(err)
	}
}

func (r *Renderer) Render() image.Image {
	r.run()
	return r.Buffer.Image(ColorChannel)
}

func (r *Renderer) IterativeRender(pathTemplate string, iterations int) image.Image {
	var wg sync.WaitGroup
	for i := 1; i <= iterations; i++ {
		r.printf("\n[Iteration %d of %d]\n", i, iterations)
		r.run()
		path := pathTemplate
		if strings.Contains(path, "%") {
			path = fmt.Sprintf(pathTemplate, i)
		}
		buf := r.Buffer.Copy()
		wg.Add(1)
		go r.writeImage(path, buf, ColorChannel, &wg)
		// wg.Add(1)
		// go r.writeImage("deviation.png", buf, StandardDeviationChannel, &wg)
	}
	wg.Wait()
	return r.Buffer.Image(ColorChannel)
}

func (r *Renderer) ChannelRender() <-chan image.Image {
	ch := make(chan image.Image)
	go func() {
		for i := 1; ; i++ {
			r.run()
			ch <- r.Buffer.Image(ColorChannel)
		}
	}()
	return ch
}

func (r *Renderer) FrameRender(path string, iterations int, wg *sync.WaitGroup) {
	for i := 1; i <= iterations; i++ {
		r.run()
	}
	buf := r.Buffer.Copy()
	wg.Add(1)
	go r.writeImage(path, buf, ColorChannel, wg)
}

func (r *Renderer) TimedRender(duration time.Duration) image.Image {
	start := time.Now()
	for {
		r.run()
		if time.Since(start) > duration {
			break
		}
	}
	return r.Buffer.Image(ColorChannel)
}
<commit_msg>Adaptive threshold and exponent<commit_after>package pt

import (
	"fmt"
	"image"
	"math"
	"math/rand"
	"runtime"
	"strings"
	"sync"
	"time"
)

type Renderer struct {
	Scene              *Scene
	Camera             *Camera
	Sampler            Sampler
	Buffer             *Buffer
	SamplesPerPixel    int
	StratifiedSampling bool
	AdaptiveSamples    int
	AdaptiveThreshold  float64
	AdaptiveExponent   float64
	FireflySamples     int
	FireflyThreshold   float64
	NumCPU             int
	Verbose            bool
}

func NewRenderer(scene *Scene, camera *Camera, sampler Sampler, w, h int) *Renderer {
	r := Renderer{}
	r.Scene = scene
	r.Camera = camera
	r.Sampler = sampler
	r.Buffer = NewBuffer(w, h)
	r.SamplesPerPixel = 1
	r.StratifiedSampling = false
	r.AdaptiveSamples = 0
	r.AdaptiveThreshold = 1
	r.AdaptiveExponent = 1
	r.FireflySamples = 0
	r.FireflyThreshold = 1
	r.NumCPU = runtime.NumCPU()
	r.Verbose = true
	return &r
}

func (r *Renderer) run() {
	scene := r.Scene
	camera := r.Camera
	sampler := r.Sampler
	buf := r.Buffer
	w, h := buf.W, buf.H
	spp := r.SamplesPerPixel
	sppRoot := int(math.Sqrt(float64(r.SamplesPerPixel)))
	ncpu := r.NumCPU

	runtime.GOMAXPROCS(ncpu)
	scene.Compile()
	ch := make(chan int, h)
	r.printf("%d x %d pixels, %d spp, %d cores\n", w, h, spp, ncpu)
	start := time.Now()
	scene.rays = 0
	for i := 0; i < ncpu; i++ {
		go func(i int) {
			rnd := rand.New(rand.NewSource(time.Now().UnixNano()))
			for y := i; y < h; y += ncpu {
				for x := 0; x < w; x++ {
					if r.StratifiedSampling {
						// stratified subsampling
						for u := 0; u < sppRoot; u++ {
							for v := 0; v < sppRoot; v++ {
								fu := (float64(u) + 0.5) / float64(sppRoot)
								fv := (float64(v) + 0.5) / float64(sppRoot)
								ray := camera.CastRay(x, y, w, h, fu, fv, rnd)
								sample := sampler.Sample(scene, ray, rnd)
								buf.AddSample(x, y, sample)
							}
						}
					} else {
						// random subsampling
						for i := 0; i < spp; i++ {
							fu := rnd.Float64()
							fv := rnd.Float64()
							ray := camera.CastRay(x, y, w, h, fu, fv, rnd)
							sample := sampler.Sample(scene, ray, rnd)
							buf.AddSample(x, y, sample)
						}
					}
					// adaptive sampling
					if r.AdaptiveSamples > 0 {
						v := buf.StandardDeviation(x, y).MaxComponent()
						v = Clamp(v/r.AdaptiveThreshold, 0, 1)
						v = math.Pow(v, r.AdaptiveExponent)
						samples := int(v * float64(r.AdaptiveSamples))
						for i := 0; i < samples; i++ {
							fu := rnd.Float64()
							fv := rnd.Float64()
							ray := camera.CastRay(x, y, w, h, fu, fv, rnd)
							sample := sampler.Sample(scene, ray, rnd)
							buf.AddSample(x, y, sample)
						}
					}
					// firefly reduction
					if r.FireflySamples > 0 {
						if buf.StandardDeviation(x, y).MaxComponent() > r.FireflyThreshold {
							for i := 0; i < r.FireflySamples; i++ {
								fu := rnd.Float64()
								fv := rnd.Float64()
								ray := camera.CastRay(x, y, w, h, fu, fv, rnd)
								sample := sampler.Sample(scene, ray, rnd)
								buf.AddSample(x, y, sample)
							}
						}
					}
				}
				ch <- 1
			}
		}(i)
	}
	r.showProgress(start, scene.RayCount(), 0, h)
	for i := 0; i < h; i++ {
		<-ch
		r.showProgress(start, scene.RayCount(), i+1, h)
	}
	r.printf("\n")
}

func (r *Renderer) printf(format string, a ...interface{}) {
	if !r.Verbose {
		return
	}
	fmt.Printf(format, a...)
}

func (r *Renderer) showProgress(start time.Time, rays uint64, i, h int) {
	if !r.Verbose {
		return
	}
	pct := int(100 * float64(i) / float64(h))
	elapsed := time.Since(start)
	rps := float64(rays) / elapsed.Seconds()
	fmt.Printf("\r%4d / %d (%3d%%) [", i, h, pct)
	for p := 0; p < 100; p += 3 {
		if pct > p {
			fmt.Print("=")
		} else {
			fmt.Print(" ")
		}
	}
	fmt.Printf("] %s %s ", DurationString(elapsed), NumberString(rps))
}

func (r *Renderer) writeImage(path string, buf *Buffer, channel Channel, wg *sync.WaitGroup) {
	defer wg.Done()
	im := buf.Image(channel)
	if err := SavePNG(path, im); err != nil {
		panic(err)
	}
}

func (r *Renderer) Render() image.Image {
	r.run()
	return r.Buffer.Image(ColorChannel)
}

func (r *Renderer) IterativeRender(pathTemplate string, iterations int) image.Image {
	var wg sync.WaitGroup
	for i := 1; i <= iterations; i++ {
		r.printf("\n[Iteration %d of %d]\n", i, iterations)
		r.run()
		path := pathTemplate
		if strings.Contains(path, "%") {
			path = fmt.Sprintf(pathTemplate, i)
		}
		buf := r.Buffer.Copy()
		wg.Add(1)
		go r.writeImage(path, buf, ColorChannel, &wg)
		// wg.Add(1)
		// go r.writeImage("deviation.png", buf, StandardDeviationChannel, &wg)
	}
	wg.Wait()
	return r.Buffer.Image(ColorChannel)
}

func (r *Renderer) ChannelRender() <-chan image.Image {
	ch := make(chan image.Image)
	go func() {
		for i := 1; ; i++ {
			r.run()
			ch <- r.Buffer.Image(ColorChannel)
		}
	}()
	return ch
}

func (r *Renderer) FrameRender(path string, iterations int, wg *sync.WaitGroup) {
	for i := 1; i <= iterations; i++ {
		r.run()
	}
	buf := r.Buffer.Copy()
	wg.Add(1)
	go r.writeImage(path, buf, ColorChannel, wg)
}

func (r *Renderer) TimedRender(duration time.Duration) image.Image {
	start := time.Now()
	for {
		r.run()
		if time.Since(start) > duration {
			break
		}
	}
	return r.Buffer.Image(ColorChannel)
}
<|endoftext|>


<commit_before>package etf

/*
Copyright © 2012 Serge Zirukin

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

import (
  "github.com/bmizerany/assert"
  "math/big"
  "testing"
)

func Test_Decode_BigInt(t *testing.T) {
  var bigint *big.Int

  size, err := Decode([]byte{131,110,15,0,0,0,0,0,16,
    159,75,179,21,7,201,123,206,151,192,
  }, &bigint)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(19), size)
}

func Test_Decode_Binary(t *testing.T) {
  var data []interface{}

  size, err := Decode([]byte{131,109,0,0,0,3,1,2,3}, &data)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(9), size)
  assert.Equal(t, byte(1), data[0])
  assert.Equal(t, byte(2), data[1])
  assert.Equal(t, byte(3), data[2])
}

func Test_Decode(t *testing.T) {
  var s string
  size, err := Decode([]byte{131,107,0,3,49,50,51}, &s)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(7), size)
  assert.Equal(t, "123", s)

  type testStruct struct {
    Atom
    X uint8
    S string
  }

  var ts testStruct

  size, err = Decode([]byte{
    131,104,3,100,0,4,98,108,97,104,97,4,108,0,0,0,4,98,
    0,0,4,68,98,0,0,4,75,98,0,0,4,50,98,0,0,4,48,106,
  }, &ts)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(38), size)
  assert.Equal(t, uint8(4), ts.X)
  assert.Equal(t, "фыва", ts.S)

  size, err = Decode([]byte{
    131,104,3,99,50,46,57,57,57,57,57,57,57,57,57,57,57,57,57,
    57,57,56,56,56,57,56,101,45,48,49,0,0,0,0,0,97,4,108,0,0,
    0,4,98,0,0,4,68,98,0,0,4,75,98,0,0,4,50,98,0,0,4,48,106,
  }, &ts)
  assert.NotEqual(t, nil, err)

  type testStruct2 struct {
    T testStruct
    Y int
  }

  var ts2 testStruct2

  size, err = Decode([]byte{
    131,104,2,104,3,100,0,4,98,108,97,104,97,4,108,0,0,0,4,98,
    0,0,4,68,98,0,0,4,75,98,0,0,4,50,98,0,0,4,48,106,98,0,0,2,154,
  }, &ts2)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(45), size)
  assert.Equal(t, uint8(4), ts2.T.X)
  assert.Equal(t, "фыва", ts2.T.S)
  assert.Equal(t, 666, ts2.Y)
}

// Local Variables:
// indent-tabs-mode: nil
// tab-width: 2
// End:
// ex: set tabstop=2 shiftwidth=2 expandtab:
<commit_msg>decode: fix test<commit_after>package etf

/*
Copyright © 2012 Serge Zirukin

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

import (
  "github.com/bmizerany/assert"
  "math/big"
  "testing"
)

func Test_Decode_BigInt(t *testing.T) {
  var bigint *big.Int

  size, err := Decode([]byte{131,110,15,0,0,0,0,0,16,
    159,75,179,21,7,201,123,206,151,192,
  }, &bigint)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(19), size)
}

func Test_Decode_Binary(t *testing.T) {
  var data []byte

  size, err := Decode([]byte{131,109,0,0,0,3,1,2,3}, &data)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(9), size)
  assert.Equal(t, byte(1), data[0])
  assert.Equal(t, byte(2), data[1])
  assert.Equal(t, byte(3), data[2])
}

func Test_Decode(t *testing.T) {
  var s string
  size, err := Decode([]byte{131,107,0,3,49,50,51}, &s)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(7), size)
  assert.Equal(t, "123", s)

  type testStruct struct {
    Atom
    X uint8
    S string
  }

  var ts testStruct

  size, err = Decode([]byte{
    131,104,3,100,0,4,98,108,97,104,97,4,108,0,0,0,4,98,
    0,0,4,68,98,0,0,4,75,98,0,0,4,50,98,0,0,4,48,106,
  }, &ts)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(38), size)
  assert.Equal(t, uint8(4), ts.X)
  assert.Equal(t, "фыва", ts.S)

  size, err = Decode([]byte{
    131,104,3,99,50,46,57,57,57,57,57,57,57,57,57,57,57,57,57,
    57,57,56,56,56,57,56,101,45,48,49,0,0,0,0,0,97,4,108,0,0,
    0,4,98,0,0,4,68,98,0,0,4,75,98,0,0,4,50,98,0,0,4,48,106,
  }, &ts)
  assert.NotEqual(t, nil, err)

  type testStruct2 struct {
    T testStruct
    Y int
  }

  var ts2 testStruct2

  size, err = Decode([]byte{
    131,104,2,104,3,100,0,4,98,108,97,104,97,4,108,0,0,0,4,98,
    0,0,4,68,98,0,0,4,75,98,0,0,4,50,98,0,0,4,48,106,98,0,0,2,154,
  }, &ts2)
  assert.Equal(t, nil, err)
  assert.Equal(t, uint(45), size)
  assert.Equal(t, uint8(4), ts2.T.X)
  assert.Equal(t, "фыва", ts2.T.S)
  assert.Equal(t, 666, ts2.Y)
}

// Local Variables:
// indent-tabs-mode: nil
// tab-width: 2
// End:
// ex: set tabstop=2 shiftwidth=2 expandtab:
<|endoftext|>