<commit_before>package xslt

import (
	"fmt"
	"math"
	"unsafe"

	"github.com/jbowtie/gokogiri/xml"
	"github.com/jbowtie/gokogiri/xpath"
)

func (style *Stylesheet) RegisterXsltFunctions() {
	// id and lang are built into libxml2, don't need to be registered
	style.Functions["{}document"] = XsltDocumentFn
	style.Functions["{}generate-id"] = XsltGenerateId
	style.Functions["{}key"] = XsltKey
	style.Functions["{}system-property"] = XsltSystemProperty
	style.Functions["{}unparsed-entity-uri"] = XsltUnparsedEntityUri
	style.Functions["{}current"] = XsltCurrent
	style.Functions["{}element-available"] = XsltElementAvailable
	style.Functions["{}function-available"] = XsltFunctionAvailable
	//format-number - requires handling decimal-format

	style.Functions["{http://xmlsoft.org/XSLT/namespace}node-set"] = EXSLTnodeset
	style.Functions["{http://exslt.org/common}node-set"] = EXSLTnodeset
	style.Functions["{http://exslt.org/math}constant"] = EXSLTmathconstant
	style.Functions["{http://exslt.org/math}sin"] = EXSLTmathsin
	style.Functions["{http://exslt.org/math}cos"] = EXSLTmathcos
}

type Key struct {
	nodes map[string]xml.Nodeset
	use   string
	match string
}

// Implementation of key() from XSLT spec
func XsltKey(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 2 {
		return nil
	}
	// always convert to string
	name := args[0].(string)
	// convert to string (TODO: unless nodeset)
	val := ""
	switch v := args[1].(type) {
	case string:
		val = v
	case []unsafe.Pointer:
		// nodeset; see xsl:key spec for how to handle this
	}
	//get the execution context
	c := context.(*ExecutionContext)
	//look up the key
	k, ok := c.Style.Keys[name]
	if !ok {
		return nil
	}
	result, _ := k.nodes[val]
	//return the nodeset
	return result.ToPointers()
}

// Implementation of system-property() from XSLT spec
func XsltSystemProperty(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	switch args[0].(string) {
	case "xsl:version":
		return 1.0
	case "xsl:vendor":
		return "John C Barstow"
	case "xsl:vendor-url":
		return "http://github.com/jbowtie/ratago"
	default:
		fmt.Println("EXEC system-property", args[0])
	}
	return nil
}

//Implementation of document() from XSLT spec
func XsltDocumentFn(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)

	switch doc := args[0].(type) {
	case string:
		if doc == "" {
			nodeset := xml.Nodeset{c.Style.Doc}
			return nodeset.ToPointers()
		}
		input := c.FetchInputDocument(doc, false)
		if input != nil {
			nodeset := xml.Nodeset{input}
			return nodeset.ToPointers()
		}
		return nil
	case []unsafe.Pointer:
		n := xml.NewNode(doc[0], nil)
		location := n.Content()
		input := c.FetchInputDocument(location, true)
		if input != nil {
			nodeset := xml.Nodeset{input}
			return nodeset.ToPointers()
		}
		fmt.Println("DOCUMENT", location)
	}
	return nil
}

// Implementation of generate-id() from XSLT spec
func XsltGenerateId(context xpath.VariableScope, args []interface{}) interface{} {
	// should be 0 or 1 argument
	if len(args) > 1 {
		return nil
	}

	//c := context.(*ExecutionContext)
	if len(args) < 1 {
		fmt.Println("GENERATE-ID for current")
		return "N" //id of context node
	}

	switch v := args[0].(type) {
	case []unsafe.Pointer:
		if len(v) == 0 {
			return nil
		}
		out := fmt.Sprintf("N%v", uintptr(v[0]))
		return out
	default:
		return nil
	}
}

// Implementation of unparsed-entity-uri() from XSLT spec
func XsltUnparsedEntityUri(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)
	name := argValToString(args[0])
	val := c.Source.UnparsedEntityURI(name)
	return val
}

// Implementation of current() from XSLT spec
func XsltCurrent(context xpath.VariableScope, args []interface{}) interface{} {
	c := context.(*ExecutionContext)
	n := xml.Nodeset{c.Current}
	return n.ToPointers()
}

// Implementation of function-available() from XSLT spec
func XsltFunctionAvailable(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)
	qname := args[0].(string)
	//TODO: resolve namespace of argument
	return c.IsFunctionRegistered("", qname)
}

// Implementation of element-available() from XSLT spec
func XsltElementAvailable(context xpath.VariableScope, args []interface{}) interface{} {
	return false
}

// util function because we can't assume we're actually getting a string
func argValToString(val interface{}) (out string) {
	if val == nil {
		return
	}
	switch v := val.(type) {
	case string:
		return v
	case []unsafe.Pointer:
		if len(v) == 0 {
			return
		}
		n := xml.NewNode(v[0], nil)
		out = n.Content()
	default:
		out = fmt.Sprintf("%v", v)
	}
	return
}

func EXSLTnodeset(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)
	nodes := args[0]
	switch v := nodes.(type) {
	case []unsafe.Pointer:
		if len(v) == 0 {
			return nil
		}
		fauxroot := c.Output.CreateElementNode("VARIABLE")
		for _, node := range v {
			n := xml.NewNode(node, nil)
			fauxroot.AddChild(n)
		}
		out := xml.Nodeset{fauxroot}
		return out.ToPointers()
	default:
		out := fmt.Sprintf("%v", v)
		fmt.Println("invalid argument to exslt:nodeset", out)
	}

	return nodes
}

func EXSLTmathconstant(context xpath.VariableScope, args []interface{}) interface{} {

	if len(args) != 2 {
		return 0
	}

	name := args[0]
	precision := int(args[1].(float64))

	switch name {
	case "PI":
		return fmt.Sprintf("%.*f", precision, 3.1415926535897932384626433832795028841971693993751)
	case "E":
		return fmt.Sprintf("%.*f", precision, 2.71828182845904523536028747135266249775724709369996)
	case "SQRRT2":
		return fmt.Sprintf("%.*f", precision, 1.41421356237309504880168872420969807856967187537694)
	case "LN2":
		return fmt.Sprintf("%.*f", precision, 0.69314718055994530941723212145817656807550013436025)
	case "LN10":
		return fmt.Sprintf("%.*f", precision, 2.30258509299404568402)
	case "LOG2E":
		return fmt.Sprintf("%.*f", precision, 1.4426950408889634074)
	case "SQRT1_2":
		return fmt.Sprintf("%.*f", precision, 0.70710678118654752440)
	}

	return 0
}

func EXSLTmathsin(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) != 1 {
		return nil
	}

	return math.Sin(args[0].(float64))
}

func EXSLTmathcos(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) != 1 {
		return nil
	}

	return math.Cos(args[0].(float64))
}
<commit_msg>Print an error message for unsupported constants<commit_after>package xslt

import (
	"fmt"
	"math"
	"unsafe"

	"github.com/jbowtie/gokogiri/xml"
	"github.com/jbowtie/gokogiri/xpath"
)

func (style *Stylesheet) RegisterXsltFunctions() {
	// id and lang are built into libxml2, don't need to be registered
	style.Functions["{}document"] = XsltDocumentFn
	style.Functions["{}generate-id"] = XsltGenerateId
	style.Functions["{}key"] = XsltKey
	style.Functions["{}system-property"] = XsltSystemProperty
	style.Functions["{}unparsed-entity-uri"] = XsltUnparsedEntityUri
	style.Functions["{}current"] = XsltCurrent
	style.Functions["{}element-available"] = XsltElementAvailable
	style.Functions["{}function-available"] = XsltFunctionAvailable
	//format-number - requires handling decimal-format

	style.Functions["{http://xmlsoft.org/XSLT/namespace}node-set"] = EXSLTnodeset
	style.Functions["{http://exslt.org/common}node-set"] = EXSLTnodeset
	style.Functions["{http://exslt.org/math}constant"] = EXSLTmathconstant
	style.Functions["{http://exslt.org/math}sin"] = EXSLTmathsin
	style.Functions["{http://exslt.org/math}cos"] = EXSLTmathcos
}

type Key struct {
	nodes map[string]xml.Nodeset
	use   string
	match string
}

// Implementation of key() from XSLT spec
func XsltKey(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 2 {
		return nil
	}
	// always convert to string
	name := args[0].(string)
	// convert to string (TODO: unless nodeset)
	val := ""
	switch v := args[1].(type) {
	case string:
		val = v
	case []unsafe.Pointer:
		// nodeset; see xsl:key spec for how to handle this
	}
	//get the execution context
	c := context.(*ExecutionContext)
	//look up the key
	k, ok := c.Style.Keys[name]
	if !ok {
		return nil
	}
	result, _ := k.nodes[val]
	//return the nodeset
	return result.ToPointers()
}

// Implementation of system-property() from XSLT spec
func XsltSystemProperty(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	switch args[0].(string) {
	case "xsl:version":
		return 1.0
	case "xsl:vendor":
		return "John C Barstow"
	case "xsl:vendor-url":
		return "http://github.com/jbowtie/ratago"
	default:
		fmt.Println("EXEC system-property", args[0])
	}
	return nil
}

//Implementation of document() from XSLT spec
func XsltDocumentFn(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)

	switch doc := args[0].(type) {
	case string:
		if doc == "" {
			nodeset := xml.Nodeset{c.Style.Doc}
			return nodeset.ToPointers()
		}
		input := c.FetchInputDocument(doc, false)
		if input != nil {
			nodeset := xml.Nodeset{input}
			return nodeset.ToPointers()
		}
		return nil
	case []unsafe.Pointer:
		n := xml.NewNode(doc[0], nil)
		location := n.Content()
		input := c.FetchInputDocument(location, true)
		if input != nil {
			nodeset := xml.Nodeset{input}
			return nodeset.ToPointers()
		}
		fmt.Println("DOCUMENT", location)
	}
	return nil
}

// Implementation of generate-id() from XSLT spec
func XsltGenerateId(context xpath.VariableScope, args []interface{}) interface{} {
	// should be 0 or 1 argument
	if len(args) > 1 {
		return nil
	}

	//c := context.(*ExecutionContext)
	if len(args) < 1 {
		fmt.Println("GENERATE-ID for current")
		return "N" //id of context node
	}

	switch v := args[0].(type) {
	case []unsafe.Pointer:
		if len(v) == 0 {
			return nil
		}
		out := fmt.Sprintf("N%v", uintptr(v[0]))
		return out
	default:
		return nil
	}
}

// Implementation of unparsed-entity-uri() from XSLT spec
func XsltUnparsedEntityUri(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)
	name := argValToString(args[0])
	val := c.Source.UnparsedEntityURI(name)
	return val
}

// Implementation of current() from XSLT spec
func XsltCurrent(context xpath.VariableScope, args []interface{}) interface{} {
	c := context.(*ExecutionContext)
	n := xml.Nodeset{c.Current}
	return n.ToPointers()
}

// Implementation of function-available() from XSLT spec
func XsltFunctionAvailable(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)
	qname := args[0].(string)
	//TODO: resolve namespace of argument
	return c.IsFunctionRegistered("", qname)
}

// Implementation of element-available() from XSLT spec
func XsltElementAvailable(context xpath.VariableScope, args []interface{}) interface{} {
	return false
}

// util function because we can't assume we're actually getting a string
func argValToString(val interface{}) (out string) {
	if val == nil {
		return
	}
	switch v := val.(type) {
	case string:
		return v
	case []unsafe.Pointer:
		if len(v) == 0 {
			return
		}
		n := xml.NewNode(v[0], nil)
		out = n.Content()
	default:
		out = fmt.Sprintf("%v", v)
	}
	return
}

func EXSLTnodeset(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)
	nodes := args[0]
	switch v := nodes.(type) {
	case []unsafe.Pointer:
		if len(v) == 0 {
			return nil
		}
		fauxroot := c.Output.CreateElementNode("VARIABLE")
		for _, node := range v {
			n := xml.NewNode(node, nil)
			fauxroot.AddChild(n)
		}
		out := xml.Nodeset{fauxroot}
		return out.ToPointers()
	default:
		out := fmt.Sprintf("%v", v)
		fmt.Println("invalid argument to exslt:nodeset", out)
	}

	return nodes
}

func EXSLTmathconstant(context xpath.VariableScope, args []interface{}) interface{} {

	if len(args) != 2 {
		return 0
	}

	name := args[0].(string)
	precision := int(args[1].(float64))

	switch name {
	case "PI":
		return fmt.Sprintf("%.*f", precision, 3.1415926535897932384626433832795028841971693993751)
	case "E":
		return fmt.Sprintf("%.*f", precision, 2.71828182845904523536028747135266249775724709369996)
	case "SQRRT2":
		return fmt.Sprintf("%.*f", precision, 1.41421356237309504880168872420969807856967187537694)
	case "LN2":
		return fmt.Sprintf("%.*f", precision, 0.69314718055994530941723212145817656807550013436025)
	case "LN10":
		return fmt.Sprintf("%.*f", precision, 2.30258509299404568402)
	case "LOG2E":
		return fmt.Sprintf("%.*f", precision, 1.4426950408889634074)
	case "SQRT1_2":
		return fmt.Sprintf("%.*f", precision, 0.70710678118654752440)
	default:
		out := fmt.Sprintf("%v", name)
		fmt.Println("unsupported constant in math:constant", out)
	}

	return 0
}

func EXSLTmathsin(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) != 1 {
		return nil
	}

	return math.Sin(args[0].(float64))
}

func EXSLTmathcos(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) != 1 {
		return nil
	}

	return math.Cos(args[0].(float64))
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"testing"
	"time"
)

const (
	ANY_DB_DRIVER = ""
	TEST_ENDPOINT = "localhost:9999"
	TEST_DB       = "goslow_test"
)

var DATA_SOURCE = map[string]string{
	"sqlite3":  ":memory:",
	"postgres": "postgres://localhost/" + TEST_DB,
}

type TestCase struct {
	createDefaultRules bool
	adminUrlPathPrefix string
	driver             string
	dataSource         string
}

type TestCases []*TestCase

type CheckFunc func(*testing.T, *httptest.Server, *TestCase)

var (
	defaultTestCases = TestCases{
		NewTestCase(true, "", ANY_DB_DRIVER),
	}

	ruleCreationTestCases = TestCases{
		NewTestCase(true, "", ANY_DB_DRIVER),
		NewTestCase(false, "/goslow", ANY_DB_DRIVER),
		NewTestCase(false, "/goslow/", ANY_DB_DRIVER),
		NewTestCase(false, "/te", ANY_DB_DRIVER),
		NewTestCase(false, "/te/", ANY_DB_DRIVER),
		NewTestCase(false, "/composite/path", ANY_DB_DRIVER),
	}
)

func NewTestCase(createDefaultRules bool, adminUrlPathPrefix string, driver string) *TestCase {
	dataSource, knownDriver := DATA_SOURCE[driver]
	if driver != ANY_DB_DRIVER && !knownDriver {
		log.Fatalf("unknown driver: <%s>", driver)
	}
	return &TestCase{
		createDefaultRules: createDefaultRules,
		adminUrlPathPrefix: adminUrlPathPrefix,
		driver:             driver,
		dataSource:         dataSource,
	}
}

func TestZeroSite(t *testing.T) {
	runAll(t, checkZeroSite, defaultTestCases)
}

func runAll(t *testing.T, checkFunc CheckFunc, testCases TestCases) {
	for _, testCase := range runnable(all(testCases)) {
		run(t, checkFunc, testCase)
	}
}

func runnable(testCases TestCases) TestCases {
	runnableTestCases := make(TestCases, 0)
	for _, testCase := range testCases {
		if !testCase.skippable() {
			runnableTestCases = append(runnableTestCases, testCase)
		}
	}
	return runnableTestCases
}

func all(testCases TestCases) TestCases {
	allTestCases := make(TestCases, 0)
	for _, testCase := range testCases {
		if testCase.driver == ANY_DB_DRIVER {
			sqlite3TestCase := NewTestCase(testCase.createDefaultRules, testCase.adminUrlPathPrefix, "sqlite3")
			postgresTestCase := NewTestCase(testCase.createDefaultRules, testCase.adminUrlPathPrefix, "postgres")
			allTestCases = append(allTestCases, sqlite3TestCase)
			allTestCases = append(allTestCases, postgresTestCase)
		} else {
			allTestCases = append(allTestCases, testCase)
		}
	}
	return allTestCases
}

func (testCase *TestCase) skippable() bool {
	return testCase.driver == "postgres" && testing.Short()
}

func run(t *testing.T, checkFunc CheckFunc, testCase *TestCase) {
	if testCase.driver == "postgres" {
		createDb(TEST_DB)
		defer dropDb(TEST_DB)
	}
	goSlowServer := newSubDomainServer(testCase)
	server := httptest.NewServer(goSlowServer)
	defer server.Close()
	defer goSlowServer.storage.db.Close()
	checkFunc(t, server, testCase)
}

func checkZeroSite(t *testing.T, server *httptest.Server, testCase *TestCase) {
	bytesShouldBeEqual(t,
		readBody(GET(server.URL, "/", makeHost("0", TEST_ENDPOINT))),
		DEFAULT_BODY)
}

func TestTooLargeDelay(t *testing.T) {
	runAll(t, checkTooLargeDelay, defaultTestCases)
}

func checkTooLargeDelay(t *testing.T, server *httptest.Server, testCase *TestCase) {
	prefix := testCase.adminUrlPathPrefix
	site := newSite(server, join(prefix, "/booya"), []byte("haha"))
	resp := addRule(server, &Rule{Site: site, Delay: time.Duration(1000) * time.Second})
	shouldHaveStatusCode(t, http.StatusBadRequest, resp)
	resp = GET(server.URL, "/", site)
	shouldHaveStatusCode(t, http.StatusNotFound, resp)
}

func TestRedefineBuiltinSites(t *testing.T) {
	runAll(t, checkRedefineBuiltinSites, defaultTestCases)
}

func checkRedefineBuiltinSites(t *testing.T, server *httptest.Server, testCase *TestCase) {
	for _, site := range []string{"0", "99", "500", "create"} {
		resp := addRule(server, &Rule{Site: site, Path: "/test", Body: []byte("hop"), Method: "GET"})
		shouldHaveStatusCode(t, http.StatusForbidden, resp)
	}
}

func TestRedefineNonExistentSite(t *testing.T) {
	runAll(t, checkRedefineNonExistentSite, defaultTestCases)
}

func checkRedefineNonExistentSite(t *testing.T, server *httptest.Server, testCase *TestCase) {
	for _, site := range []string{"", "ha", "admin-500"} {
		resp := addRule(server, &Rule{Site: site})
		shouldHaveStatusCode(t, http.StatusNotFound, resp)
	}
}

func TestDelay(t *testing.T) {
	runAll(t, checkDelay, defaultTestCases)
}

func checkDelay(t *testing.T, server *httptest.Server, testCase *TestCase) {
	shouldRespondIn(t,
		createGET(server.URL, "/", makeHost("0", TEST_ENDPOINT)),
		0, 0.1) // seconds
	shouldRespondIn(t,
		createGET(server.URL, "/", makeHost("1", TEST_ENDPOINT)),
		1, 1.1) // seconds
}

func TestStatus(t *testing.T) {
	runAll(t, checkStatus, defaultTestCases)
}

func checkStatus(t *testing.T, server *httptest.Server, testCase *TestCase) {
	for _, statusCode := range []int{200, 404, 500} {
		resp := GET(server.URL, "/", makeHost(strconv.Itoa(statusCode), TEST_ENDPOINT))
		shouldHaveStatusCode(t, statusCode, resp)
	}
}

func TestRuleCreation(t *testing.T) {
	runAll(t, checkRuleCreationTestCase, ruleCreationTestCases)
}

// TODO: refactor
func checkRuleCreationTestCase(t *testing.T, server *httptest.Server, testCase *TestCase) {
	prefix := testCase.adminUrlPathPrefix
	site := ""
	if prefix == "" {
		site = newSite(server, join(prefix, "/"), []byte("haha"))
	} else {
		resp := addRule(server, &Rule{Path: join(prefix, "/"), Body: []byte("haha")})
		shouldHaveStatusCode(t, http.StatusOK, resp)
	}
	bytesShouldBeEqual(t, readBody(GET(server.URL, "/", site)), []byte("haha"))
	resp := addRule(server, &Rule{Site: site, Path: join(prefix, "/test"), Body: []byte("hop"), Method: "GET"})
	shouldHaveStatusCode(t, http.StatusOK, resp)

	bytesShouldBeEqual(t, readBody(GET(server.URL, "/test", site)), []byte("hop"))
	resp = POST(server.URL, "/test", site, []byte(""))
	intShouldBeEqual(t, 404, resp.StatusCode)
	resp = addRule(server, &Rule{Site: site, Path: join(prefix, "/test"), Body: []byte("for POST"), Method: "POST",
		Delay: time.Duration(100) * time.Millisecond})
	shouldHaveStatusCode(t, http.StatusOK, resp)
	bytesShouldBeEqual(t, readBody(GET(server.URL, "/test", site)), []byte("hop"))
	bytesShouldBeEqual(t, readBody(POST(server.URL, "/test", site, []byte(""))), []byte("for POST"))
	shouldRespondIn(t, createPOST(server.URL, "/test", site, []byte("")), 0.1, 0.15)
}

func newSubDomainServer(testCase *TestCase) *Server {
	config := DEFAULT_CONFIG // copies DEFAULT_CONFIG
	config.endpoint = TEST_ENDPOINT
	config.createDefaultRules = testCase.createDefaultRules
	config.adminUrlPathPrefix = testCase.adminUrlPathPrefix
	config.driver = testCase.driver
	config.dataSource = testCase.dataSource
	return NewServer(&config)
}

func GET(url, path, host string) *http.Response {
	req := createGET(url, path, host)
	return do(req)
}

func do(req *http.Request) *http.Response {
	resp, err := new(http.Client).Do(req)
	if err != nil {
		log.Fatal(err)
	}
	return resp
}

func createGET(url, path, host string) *http.Request {
	req, err := http.NewRequest("GET", url+path, nil)
	if err != nil {
		log.Fatal(err)
	}
	req.Host = host
	return req
}

func readBody(resp *http.Response) []byte {
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	return body
}

func makeHost(subdomain, host string) string {
	return fmt.Sprintf("%s.%s", subdomain, host)
}

func bytesShouldBeEqual(t *testing.T, expected, actual []byte) {
	if !bytes.Equal(expected, actual) {
		t.Fatalf("<<%v>> != <<%v>>", string(expected), string(actual))
	}
}

func intShouldBeEqual(t *testing.T, expected, actual int) {
	if expected != actual {
		t.Fatalf("<<%v>> != <<%v>>", expected, actual)
	}
}

func shouldRespondIn(t *testing.T, req *http.Request, minSeconds, maxSeconds float64) {
	start := time.Now()
	resp := do(req)
	readBody(resp)
	duration := time.Since(start)
	minDuration := toDuration(minSeconds)
	maxDuration := toDuration(maxSeconds)
	if duration < minDuration || duration > maxDuration {
		t.Fatalf("%s%s answered in %v. Not in the interval [%v; %v]",
			req.Host, req.URL.Path, duration, minDuration, maxDuration)
	}
}

func shouldHaveStatusCode(t *testing.T, statusCode int, resp *http.Response) {
	intShouldBeEqual(t, statusCode, resp.StatusCode)
}

func toDuration(seconds float64) time.Duration {
	return time.Duration(seconds*1000) * time.Millisecond
}

// TODO: rename, it returns a full domain, not a site
func newSite(server *httptest.Server, path string, response []byte) string {
	resp := POST(server.URL, fmt.Sprintf("%s?output=short&method=GET", path),
		makeHost("create", TEST_ENDPOINT), response)
	return string(readBody(resp))
}

func POST(url, path, host string, payload []byte) *http.Response {
	req := createPOST(url, path, host, payload)
	return do(req)
}

func createPOST(url, path, host string, payload []byte) *http.Request {
	req, err := http.NewRequest("POST", url+path, bytes.NewReader(payload))
	if err != nil {
		log.Fatal(err)
	}
	req.Host = host
	return req
}

// TODO: use normal GET-parameters builder, not Sprintf and string concatenation
func addRule(server *httptest.Server, rule *Rule) *http.Response {
	path := rule.Path
	path += "?method=" + rule.Method
	path += fmt.Sprintf("&delay=%f", rule.Delay.Seconds())
	return POST(server.URL, path, makeHost("admin-"+rule.Site, TEST_ENDPOINT),
		rule.Body)
}

// Wrapper around path.Join. Preserves trailing slash.
func join(elem ...string) string {
	lastElem := elem[len(elem)-1]
	shouldEndWithSlash := strings.HasSuffix(lastElem, "/")
	joined := path.Join(elem...)
	if shouldEndWithSlash && !strings.HasSuffix(joined, "/") {
		joined += "/"
	}
	return joined
}

func createDb(name string) {
	cmd := exec.Command("createdb", name)
	err := cmd.Run()
	if err != nil {
		log.Fatalf("createdb error: %s", err)
	}
}

func dropDb(name string) {
	cmd := exec.Command("dropdb", name)
	err := cmd.Run()
	if err != nil {
		log.Fatalf("dropdb error: %s", err)
	}
}
<commit_msg>newSite -> newDomain in tests<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"testing"
	"time"
)

const (
	ANY_DB_DRIVER = ""
	TEST_ENDPOINT = "localhost:9999"
	TEST_DB       = "goslow_test"
)

var DATA_SOURCE = map[string]string{
	"sqlite3":  ":memory:",
	"postgres": "postgres://localhost/" + TEST_DB,
}

type TestCase struct {
	createDefaultRules bool
	adminUrlPathPrefix string
	driver             string
	dataSource         string
}

type TestCases []*TestCase

type CheckFunc func(*testing.T, *httptest.Server, *TestCase)

var (
	defaultTestCases = TestCases{
		NewTestCase(true, "", ANY_DB_DRIVER),
	}

	ruleCreationTestCases = TestCases{
		NewTestCase(true, "", ANY_DB_DRIVER),
		NewTestCase(false, "/goslow", ANY_DB_DRIVER),
		NewTestCase(false, "/goslow/", ANY_DB_DRIVER),
		NewTestCase(false, "/te", ANY_DB_DRIVER),
		NewTestCase(false, "/te/", ANY_DB_DRIVER),
		NewTestCase(false, "/composite/path", ANY_DB_DRIVER),
	}
)

func NewTestCase(createDefaultRules bool, adminUrlPathPrefix string, driver string) *TestCase {
	dataSource, knownDriver := DATA_SOURCE[driver]
	if driver != ANY_DB_DRIVER && !knownDriver {
		log.Fatalf("unknown driver: <%s>", driver)
	}
	return &TestCase{
		createDefaultRules: createDefaultRules,
		adminUrlPathPrefix: adminUrlPathPrefix,
		driver:             driver,
		dataSource:         dataSource,
	}
}

func TestZeroSite(t *testing.T) {
	runAll(t, checkZeroSite, defaultTestCases)
}

func runAll(t *testing.T, checkFunc CheckFunc, testCases TestCases) {
	for _, testCase := range runnable(all(testCases)) {
		run(t, checkFunc, testCase)
	}
}

func runnable(testCases TestCases) TestCases {
	runnableTestCases := make(TestCases, 0)
	for _, testCase := range testCases {
		if !testCase.skippable() {
			runnableTestCases = append(runnableTestCases, testCase)
		}
	}
	return runnableTestCases
}

func all(testCases TestCases) TestCases {
	allTestCases := make(TestCases, 0)
	for _, testCase := range testCases {
		if testCase.driver == ANY_DB_DRIVER {
			sqlite3TestCase := NewTestCase(testCase.createDefaultRules, testCase.adminUrlPathPrefix, "sqlite3")
			postgresTestCase := NewTestCase(testCase.createDefaultRules, testCase.adminUrlPathPrefix, "postgres")
			allTestCases = append(allTestCases, sqlite3TestCase)
			allTestCases = append(allTestCases, postgresTestCase)
		} else {
			allTestCases = append(allTestCases, testCase)
		}
	}
	return allTestCases
}

func (testCase *TestCase) skippable() bool {
	return testCase.driver == "postgres" && testing.Short()
}

func run(t *testing.T, checkFunc CheckFunc, testCase *TestCase) {
	if testCase.driver == "postgres" {
		createDb(TEST_DB)
		defer dropDb(TEST_DB)
	}
	goSlowServer := newSubDomainServer(testCase)
	server := httptest.NewServer(goSlowServer)
	defer server.Close()
	defer goSlowServer.storage.db.Close()
	checkFunc(t, server, testCase)
}

func checkZeroSite(t *testing.T, server *httptest.Server, testCase *TestCase) {
	bytesShouldBeEqual(t,
		readBody(GET(server.URL, "/", makeHost("0", TEST_ENDPOINT))),
		DEFAULT_BODY)
}

func TestTooLargeDelay(t *testing.T) {
	runAll(t, checkTooLargeDelay, defaultTestCases)
}

func checkTooLargeDelay(t *testing.T, server *httptest.Server, testCase *TestCase) {
	prefix := testCase.adminUrlPathPrefix
	domain := newDomain(server, join(prefix, "/booya"), []byte("haha"))
	site := getSite(domain)
	resp := addRule(server, &Rule{Site: site, Delay: time.Duration(1000) * time.Second})
	shouldHaveStatusCode(t, http.StatusBadRequest, resp)
	resp = GET(server.URL, "/", domain)
	shouldHaveStatusCode(t, http.StatusNotFound, resp)
}

func TestRedefineBuiltinSites(t *testing.T) {
	runAll(t, checkRedefineBuiltinSites, defaultTestCases)
}

func checkRedefineBuiltinSites(t *testing.T, server *httptest.Server, testCase *TestCase) {
	for _, site := range []string{"0", "99", "500", "create"} {
		resp := addRule(server, &Rule{Site: site, Path: "/test", Body: []byte("hop"), Method: "GET"})
		shouldHaveStatusCode(t, http.StatusForbidden, resp)
	}
}

func TestRedefineNonExistentSite(t *testing.T) {
	runAll(t, checkRedefineNonExistentSite, defaultTestCases)
}

func checkRedefineNonExistentSite(t *testing.T, server *httptest.Server, testCase *TestCase) {
	for _, site := range []string{"", "ha", "admin-500"} {
		resp := addRule(server, &Rule{Site: site})
		shouldHaveStatusCode(t, http.StatusNotFound, resp)
	}
}

func TestDelay(t *testing.T) {
	runAll(t, checkDelay, defaultTestCases)
}

func checkDelay(t *testing.T, server *httptest.Server, testCase *TestCase) {
	shouldRespondIn(t,
		createGET(server.URL, "/", makeHost("0", TEST_ENDPOINT)),
		0, 0.1) // seconds
	shouldRespondIn(t,
		createGET(server.URL, "/", makeHost("1", TEST_ENDPOINT)),
		1, 1.1) // seconds
}

func TestStatus(t *testing.T) {
	runAll(t, checkStatus, defaultTestCases)
}

func checkStatus(t *testing.T, server *httptest.Server, testCase *TestCase) {
	for _, statusCode := range []int{200, 404, 500} {
		resp := GET(server.URL, "/", makeHost(strconv.Itoa(statusCode), TEST_ENDPOINT))
		shouldHaveStatusCode(t, statusCode, resp)
	}
}

func TestRuleCreation(t *testing.T) {
	runAll(t, checkRuleCreationTestCase, ruleCreationTestCases)
}

// TODO: refactor
func checkRuleCreationTestCase(t *testing.T, server *httptest.Server, testCase *TestCase) {
	prefix := testCase.adminUrlPathPrefix
	domain := TEST_ENDPOINT
	site := ""
	if prefix == "" {
		domain = newDomain(server, join(prefix, "/"), []byte("haha"))
		site = getSite(domain)
	} else {
		resp := addRule(server, &Rule{Path: join(prefix, "/"), Body: []byte("haha")})
		shouldHaveStatusCode(t, http.StatusOK, resp)
	}
	bytesShouldBeEqual(t, readBody(GET(server.URL, "/", domain)), []byte("haha"))
	resp := addRule(server, &Rule{Site: site, Path: join(prefix, "/test"), Body: []byte("hop"), Method: "GET"})
	shouldHaveStatusCode(t, http.StatusOK, resp)

	bytesShouldBeEqual(t, readBody(GET(server.URL, "/test", domain)), []byte("hop"))
	resp = POST(server.URL, "/test", domain, []byte(""))
	intShouldBeEqual(t, 404, resp.StatusCode)
	resp = addRule(server, &Rule{Site: site, Path: join(prefix, "/test"), Body: []byte("for POST"), Method: "POST",
		Delay: time.Duration(100) * time.Millisecond})
	shouldHaveStatusCode(t, http.StatusOK, resp)
	bytesShouldBeEqual(t, readBody(GET(server.URL, "/test", domain)), []byte("hop"))
	bytesShouldBeEqual(t, readBody(POST(server.URL, "/test", domain, []byte(""))), []byte("for POST"))
	shouldRespondIn(t, createPOST(server.URL, "/test", domain, []byte("")), 0.1, 0.15)
}

func newSubDomainServer(testCase *TestCase) *Server {
	config := DEFAULT_CONFIG // copies DEFAULT_CONFIG
	config.endpoint = TEST_ENDPOINT
	config.createDefaultRules = testCase.createDefaultRules
	config.adminUrlPathPrefix = testCase.adminUrlPathPrefix
	config.driver = testCase.driver
	config.dataSource = testCase.dataSource
	return NewServer(&config)
}

func GET(url, path, host string) *http.Response {
	req := createGET(url, path, host)
	return do(req)
}

func do(req *http.Request) *http.Response {
	resp, err := new(http.Client).Do(req)
	if err != nil {
		log.Fatal(err)
	}
	return resp
}

func createGET(url, path, host string) *http.Request {
	req, err := http.NewRequest("GET", url+path, nil)
	if err != nil {
		log.Fatal(err)
	}
	req.Host = host
	return req
}

func readBody(resp *http.Response) []byte {
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	return body
}

func makeHost(subdomain, host string) string {
	return fmt.Sprintf("%s.%s", subdomain, host)
}

func bytesShouldBeEqual(t *testing.T, expected, actual []byte) {
	if !bytes.Equal(expected, actual) {
		t.Fatalf("<<%v>> != <<%v>>", string(expected), string(actual))
	}
}

func intShouldBeEqual(t *testing.T, expected, actual int) {
	if expected != actual {
		t.Fatalf("<<%v>> != <<%v>>", expected, actual)
	}
}

func shouldRespondIn(t *testing.T, req *http.Request, minSeconds, maxSeconds float64) {
	start := time.Now()
	resp := do(req)
	readBody(resp)
	duration := time.Since(start)
	minDuration := toDuration(minSeconds)
	maxDuration := toDuration(maxSeconds)
	if duration < minDuration || duration > maxDuration {
		t.Fatalf("%s%s answered in %v. Not in the interval [%v; %v]",
			req.Host, req.URL.Path, duration, minDuration, maxDuration)
	}
}

func shouldHaveStatusCode(t *testing.T, statusCode int, resp *http.Response) {
	intShouldBeEqual(t, statusCode, resp.StatusCode)
}

func toDuration(seconds float64) time.Duration {
	return time.Duration(seconds*1000) * time.Millisecond
}

// TODO: rename, it returns a full domain, not a site
func newDomain(server *httptest.Server, path string, response []byte) string {
	resp := POST(server.URL, fmt.Sprintf("%s?output=short&method=GET", path),
		makeHost("create", TEST_ENDPOINT), response)
	return string(readBody(resp))
}

func POST(url, path, host string, payload []byte) *http.Response {
	req := createPOST(url, path, host, payload)
	return do(req)
}

func createPOST(url, path, host string, payload []byte) *http.Request {
	req, err := http.NewRequest("POST", url+path, bytes.NewReader(payload))
	if err != nil {
		log.Fatal(err)
	}
	req.Host = host
	return req
}

func getSite(domain string) string {
	return strings.Split(domain, ".")[0]
}

// TODO: use normal GET-parameters builder, not Sprintf and string concatenation
func addRule(server *httptest.Server, rule *Rule) *http.Response {
	path := rule.Path
	path += "?method=" + rule.Method
	path += fmt.Sprintf("&delay=%f", rule.Delay.Seconds())
	return POST(server.URL, path, makeHost("admin-"+rule.Site, TEST_ENDPOINT),
		rule.Body)
}

// Wrapper around path.Join. Preserves trailing slash.
func join(elem ...string) string {
	lastElem := elem[len(elem)-1]
	shouldEndWithSlash := strings.HasSuffix(lastElem, "/")
	joined := path.Join(elem...)
	if shouldEndWithSlash && !strings.HasSuffix(joined, "/") {
		joined += "/"
	}
	return joined
}

func createDb(name string) {
	cmd := exec.Command("createdb", name)
	err := cmd.Run()
	if err != nil {
		log.Fatalf("createdb error: %s", err)
	}
}

func dropDb(name string) {
	cmd := exec.Command("dropdb", name)
	err := cmd.Run()
	if err != nil {
		log.Fatalf("dropdb error: %s", err)
	}
}
<|endoftext|>


<commit_before>package pqstream

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"regexp"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

var testConnectionString = "postgres://localhost?sslmode=disable"
var testConnectionStringTemplate = "postgres://localhost/%s?sslmode=disable"

var testDatabaseDDL = `create table notes (id serial, created_at timestamp, note text)`
var testInsert = `insert into notes values (default, default, 'here is a sample note')`
var testUpdate = `update notes set note = 'here is an updated note' where id=1`

func init() {
	if s := os.Getenv("PQSTREAM_TEST_DB_URL"); s != "" {
		testConnectionString = s
	}
	if s := os.Getenv("PQSTREAM_TEST_DB_TMPL_URL"); s != "" {
		testConnectionStringTemplate = s
	}
}

func TestWithTableRegexp(t *testing.T) {
	re := regexp.MustCompile(".*")
	tests := []struct {
		name string
		want *regexp.Regexp
	}{
		{"basic", re},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s, err := NewServer(testConnectionString, WithTableRegexp(re))
			if err != nil {
				t.Fatal(err)
			}
			if got := s.tableRe; got != tt.want {
				t.Errorf("WithTableRegexp() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestNewServer(t *testing.T) {
	type args struct {
		connectionString string
		opts             []ServerOption
	}
	tests := []struct {
		name    string
		args    args
		check   func(t *testing.T, s *Server)
		wantErr bool
	}{
		{"bad", args{
			connectionString: "this is an invalid connection string",
		}, nil, true},
		{"good", args{
			connectionString: testConnectionString,
		}, nil, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := NewServer(tt.args.connectionString, tt.args.opts...)
			if (err != nil) != tt.wantErr {
				t.Errorf("NewServer() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if tt.check != nil {
				tt.check(t, got)
			}
		})
	}
}

func dbOrSkip(t *testing.T) *sql.DB {
	t.Helper()
	db, err := sql.Open("postgres", testConnectionString)
	if err != nil {
		t.Skip(err)
	}
	if err := db.Ping(); err != nil {
		t.Skip(errors.Wrap(err, "ping"))
	}
	return db
}

func testDBConn(t *testing.T, db *sql.DB, testcase string) (connectionString string, cleanup func()) {
	s := fmt.Sprintf("test_pqstream_%s", testcase)
	db.Exec(fmt.Sprintf("drop database %s", s))
	_, err := db.Exec(fmt.Sprintf("create database %s", s))
	if err != nil {
		t.Fatal(err)
	}
	dsn := fmt.Sprintf(testConnectionStringTemplate, s)
	newDB, err := sql.Open("postgres", dsn)
	if err != nil {
		t.Skip(err)
	}
	if err := db.Ping(); err != nil {
		t.Skip(errors.Wrap(err, "ping"))
	}
	defer newDB.Close()
	_, err = newDB.Exec(testDatabaseDDL)
	if err != nil {
		t.Fatal(err)
	}
	return dsn, func() {
		_, err := db.Exec(fmt.Sprintf("drop database %s", s))
		if err != nil {
			t.Fatal(err)
		}
	}
}

func TestServer_HandleEvents(t *testing.T) {
	db := dbOrSkip(t)
	logger := logrus.New()
	logger.Level = logrus.DebugLevel
	type testCase struct {
		name    string
		opts    []ServerOption
		fn      func(*testing.T, *Server)
		wantErr bool
	}
	tests := []testCase{
		{"basics", []ServerOption{WithLogger(logger)}, nil, false},
		{"basic_insert", nil, func(t *testing.T, s *Server) {
			if _, err := s.db.Exec(testInsert); err != nil {
				t.Fatal(err)
			}
		}, false},
		{"basic_insert_and_update", nil, func(t *testing.T, s *Server) {
			if _, err := s.db.Exec(testInsert); err != nil {
				t.Fatal(err)
			}
			time.Sleep(time.Second)
			if _, err := s.db.Exec(testUpdate); err != nil {
				t.Fatal(err)
			}
		}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			caseName := tt.name
			t.Parallel()
			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
			defer cancel()
			cs, cleanup := testDBConn(t, db, caseName)
			defer cleanup()
			s, err := NewServer(cs, tt.opts...)
			s.listenerPingInterval = time.Second // move into a helper?
			if err != nil {
				t.Fatal(err)
			}
			s.InstallTriggers()
			defer s.RemoveTriggers()
			defer s.Close()
			go func(t *testing.T, tt testCase) {
				if err := s.HandleEvents(ctx); (err != nil) != tt.wantErr {
					t.Errorf("Server.HandleEvents(%v) error = %v, wantErr %v", ctx, err, tt.wantErr)
				}
			}(t, tt)
			if tt.fn != nil {
				tt.fn(t, s)
			}
			<-ctx.Done()

		})
	}
}
<commit_msg>test: test empty conn string<commit_after>package pqstream

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"regexp"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

var testConnectionString = "postgres://localhost?sslmode=disable"
var testConnectionStringTemplate = "postgres://localhost/%s?sslmode=disable"

var testDatabaseDDL = `create table notes (id serial, created_at timestamp, note text)`
var testInsert = `insert into notes values (default, default, 'here is a sample note')`
var testUpdate = `update notes set note = 'here is an updated note' where id=1`

func init() {
	if s := os.Getenv("PQSTREAM_TEST_DB_URL"); s != "" {
		testConnectionString = s
	}
	if s := os.Getenv("PQSTREAM_TEST_DB_TMPL_URL"); s != "" {
		testConnectionStringTemplate = s
	}
}

func TestWithTableRegexp(t *testing.T) {
	re := regexp.MustCompile(".*")
	tests := []struct {
		name string
		want *regexp.Regexp
	}{
		{"basic", re},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s, err := NewServer(testConnectionString, WithTableRegexp(re))
			if err != nil {
				t.Fatal(err)
			}
			if got := s.tableRe; got != tt.want {
				t.Errorf("WithTableRegexp() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestNewServer(t *testing.T) {
	type args struct {
		connectionString string
		opts             []ServerOption
	}
	tests := []struct {
		name    string
		args    args
		check   func(t *testing.T, s *Server)
		wantErr bool
	}{
		{"bad", args{
			connectionString: "this is an invalid connection string",
		}, nil, true},
		{"empty", args{
			connectionString: "",
		}, nil, true},
		{"good", args{
			connectionString: testConnectionString,
		}, nil, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := NewServer(tt.args.connectionString, tt.args.opts...)
			if (err != nil) != tt.wantErr {
				t.Errorf("NewServer() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if tt.check != nil {
				tt.check(t, got)
			}
		})
	}
}

func dbOrSkip(t *testing.T) *sql.DB {
	t.Helper()
	db, err := sql.Open("postgres", testConnectionString)
	if err != nil {
		t.Skip(err)
	}
	if err := db.Ping(); err != nil {
		t.Skip(errors.Wrap(err, "ping"))
	}
	return db
}

func testDBConn(t *testing.T, db *sql.DB, testcase string) (connectionString string, cleanup func()) {
	s := fmt.Sprintf("test_pqstream_%s", testcase)
	db.Exec(fmt.Sprintf("drop database %s", s))
	_, err := db.Exec(fmt.Sprintf("create database %s", s))
	if err != nil {
		t.Fatal(err)
	}
	dsn := fmt.Sprintf(testConnectionStringTemplate, s)
	newDB, err := sql.Open("postgres", dsn)
	if err != nil {
		t.Skip(err)
	}
	if err := db.Ping(); err != nil {
		t.Skip(errors.Wrap(err, "ping"))
	}
	defer newDB.Close()
	_, err = newDB.Exec(testDatabaseDDL)
	if err != nil {
		t.Fatal(err)
	}
	return dsn, func() {
		_, err := db.Exec(fmt.Sprintf("drop database %s", s))
		if err != nil {
			t.Fatal(err)
		}
	}
}

func TestServer_HandleEvents(t *testing.T) {
	db := dbOrSkip(t)
	logger := logrus.New()
	logger.Level = logrus.DebugLevel
	type testCase struct {
		name    string
		opts    []ServerOption
		fn      func(*testing.T, *Server)
		wantErr bool
	}
	tests := []testCase{
		{"basics", []ServerOption{WithLogger(logger)}, nil, false},
		{"basic_insert", nil, func(t *testing.T, s *Server) {
			if _, err := s.db.Exec(testInsert); err != nil {
				t.Fatal(err)
			}
		}, false},
		{"basic_insert_and_update", nil, func(t *testing.T, s *Server) {
			if _, err := s.db.Exec(testInsert); err != nil {
				t.Fatal(err)
			}
			time.Sleep(time.Second)
			if _, err := s.db.Exec(testUpdate); err != nil {
				t.Fatal(err)
			}
		}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			caseName := tt.name
			t.Parallel()
			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
			defer cancel()
			cs, cleanup := testDBConn(t, db, caseName)
			defer cleanup()
			s, err := NewServer(cs, tt.opts...)
			s.listenerPingInterval = time.Second // move into a helper?
			if err != nil {
				t.Fatal(err)
			}
			s.InstallTriggers()
			defer s.RemoveTriggers()
			defer s.Close()
			go func(t *testing.T, tt testCase) {
				if err := s.HandleEvents(ctx); (err != nil) != tt.wantErr {
					t.Errorf("Server.HandleEvents(%v) error = %v, wantErr %v", ctx, err, tt.wantErr)
				}
			}(t, tt)
			if tt.fn != nil {
				tt.fn(t, s)
			}
			<-ctx.Done()

		})
	}
}
<|endoftext|>


<commit_before>package emil

import (
	"errors"
	"fmt"
	"time"
)

var errNowNewAnalysis = errors.New("errNowNewAnalysis")

type Analysis struct {
	dtm int // Depth to mate

	board *Board
	move  *Move
}

func (a *Analysis) Move() *Move {
	return a.move
}
func (a *Analysis) Board() *Board {
	return a.board
}

// EndGameDb to query for mate in 1,2, etc.
type EndGameDb struct {
	positionDb map[string]*Analysis

	dtmDb []map[string]bool

	searchedPositions int
}

func (db *EndGameDb) Find(board *Board) (bestMove *Move) {
	if DEBUG {
		fmt.Printf("Find:\n%s\n", board.String())
	}
	a := db.positionDb[board.String()]
	if DEBUG {
		fmt.Printf("Found: positionDb with dtm %d\n", a.dtm)
	}
	return a.move
}
func (db *EndGameDb) FindMatesIn(dtm int) (as []*Analysis) {
	if dtm == -1 {
		for _, a := range db.positionDb {
			if a.dtm == -1 {
				as = append(as, a)
			}
		}
	} else {
		for str := range db.dtmDb[dtm] {
			as = append(as, db.positionDb[str])
		}
	}
	return as
}

func (db *EndGameDb) FindMates() (as []*Analysis) {
	return db.FindMatesIn(0)
}

func (db *EndGameDb) FindMate(piece, square int) (boards []*Board) {
	for str := range db.dtmDb[0] {
		a := db.positionDb[str]
		if a.board.squares[square] == piece {
			boards = append(boards, a.board)
		}
	}
	return boards
}

func (db *EndGameDb) addPosition(board *Board) {
	db.addAnalysis(board, -1, nil)
}

func (db *EndGameDb) addAnalysis(board *Board, dtm int, move *Move) {
	a := &Analysis{
		dtm:   dtm,
		board: board}
	if move != nil {
		a.move = move.reverse()
	}
	if dtm >= 0 {
		if move != nil {
			playerForStep := playerForStepN(dtm)
			if playerForStep != move.player {
				panic("playerForStep != move.player")
			}
		}

		db.dtmDb[dtm][a.board.String()] = true
	}

	db.positionDb[a.board.String()] = a
}

func (db *EndGameDb) positions() int {
	return len(db.positionDb)
}

// find positions where black is checkmate
func (db *EndGameDb) retrogradeAnalysisStep1() {
	db.dtmDb = append(db.dtmDb, make(map[string]bool))

	start := time.Now()

	player := BLACK
	for boardStr, a := range db.positionDb {
		// mate only on border square
		blackKingSquare := BoardSquares[a.board.blackKing]
		if !blackKingSquare.isBorder {
			continue
		}
		// mate only with help from king
		if squaresDistances[a.board.blackKing][a.board.whiteKing] > 2 {
			continue
		}

		p := NewPosition(a.board, player)

		move := Search(p)
		db.searchedPositions++
		if move == nil {
			if isKingInCheck(p) {
				a.dtm = 0
				db.addAnalysis(a.board, 0, nil)
				if DEBUG {
					fmt.Printf("mate:\n%s\n", boardStr)
				}
			}
		}
	}
	end := time.Now()
	if DEBUG {
		fmt.Printf("searchedPositions %d\n", db.searchedPositions)
		fmt.Printf("db.dtmDb[0] %d\n", len(db.dtmDb[0]))
		fmt.Printf("duration %v\n\n\n", end.Sub(start))
	}
}
func playerForStepN(dtm int) (player int) {
	if dtm%2 == 0 {
		return BLACK
	}
	return WHITE
}

func (db *EndGameDb) retrogradeAnalysisStepN(dtm int) (noError error) {
	start := time.Now()
	db.dtmDb = append(db.dtmDb, make(map[string]bool))

	player := playerForStepN(dtm)

	positions := 0
	if player == WHITE {
		for _, a := range db.positionDb {
			if db.isMateIn1357(a.board, dtm) >= 0 {
				positions++
			}
		}
		if DEBUG {
			fmt.Printf("WHITE Start positions %d\n", len(db.positionDb)-positions)
		}
		for _, a := range db.positionDb {
			if db.isMateIn1357(a.board, dtm) >= 0 {
				continue
			}
			p := NewPosition(a.board, player)
			if IsTheKingInCheck(p) {
				continue
			}
			moves := generateMoves(p)

			for _, m := range moves {
				newBoard := a.board.doMove(m)
				newDtm := db.isMateIn0246(newBoard, dtm)
				if newDtm == dtm-1 {
					db.addAnalysis(a.board, dtm, m)
				}
			}
		}
	} else {
		// Suche alle Stellungen, bei denen Schwarz am Zug ist und
		// **jeder** Zug von ihm zu einer Stellung unter 2. führt.
		// Schwarz kann hier Matt in einem Zug nicht verhindern.
		// Markiere diese Stellungen in der Datei.
		for _, a := range db.positionDb {
			if db.isMateIn0246(a.board, dtm) >= 0 {
				positions++
			}
		}
		if DEBUG {
			fmt.Printf("BLACK Start positions %d\n", len(db.positionDb)-positions)
		}
		for _, a := range db.positionDb {
			if db.isMateIn0246(a.board, dtm) >= 0 {
				continue
			}
			p := NewPosition(a.board, player)
			moves := GenerateMoves(p)

			found := 0
			maxDTM := -1
			for _, m := range moves {
				newBoard := a.board.doMove(m)
				newDtm := db.isMateIn1357(newBoard, dtm)
				if newDtm > maxDTM {
					maxDTM = newDtm
				}
				if db.isMateIn1357(newBoard, dtm) >= 0 {
					found++
				}
			}

			if found == len(moves) {
				for _, m := range moves {
					db.addAnalysis(a.board, maxDTM+1, m)
				}
			}
		}
	}
	end := time.Now()

	if DEBUG {
		fmt.Printf("db.dtmDb[%d] %d\n", dtm, len(db.dtmDb[dtm]))
		fmt.Printf("duration %v\n\n\n", end.Sub(start))
	}

	if len(db.dtmDb[dtm]) == 0 {
		return errNowNewAnalysis
	}
	return noError
}
func (db *EndGameDb) isMateIn0246(board *Board, maxDtm int) int {
	for dtm := 0; dtm < maxDtm; dtm += 2 {
		_, ok := db.dtmDb[dtm][board.String()]
		if ok {
			return dtm
		}
	}
	return -1
}
func (db *EndGameDb) isMateIn1357(board *Board, maxDtm int) int {
	for dtm := 1; dtm < maxDtm; dtm += 2 {
		_, ok := db.dtmDb[dtm][board.String()]
		if ok {
			return dtm
		}
	}
	return -1
}

func (db *EndGameDb) MaxDtm() int {
	return len(db.dtmDb)
}

func (db *EndGameDb) retrogradeAnalysis() {
	// find positions where black is checkmate
	db.retrogradeAnalysisStep1()
	dtm := 1
	for {
		err := db.retrogradeAnalysisStepN(dtm)
		if err != nil {
			break
		}
		dtm++
	}
}
func GenerateMoves(p *position) (list []*Move) {
	for _, m := range generateMoves(p) {
		b := p.board.DoMove(m)
		if !IsTheKingInCheck(NewPosition(b, WHITE)) {
			list = append(list, m)
		}
	}
	return list
}
func generateMoves(p *position) (list []*Move) {
	for src, piece := range p.board.squares {
		if isOwnPiece(p.player, piece) {
			switch abs(piece) {
			case kingValue:
				for _, dst := range kingDestinationsFrom(src) {
					capture := p.board.squares[dst]
					if isOtherKing(p.player, capture) {
						continue
					}
					if capture == Empty {
						list = append(list, newSilentMove(p.player, piece, src, dst))
					} else if !isOwnPiece(p.player, capture) {
						list = append(list, newCaptureMove(p.player, piece, capture, src, dst))
					}
				}
			case rockValue:
				for _, dsts := range rockDestinationsFrom(src) {
					for _, dst := range dsts {
						capture := p.board.squares[dst]
						if isOtherKing(p.player, capture) {
							break
						}
						if capture == Empty {
							list = append(list, newSilentMove(p.player, piece, src, dst))
						} else if !isOwnPiece(p.player, capture) {
							list = append(list, newCaptureMove(p.player, piece, capture, src, dst))
							break
						} else {
							break // onOwnPiece
						}
					}
				}
			}
		}
	}
	return list
}

// NewEndGameDb generates an end game DB for KRK
func NewEndGameDb() *EndGameDb {
	var err error
	start := time.Now()

	endGames := &EndGameDb{
		positionDb: make(map[string]*Analysis),
		dtmDb:      make([]map[string]bool, 0)}

	for wk := A1; wk <= H8; wk++ {
		//for wk := E3; wk <= E3; wk++ {
		if DEBUG {
			fmt.Printf("White king on %s\n", BoardSquares[wk])
		}
		for wr := A1; wr <= H8; wr++ {
			for bk := A1; bk <= H8; bk++ {

				board := NewBoard()

				err = board.Setup(WhiteKing, wk)
				if err != nil {
					continue
				}

				err = board.Setup(WhiteRock, wr)
				if err != nil {
					continue
				}

				err = board.Setup(BlackKing, bk)
				if err != nil {
					continue
				}

				err = board.kingsToClose()
				if err != nil {
					continue
				}

				endGames.addPosition(board)
			}
		}
	}
	end := time.Now()
	if DEBUG {
		fmt.Printf("all positions %d\n", 64*63*62)
		fmt.Printf("endGames.positions() %d\n", endGames.positions())
		fmt.Printf("difference %d\n", 64*63*62-endGames.positions())
		fmt.Printf("duration %v\n", end.Sub(start))
	}
	endGames.retrogradeAnalysis()

	return endGames
}
<commit_msg>refactor retrogradeAnalysisStepN<commit_after>package emil

import (
	"errors"
	"fmt"
	"time"
)

var errNowNewAnalysis = errors.New("errNowNewAnalysis")

type Analysis struct {
	dtm int // Depth to mate

	board *Board
	move  *Move
}

func (a *Analysis) Move() *Move {
	return a.move
}
func (a *Analysis) Board() *Board {
	return a.board
}

// EndGameDb to query for mate in 1,2, etc.
type EndGameDb struct {
	positionDb map[string]*Analysis

	dtmDb []map[string]bool

	searchedPositions int
}

func (db *EndGameDb) Find(board *Board) (bestMove *Move) {
	if DEBUG {
		fmt.Printf("Find:\n%s\n", board.String())
	}
	a := db.positionDb[board.String()]
	if DEBUG {
		fmt.Printf("Found: positionDb with dtm %d\n", a.dtm)
	}
	return a.move
}
func (db *EndGameDb) FindMatesIn(dtm int) (as []*Analysis) {
	if dtm == -1 {
		for _, a := range db.positionDb {
			if a.dtm == -1 {
				as = append(as, a)
			}
		}
	} else {
		for str := range db.dtmDb[dtm] {
			as = append(as, db.positionDb[str])
		}
	}
	return as
}

func (db *EndGameDb) FindMates() (as []*Analysis) {
	return db.FindMatesIn(0)
}

func (db *EndGameDb) FindMate(piece, square int) (boards []*Board) {
	for str := range db.dtmDb[0] {
		a := db.positionDb[str]
		if a.board.squares[square] == piece {
			boards = append(boards, a.board)
		}
	}
	return boards
}

func (db *EndGameDb) addPosition(board *Board) {
	db.addAnalysis(board, -1, nil)
}

func (db *EndGameDb) addAnalysis(board *Board, dtm int, move *Move) {
	a := &Analysis{
		dtm:   dtm,
		board: board}
	if move != nil {
		a.move = move.reverse()
	}
	if dtm >= 0 {
		if move != nil {
			playerForStep := playerForStepN(dtm)
			if playerForStep != move.player {
				panic("playerForStep != move.player")
			}
		}

		db.dtmDb[dtm][a.board.String()] = true
	}

	db.positionDb[a.board.String()] = a
}

func (db *EndGameDb) positions() int {
	return len(db.positionDb)
}

// find positions where black is checkmate
func (db *EndGameDb) retrogradeAnalysisStep1() {
	db.dtmDb = append(db.dtmDb, make(map[string]bool))

	start := time.Now()

	player := BLACK
	for boardStr, a := range db.positionDb {
		// mate only on border square
		blackKingSquare := BoardSquares[a.board.blackKing]
		if !blackKingSquare.isBorder {
			continue
		}
		// mate only with help from king
		if squaresDistances[a.board.blackKing][a.board.whiteKing] > 2 {
			continue
		}

		p := NewPosition(a.board, player)

		move := Search(p)
		db.searchedPositions++
		if move == nil {
			if isKingInCheck(p) {
				a.dtm = 0
				db.addAnalysis(a.board, 0, nil)
				if DEBUG {
					fmt.Printf("mate:\n%s\n", boardStr)
				}
			}
		}
	}
	end := time.Now()
	if DEBUG {
		fmt.Printf("searchedPositions %d\n", db.searchedPositions)
		fmt.Printf("db.dtmDb[0] %d\n", len(db.dtmDb[0]))
		fmt.Printf("duration %v\n\n\n", end.Sub(start))
	}
}
func playerForStepN(dtm int) (player int) {
	if dtm%2 == 0 {
		return BLACK
	}
	return WHITE
}

func (db *EndGameDb) retrogradeAnalysisStepN(dtm int) (noError error) {
	start := time.Now()
	db.dtmDb = append(db.dtmDb, make(map[string]bool))

	player := playerForStepN(dtm)

	positions := 0
	if player == WHITE {
		for _, a := range db.positionDb {
			if db.isMateIn1357(a.board, dtm) >= 0 {
				positions++
			}
		}
		if DEBUG {
			fmt.Printf("WHITE Start positions %d\n", len(db.positionDb)-positions)
		}
		for _, a := range db.positionDb {
			if db.isMateIn1357(a.board, dtm) >= 0 {
				continue
			}
			p := NewPosition(a.board, player)
			if IsTheKingInCheck(p) {
				continue
			}
			moves := generateMoves(p)

			for _, m := range moves {
				newBoard := a.board.doMove(m)
				newDtm := db.isMateIn0246(newBoard, dtm)
				if newDtm == dtm-1 {
					db.addAnalysis(a.board, dtm, m)
				}
			}
		}
	} else {
		for _, a := range db.positionDb {
			if db.isMateIn0246(a.board, dtm) >= 0 {
				positions++
			}
		}
		if DEBUG {
			fmt.Printf("BLACK Start positions %d\n", len(db.positionDb)-positions)
		}
		for _, a := range db.positionDb {
			if db.isMateIn0246(a.board, dtm) >= 0 {
				continue
			}
			p := NewPosition(a.board, player)
			moves := GenerateMoves(p)

			found := 0
			maxDTM := -1
			for _, m := range moves {
				newBoard := a.board.doMove(m)
				newDtm := db.isMateIn1357(newBoard, dtm)
				if newDtm > maxDTM {
					maxDTM = newDtm
				}
				if db.isMateIn1357(newBoard, dtm) >= 0 {
					found++
				}
			}

			if found == len(moves) {
				for _, m := range moves {
					db.addAnalysis(a.board, maxDTM+1, m)
				}
			}
		}
	}
	end := time.Now()

	if DEBUG {
		fmt.Printf("db.dtmDb[%d] %d\n", dtm, len(db.dtmDb[dtm]))
		fmt.Printf("duration %v\n\n\n", end.Sub(start))
	}

	if len(db.dtmDb[dtm]) == 0 {
		return errNowNewAnalysis
	}
	return noError
}
func (db *EndGameDb) isMateIn0246(board *Board, maxDtm int) int {
	for dtm := 0; dtm < maxDtm; dtm += 2 {
		_, ok := db.dtmDb[dtm][board.String()]
		if ok {
			return dtm
		}
	}
	return -1
}
func (db *EndGameDb) isMateIn1357(board *Board, maxDtm int) int {
	for dtm := 1; dtm < maxDtm; dtm += 2 {
		_, ok := db.dtmDb[dtm][board.String()]
		if ok {
			return dtm
		}
	}
	return -1
}

func (db *EndGameDb) MaxDtm() int {
	return len(db.dtmDb)
}

func (db *EndGameDb) retrogradeAnalysis() {
	// find positions where black is checkmate
	db.retrogradeAnalysisStep1()
	dtm := 1
	for {
		err := db.retrogradeAnalysisStepN(dtm)
		if err != nil {
			break
		}
		dtm++
	}
}
func GenerateMoves(p *position) (list []*Move) {
	for _, m := range generateMoves(p) {
		b := p.board.DoMove(m)
		if !IsTheKingInCheck(NewPosition(b, WHITE)) {
			list = append(list, m)
		}
	}
	return list
}
func generateMoves(p *position) (list []*Move) {
	for src, piece := range p.board.squares {
		if isOwnPiece(p.player, piece) {
			switch abs(piece) {
			case kingValue:
				for _, dst := range kingDestinationsFrom(src) {
					capture := p.board.squares[dst]
					if isOtherKing(p.player, capture) {
						continue
					}
					if capture == Empty {
						list = append(list, newSilentMove(p.player, piece, src, dst))
					} else if !isOwnPiece(p.player, capture) {
						list = append(list, newCaptureMove(p.player, piece, capture, src, dst))
					}
				}
			case rockValue:
				for _, dsts := range rockDestinationsFrom(src) {
					for _, dst := range dsts {
						capture := p.board.squares[dst]
						if isOtherKing(p.player, capture) {
							break
						}
						if capture == Empty {
							list = append(list, newSilentMove(p.player, piece, src, dst))
						} else if !isOwnPiece(p.player, capture) {
							list = append(list, newCaptureMove(p.player, piece, capture, src, dst))
							break
						} else {
							break // onOwnPiece
						}
					}
				}
			}
		}
	}
	return list
}

// NewEndGameDb generates an end game DB for KRK
func NewEndGameDb() *EndGameDb {
	var err error
	start := time.Now()

	endGames := &EndGameDb{
		positionDb: make(map[string]*Analysis),
		dtmDb:      make([]map[string]bool, 0)}

	for wk := A1; wk <= H8; wk++ {
		//for wk := E3; wk <= E3; wk++ {
		if DEBUG {
			fmt.Printf("White king on %s\n", BoardSquares[wk])
		}
		for wr := A1; wr <= H8; wr++ {
			for bk := A1; bk <= H8; bk++ {

				board := NewBoard()

				err = board.Setup(WhiteKing, wk)
				if err != nil {
					continue
				}

				err = board.Setup(WhiteRock, wr)
				if err != nil {
					continue
				}

				err = board.Setup(BlackKing, bk)
				if err != nil {
					continue
				}

				err = board.kingsToClose()
				if err != nil {
					continue
				}

				endGames.addPosition(board)
			}
		}
	}
	end := time.Now()
	if DEBUG {
		fmt.Printf("all positions %d\n", 64*63*62)
		fmt.Printf("endGames.positions() %d\n", endGames.positions())
		fmt.Printf("difference %d\n", 64*63*62-endGames.positions())
		fmt.Printf("duration %v\n", end.Sub(start))
	}
	endGames.retrogradeAnalysis()

	return endGames
}
<|endoftext|>


<commit_before>// Copyright 2017 Apcera Inc. All rights reserved.

package server

import (
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"testing"
	"time"

	"github.com/nats-io/go-nats"
)

// Ensure Reload returns an error when attempting to reload a server that did
// not start with a config file.
func TestConfigReloadNoConfigFile(t *testing.T) {
	server := New(&Options{})
	if server.Reload() == nil {
		t.Fatal("Expected Reload to return an error")
	}
}

// Ensure Reload returns an error when attempting to change an option which
// does not support reloading.
func TestConfigReloadUnsupported(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	golden := &Options{
		ConfigFile:     config,
		Host:           "0.0.0.0",
		Port:           4222,
		AuthTimeout:    1.0,
		Debug:          false,
		Trace:          false,
		Logtime:        false,
		MaxControlLine: 1024,
		MaxPayload:     1048576,
		MaxConn:        65536,
		PingInterval:   2 * time.Minute,
		MaxPingsOut:    2,
		WriteDeadline:  2 * time.Second,
	}
	processOptions(golden)

	if err := os.Symlink("./configs/reload/test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v (ensure you have privileges)", err)
	}
	defer os.Remove(config)
	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}
	server := New(opts)

	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}

	// Change config file to bad config by replacing symlink.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/reload/reload_unsupported.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v (ensure you have privileges)", err)
	}

	// This should fail because `debug` cannot be changed.
	if err := server.Reload(); err == nil {
		t.Fatal("Expected Reload to return an error")
	}

	// Ensure config didn't change.
	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}
}

// Ensure Reload returns an error when reloading from a bad config file.
func TestConfigReloadInvalidConfig(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	golden := &Options{
		ConfigFile:     config,
		Host:           "0.0.0.0",
		Port:           4222,
		AuthTimeout:    1.0,
		Debug:          false,
		Trace:          false,
		Logtime:        false,
		MaxControlLine: 1024,
		MaxPayload:     1048576,
		MaxConn:        65536,
		PingInterval:   2 * time.Minute,
		MaxPingsOut:    2,
		WriteDeadline:  2 * time.Second,
	}
	processOptions(golden)

	if err := os.Symlink("./configs/reload/test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)
	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}
	server := New(opts)

	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}

	// Change config file to bad config by replacing symlink.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/reload/invalid.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}

	// This should fail because the new config should not parse.
	if err := server.Reload(); err == nil {
		t.Fatal("Expected Reload to return an error")
	}

	// Ensure config didn't change.
	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}
}

// Ensure Reload returns nil and the config is changed on success.
func TestConfigReload(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	golden := &Options{
		ConfigFile:     config,
		Host:           "0.0.0.0",
		Port:           4222,
		AuthTimeout:    1.0,
		Debug:          false,
		Trace:          false,
		Logtime:        false,
		MaxControlLine: 1024,
		MaxPayload:     1048576,
		MaxConn:        65536,
		PingInterval:   2 * time.Minute,
		MaxPingsOut:    2,
		WriteDeadline:  2 * time.Second,
	}
	processOptions(golden)

	if err := os.Symlink("./configs/reload/test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)
	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}
	server := New(opts)

	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}

	// Change config file to new config by replacing symlink.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/reload/reload.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}

	if err := server.Reload(); err != nil {
		t.Fatalf("Error reloading config: %v", err)
	}

	// Ensure config changed.
	updated := server.getOpts()
	if !updated.Trace {
		t.Fatal("Expected Trace to be true")
	}
	if !updated.Debug {
		t.Fatal("Expected Debug to be true")
	}
	if updated.TLSConfig == nil {
		t.Fatal("Expected TLSConfig to be non-nil")
	}
	if !server.info.TLSRequired {
		t.Fatal("Expected TLSRequired to be true")
	}
	if !server.info.TLSVerify {
		t.Fatal("Expected TLSVerify to be true")
	}
}

// Ensure Reload supports TLS config changes. Test this by starting a server
// with TLS enabled, connect to it to verify, reload config using a different
// key pair and client verification enabled, ensure reconnect fails, then
// ensure reconnect succeeds when the client provides a cert.
func TestConfigReloadRotateTLS(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	if err := os.Symlink("./configs/tls_test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)

	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}

	server := RunServer(opts)
	defer server.Shutdown()

	// Ensure we can connect as a sanity check.
	addr := fmt.Sprintf("nats://%s:%d", opts.Host, opts.Port)
	nc, err := nats.Connect(addr, nats.Secure())
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()

	// Rotate cert and enable client verification.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/tls_verify_test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	if err := server.Reload(); err != nil {
		t.Fatalf("Error reloading config: %v", err)
	}

	// Ensure connecting fails.
	if _, err := nats.Connect(addr, nats.Secure()); err == nil {
		t.Fatal("Expected connect to fail")
	}

	// Ensure connecting succeeds when client presents cert.
	cert := nats.ClientCert("./configs/certs/cert.new.pem", "./configs/certs/key.new.pem")
	nc, err = nats.Connect(addr, cert, nats.RootCAs("./configs/certs/cert.new.pem"))
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()
}

// Ensure Reload supports enabling TLS. Test this by starting a server without
// TLS enabled, connect to it to verify, reload config with TLS enabled, ensure
// reconnect fails, then ensure reconnect succeeds when using secure.
func TestConfigReloadEnableTLS(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	if err := os.Symlink("./configs/basic.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)

	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}

	server := RunServer(opts)
	defer server.Shutdown()

	// Ensure we can connect as a sanity check.
	addr := fmt.Sprintf("nats://%s:%d", opts.Host, opts.Port)
	nc, err := nats.Connect(addr)
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()

	// Enable TLS.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/tls_test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	if err := server.Reload(); err != nil {
		t.Fatalf("Error reloading config: %v", err)
	}

	// Ensure connecting fails.
	if _, err := nats.Connect(addr); err == nil {
		t.Fatal("Expected connect to fail")
	}

	// Ensure connecting succeeds when using secure.
	nc, err = nats.Connect(addr, nats.Secure())
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()
}

// Ensure Reload supports disabling TLS. Test this by starting a server with
// TLS enabled, connect to it to verify, reload config with TLS disabled,
// ensure reconnect fails, then ensure reconnect succeeds when connecting
// without secure.
func TestConfigReloadDisableTLS(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	if err := os.Symlink("./configs/tls_test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)

	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}

	server := RunServer(opts)
	defer server.Shutdown()

	// Ensure we can connect as a sanity check.
	addr := fmt.Sprintf("nats://%s:%d", opts.Host, opts.Port)
	nc, err := nats.Connect(addr, nats.Secure())
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()

	// Disable TLS.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/basic.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	if err := server.Reload(); err != nil {
		t.Fatalf("Error reloading config: %v", err)
	}

	// Ensure connecting fails.
	if _, err := nats.Connect(addr, nats.Secure()); err == nil {
		t.Fatal("Expected connect to fail")
	}

	// Ensure connecting succeeds when not using secure.
	nc, err = nats.Connect(addr)
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()
}
<commit_msg>Ensure old connection still works in TLS reload test<commit_after>// Copyright 2017 Apcera Inc. All rights reserved.

package server

import (
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"testing"
	"time"

	"github.com/nats-io/go-nats"
)

// Ensure Reload returns an error when attempting to reload a server that did
// not start with a config file.
func TestConfigReloadNoConfigFile(t *testing.T) {
	server := New(&Options{})
	if server.Reload() == nil {
		t.Fatal("Expected Reload to return an error")
	}
}

// Ensure Reload returns an error when attempting to change an option which
// does not support reloading.
func TestConfigReloadUnsupported(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	golden := &Options{
		ConfigFile:     config,
		Host:           "0.0.0.0",
		Port:           4222,
		AuthTimeout:    1.0,
		Debug:          false,
		Trace:          false,
		Logtime:        false,
		MaxControlLine: 1024,
		MaxPayload:     1048576,
		MaxConn:        65536,
		PingInterval:   2 * time.Minute,
		MaxPingsOut:    2,
		WriteDeadline:  2 * time.Second,
	}
	processOptions(golden)

	if err := os.Symlink("./configs/reload/test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v (ensure you have privileges)", err)
	}
	defer os.Remove(config)
	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}
	server := New(opts)

	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}

	// Change config file to bad config by replacing symlink.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/reload/reload_unsupported.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v (ensure you have privileges)", err)
	}

	// This should fail because `debug` cannot be changed.
	if err := server.Reload(); err == nil {
		t.Fatal("Expected Reload to return an error")
	}

	// Ensure config didn't change.
	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}
}

// Ensure Reload returns an error when reloading from a bad config file.
func TestConfigReloadInvalidConfig(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	golden := &Options{
		ConfigFile:     config,
		Host:           "0.0.0.0",
		Port:           4222,
		AuthTimeout:    1.0,
		Debug:          false,
		Trace:          false,
		Logtime:        false,
		MaxControlLine: 1024,
		MaxPayload:     1048576,
		MaxConn:        65536,
		PingInterval:   2 * time.Minute,
		MaxPingsOut:    2,
		WriteDeadline:  2 * time.Second,
	}
	processOptions(golden)

	if err := os.Symlink("./configs/reload/test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)
	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}
	server := New(opts)

	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}

	// Change config file to bad config by replacing symlink.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/reload/invalid.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}

	// This should fail because the new config should not parse.
	if err := server.Reload(); err == nil {
		t.Fatal("Expected Reload to return an error")
	}

	// Ensure config didn't change.
	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}
}

// Ensure Reload returns nil and the config is changed on success.
func TestConfigReload(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	golden := &Options{
		ConfigFile:     config,
		Host:           "0.0.0.0",
		Port:           4222,
		AuthTimeout:    1.0,
		Debug:          false,
		Trace:          false,
		Logtime:        false,
		MaxControlLine: 1024,
		MaxPayload:     1048576,
		MaxConn:        65536,
		PingInterval:   2 * time.Minute,
		MaxPingsOut:    2,
		WriteDeadline:  2 * time.Second,
	}
	processOptions(golden)

	if err := os.Symlink("./configs/reload/test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)
	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}
	server := New(opts)

	if !reflect.DeepEqual(golden, server.getOpts()) {
		t.Fatalf("Options are incorrect.\nexpected: %+v\ngot: %+v",
			golden, opts)
	}

	// Change config file to new config by replacing symlink.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/reload/reload.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}

	if err := server.Reload(); err != nil {
		t.Fatalf("Error reloading config: %v", err)
	}

	// Ensure config changed.
	updated := server.getOpts()
	if !updated.Trace {
		t.Fatal("Expected Trace to be true")
	}
	if !updated.Debug {
		t.Fatal("Expected Debug to be true")
	}
	if updated.TLSConfig == nil {
		t.Fatal("Expected TLSConfig to be non-nil")
	}
	if !server.info.TLSRequired {
		t.Fatal("Expected TLSRequired to be true")
	}
	if !server.info.TLSVerify {
		t.Fatal("Expected TLSVerify to be true")
	}
}

// Ensure Reload supports TLS config changes. Test this by starting a server
// with TLS enabled, connect to it to verify, reload config using a different
// key pair and client verification enabled, ensure reconnect fails, then
// ensure reconnect succeeds when the client provides a cert.
func TestConfigReloadRotateTLS(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	if err := os.Symlink("./configs/tls_test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)

	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}

	server := RunServer(opts)
	defer server.Shutdown()

	// Ensure we can connect as a sanity check.
	addr := fmt.Sprintf("nats://%s:%d", opts.Host, opts.Port)
	nc, err := nats.Connect(addr, nats.Secure())
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	defer nc.Close()
	sub, err := nc.SubscribeSync("foo")
	if err != nil {
		t.Fatalf("Error subscribing: %v", err)
	}
	defer sub.Unsubscribe()

	// Rotate cert and enable client verification.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/tls_verify_test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	if err := server.Reload(); err != nil {
		t.Fatalf("Error reloading config: %v", err)
	}

	// Ensure connecting fails.
	if _, err := nats.Connect(addr, nats.Secure()); err == nil {
		t.Fatal("Expected connect to fail")
	}

	// Ensure connecting succeeds when client presents cert.
	cert := nats.ClientCert("./configs/certs/cert.new.pem", "./configs/certs/key.new.pem")
	conn, err := nats.Connect(addr, cert, nats.RootCAs("./configs/certs/cert.new.pem"))
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	conn.Close()

	// Ensure the original connection can still publish/receive.
	if err := nc.Publish("foo", []byte("hello")); err != nil {
		t.Fatalf("Error publishing: %v", err)
	}
	msg, err := sub.NextMsg(time.Second)
	if err != nil {
		t.Fatalf("Error receiving msg: %v", err)
	}
	if string(msg.Data) != "hello" {
		t.Fatalf("Msg is incorrect.\nexpected: %+v\ngot: %+v", []byte("hello"), msg.Data)
	}
}

// Ensure Reload supports enabling TLS. Test this by starting a server without
// TLS enabled, connect to it to verify, reload config with TLS enabled, ensure
// reconnect fails, then ensure reconnect succeeds when using secure.
func TestConfigReloadEnableTLS(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	if err := os.Symlink("./configs/basic.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)

	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}

	server := RunServer(opts)
	defer server.Shutdown()

	// Ensure we can connect as a sanity check.
	addr := fmt.Sprintf("nats://%s:%d", opts.Host, opts.Port)
	nc, err := nats.Connect(addr)
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()

	// Enable TLS.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/tls_test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	if err := server.Reload(); err != nil {
		t.Fatalf("Error reloading config: %v", err)
	}

	// Ensure connecting fails.
	if _, err := nats.Connect(addr); err == nil {
		t.Fatal("Expected connect to fail")
	}

	// Ensure connecting succeeds when using secure.
	nc, err = nats.Connect(addr, nats.Secure())
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()
}

// Ensure Reload supports disabling TLS. Test this by starting a server with
// TLS enabled, connect to it to verify, reload config with TLS disabled,
// ensure reconnect fails, then ensure reconnect succeeds when connecting
// without secure.
func TestConfigReloadDisableTLS(t *testing.T) {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Error getting working directory: %v", err)
	}
	config := filepath.Join(dir, "tmp.conf")

	if err := os.Symlink("./configs/tls_test.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	defer os.Remove(config)

	opts, err := ProcessConfigFile(config)
	if err != nil {
		t.Fatalf("Error processing config file: %v", err)
	}

	server := RunServer(opts)
	defer server.Shutdown()

	// Ensure we can connect as a sanity check.
	addr := fmt.Sprintf("nats://%s:%d", opts.Host, opts.Port)
	nc, err := nats.Connect(addr, nats.Secure())
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()

	// Disable TLS.
	if err := os.Remove(config); err != nil {
		t.Fatalf("Error deleting symlink: %v", err)
	}
	if err := os.Symlink("./configs/basic.conf", config); err != nil {
		t.Fatalf("Error creating symlink: %v", err)
	}
	if err := server.Reload(); err != nil {
		t.Fatalf("Error reloading config: %v", err)
	}

	// Ensure connecting fails.
	if _, err := nats.Connect(addr, nats.Secure()); err == nil {
		t.Fatal("Expected connect to fail")
	}

	// Ensure connecting succeeds when not using secure.
	nc, err = nats.Connect(addr)
	if err != nil {
		t.Fatalf("Error creating client: %v", err)
	}
	nc.Close()
}
<|endoftext|>


<commit_before>package syntaxcolor

import (
	"fmt"
	"path"
	"regexp"

	"github.com/nsf/termbox-go"
)

// Color specifies a foreground and background color.
type Color struct {
	fg, bg termbox.Attribute
}

// SyntaxRule simply maps a regexp to a color.
type SyntaxRule struct {
	re    *regexp.Regexp
	color Color
}

// SyntaxRules is the full collection of all syntax rules.
type SyntaxRules []SyntaxRule

// NewSyntaxRules creates a new SyntaxRules object, and initializes
// the syntax rule sets.
func NewSyntaxRules(filename string) *SyntaxRules {
	rules := SyntaxRules{}
	filetype := rules.GetFileType(filename)
	rules.addQuoteRules(filetype)
	rules.addCommentRules(filetype)
	rules.addMiscRules(filetype)
	rules.addWhitespaceRules(filetype)
	return &rules
}

func (rules *SyntaxRules) addQuoteRules(filetype string) {
	switch filetype {
	case "sh", "py", "yaml", "js", "c", "coffee", "html", "rb":
		rules.addSingleQuoteRule(termbox.ColorYellow)
		rules.addDoubleQuoteRule(termbox.ColorYellow)
	case "go":
		rules.addSingleQuoteRule(termbox.ColorRed)
		rules.addDoubleQuoteRule(termbox.ColorYellow)
		rules.addBackticQuoteRule(termbox.ColorYellow)
	}
}

func (rules *SyntaxRules) addCommentRules(filetype string) {
	switch filetype {
	case "go":
		rules.addLineCommentRule("//", termbox.ColorCyan)
	case "rb", "sh", "py", "yaml", "coffee", "git":
		rules.addLineCommentRule("#", termbox.ColorCyan)
	case "c":
		rules.addLineCommentRule("//", termbox.ColorCyan)
		rules.addRule(`/\*.*?\*/`, Color{fg: termbox.ColorCyan})
	case "js":
		rules.addLineCommentRule("//", termbox.ColorCyan)
	case "html":
		rules.addRule("<!--.*?-->", Color{fg: termbox.ColorCyan})
	}
}

func (rules *SyntaxRules) addMiscRules(filetype string) {
	switch filetype {
	case "md":
		rules.addRule("^#+.*$", Color{fg: termbox.ColorGreen})
		rules.addRule("^===*$", Color{fg: termbox.ColorGreen})
		rules.addRule("^---*$", Color{fg: termbox.ColorGreen})
	}
}

func (rules *SyntaxRules) addWhitespaceRules(filetype string) {
	rules.addRule("[ \t]+$", Color{bg: termbox.ColorYellow})
}

func (rules *SyntaxRules) addRule(reStr string, color Color) {
	re, _ := regexp.Compile(reStr)
	*rules = append(*rules, SyntaxRule{re, color})
}

func (rules *SyntaxRules) addLineCommentRule(commStr string, fg termbox.Attribute) {
	reStr := fmt.Sprintf("%s.*$", commStr)
	rules.addRule(reStr, Color{fg: fg})
}

func (rules *SyntaxRules) addBackticQuoteRule(fg termbox.Attribute) {
	rules.addRule("`.*?`", Color{fg: fg})
}

func (rules *SyntaxRules) addSingleQuoteRule(fg termbox.Attribute) {
	rules.addRule("'.*?'", Color{fg: fg})
}

func (rules *SyntaxRules) addDoubleQuoteRule(fg termbox.Attribute) {
	rules.addRule(`".*?"`, Color{fg: fg})
}

// GetFileType maps common extensions onto a type (e.g. c, C, cpp, etc
// onto 'c'), but otherwise just returns the extension.
func (rules SyntaxRules) GetFileType(filename string) string {
	ext := rules.getFileExt(filename)
	switch ext {
	case "c", "C", "cpp", "c++", "h", "hpp", "cc", "hh":
		return "c"
	case "sh", "csh":
		return "sh"
	case "":
		return ""
	default:
		return ext
	}
}

func (rules SyntaxRules) getFileExt(filename string) string {
	ext := path.Ext(filename)
	if len(ext) == 0 {
		basename := path.Base(filename)
		if basename == "COMMIT_EDITMSG" {
			return "git"
		}
		return ""
	}
	ext = ext[1:]
	return ext
}

// LineColor object maps a color pair (bg, fg) with start/end indices
// within a string.
type LineColor struct {
	Fg, Bg     termbox.Attribute
	Start, End int
}

// NextMatch finds the next match accross all the rules.
func (rules SyntaxRules) NextMatch(str string, idx int) (LineColor, error) {
	subStr := str[idx:]
	i0 := len(subStr) + 1
	lc := LineColor{}
	for _, rule := range rules {
		match := rule.re.FindStringIndex(subStr)
		if match == nil {
			continue
		}
		if match[0] < i0 {
			i0 = match[0]
			lc.Fg, lc.Bg = rule.color.fg, rule.color.bg
			lc.Start, lc.End = idx+match[0], idx+match[1]
		}
	}
	if i0 < len(subStr)+1 {
		return lc, nil
	}
	return lc, fmt.Errorf("no match")
}

// Colorize takes in a string and outputs an array of LineColor objects.
func (rules SyntaxRules) Colorize(str string) []LineColor {
	lcs := []LineColor{}
	idx := 0
	for {
		lc, err := rules.NextMatch(str, idx)
		if err != nil {
			break
		}
		lcs = append(lcs, lc)
		idx = lc.End
	}
	return lcs
}
<commit_msg>make trailing whitespace syntax coloring always display<commit_after>package syntaxcolor

import (
	"fmt"
	"path"
	"regexp"

	"github.com/nsf/termbox-go"
)

// Color specifies a foreground and background color.
type Color struct {
	fg, bg termbox.Attribute
}

// SyntaxRule simply maps a regexp to a color.
type SyntaxRule struct {
	re    *regexp.Regexp
	color Color
}

// SyntaxRules is the full collection of all syntax rules.
type SyntaxRules struct {
	list       []SyntaxRule
	whitespace *regexp.Regexp
}

// NewSyntaxRules creates a new SyntaxRules object, and initializes
// the syntax rule sets.
func NewSyntaxRules(filename string) *SyntaxRules {
	rules := SyntaxRules{}
	filetype := rules.GetFileType(filename)
	rules.addQuoteRules(filetype)
	rules.addCommentRules(filetype)
	rules.addMiscRules(filetype)
	rules.addWhitespaceRule()
	return &rules
}

func (rules *SyntaxRules) addWhitespaceRule() {
	rules.whitespace = regexp.MustCompile("[ \t]+$")
}

func (rules *SyntaxRules) addQuoteRules(filetype string) {
	switch filetype {
	case "sh", "py", "yaml", "js", "c", "coffee", "html", "rb":
		rules.addSingleQuoteRule(termbox.ColorYellow)
		rules.addDoubleQuoteRule(termbox.ColorYellow)
	case "go":
		rules.addSingleQuoteRule(termbox.ColorRed)
		rules.addDoubleQuoteRule(termbox.ColorYellow)
		rules.addBackticQuoteRule(termbox.ColorYellow)
	}
}

func (rules *SyntaxRules) addCommentRules(filetype string) {
	switch filetype {
	case "go":
		rules.addLineCommentRule("//", termbox.ColorCyan)
	case "rb", "sh", "py", "yaml", "coffee", "git":
		rules.addLineCommentRule("#", termbox.ColorCyan)
	case "c":
		rules.addLineCommentRule("//", termbox.ColorCyan)
		rules.addRule(`/\*.*?\*/`, Color{fg: termbox.ColorCyan})
	case "js":
		rules.addLineCommentRule("//", termbox.ColorCyan)
	case "html":
		rules.addRule("<!--.*?-->", Color{fg: termbox.ColorCyan})
	}
}

func (rules *SyntaxRules) addMiscRules(filetype string) {
	switch filetype {
	case "md":
		rules.addRule("^#+.*$", Color{fg: termbox.ColorGreen})
		rules.addRule("^===*$", Color{fg: termbox.ColorGreen})
		rules.addRule("^---*$", Color{fg: termbox.ColorGreen})
	}
}

func (rules *SyntaxRules) addRule(reStr string, color Color) {
	re, _ := regexp.Compile(reStr)
	rules.list = append(rules.list, SyntaxRule{re, color})
}

func (rules *SyntaxRules) addLineCommentRule(commStr string, fg termbox.Attribute) {
	reStr := fmt.Sprintf("%s.*$", commStr)
	rules.addRule(reStr, Color{fg: fg})
}

func (rules *SyntaxRules) addBackticQuoteRule(fg termbox.Attribute) {
	rules.addRule("`.*?`", Color{fg: fg})
}

func (rules *SyntaxRules) addSingleQuoteRule(fg termbox.Attribute) {
	rules.addRule("'.*?'", Color{fg: fg})
}

func (rules *SyntaxRules) addDoubleQuoteRule(fg termbox.Attribute) {
	rules.addRule(`".*?"`, Color{fg: fg})
}

// GetFileType maps common extensions onto a type (e.g. c, C, cpp, etc
// onto 'c'), but otherwise just returns the extension.
func (rules SyntaxRules) GetFileType(filename string) string {
	ext := rules.getFileExt(filename)
	switch ext {
	case "c", "C", "cpp", "c++", "h", "hpp", "cc", "hh":
		return "c"
	case "sh", "csh":
		return "sh"
	case "":
		return ""
	default:
		return ext
	}
}

func (rules SyntaxRules) getFileExt(filename string) string {
	ext := path.Ext(filename)
	if len(ext) == 0 {
		basename := path.Base(filename)
		if basename == "COMMIT_EDITMSG" {
			return "git"
		}
		return ""
	}
	ext = ext[1:]
	return ext
}

// LineColor object maps a color pair (bg, fg) with start/end indices
// within a string.
type LineColor struct {
	Fg, Bg     termbox.Attribute
	Start, End int
}

// NextMatch finds the next match accross all the rules.
func (rules SyntaxRules) NextMatch(str string, idx int) (LineColor, error) {
	subStr := str[idx:]
	i0 := len(subStr) + 1
	lc := LineColor{}
	for _, rule := range rules.list {
		match := rule.re.FindStringIndex(subStr)
		if match == nil {
			continue
		}
		if match[0] < i0 {
			i0 = match[0]
			lc.Fg, lc.Bg = rule.color.fg, rule.color.bg
			lc.Start, lc.End = idx+match[0], idx+match[1]
		}
	}
	if i0 < len(subStr)+1 {
		return lc, nil
	}
	return lc, fmt.Errorf("no match")
}

// Colorize takes in a string and outputs an array of LineColor objects.
func (rules SyntaxRules) Colorize(str string) []LineColor {
	lcs := []LineColor{}
	idx := 0
	for {
		lc, err := rules.NextMatch(str, idx)
		if err != nil {
			break
		}
		lcs = append(lcs, lc)
		idx = lc.End
	}
	match := rules.whitespace.FindStringIndex(str)
	if match != nil {
		lc := LineColor{}
		lc.Bg = termbox.ColorYellow
		lc.Start, lc.End = match[0], match[1]
		lcs = append(lcs, lc)
	}
	return lcs
}

func (rules *SyntaxRules) addWhitespaceRules(filetype string) {
	rules.addRule("[ \t]+$", Color{bg: termbox.ColorYellow})
}
<|endoftext|>


<commit_before>package goenvy

import (
	"errors"
	"fmt"
	"os"
	"strconv"
)

var (
	ErrStringEnvNotDefined = errors.New("empty string value")
)

// Interface that refers to the environment
// in which config values are stored
type Env interface {
	// return a string from the environment
	// returns an error if the the value is not available
	GetString(string) (string, error)

	// return an integer from the environment
	// returns an error if the value is not available or
	// not the correct type
	GetInt(string) (int, error)

	// return an integer from the environment
	// returns an error if the value is not available or
	// not the correct type
	GetBool(string) (bool, error)
}

type SimpleEnv interface {
	Get(string) string
}

type ParsingEnv struct {
	SimpleEnv
}

// returns a string from the underlying env
// TODO: determine what exactly the error means
func (p *ParsingEnv) GetString(key string) (string, error) {
	value := p.Get(key)
	if value == "" {
		return "", ErrStringEnvNotDefined
	}
	return value, nil
}

// returns an int from the underlying env
// error is returned when it is not a valid integer
func (p *ParsingEnv) GetInt(key string) (int, error) {
	value := p.Get(key)
	return strconv.Atoi(value)
}

// returns an int from the underlying env
// error is returned when it is not a valid bool
// Expects either "true" or "false"
func (p *ParsingEnv) GetBool(key string) (bool, error) {
	value := p.Get(key)
	switch value {
	case "true":
		return true, nil
	case "false":
		return false, nil
	default:
		return false, fmt.Errorf("bool string %q not parsable", value)
	}
}

// Simple type to wrap an existing Env implementation,
// and call all methods with a prefix
// Note: this is very useful for the tests
type PrefixEnv struct {
	prefix string
	Env
}

func (p *PrefixEnv) GetString(key string) (string, error) {
	return p.Env.GetString(p.prefix + key)
}

func (p *PrefixEnv) GetInt(key string) (int, error) {
	return p.Env.GetInt(p.prefix + key)
}

func (p *PrefixEnv) GetBool(key string) (int, error) {
	return p.Env.GetBool(p.prefix + key)
}

type OsEnvironParser struct{}

func (o *OsEnvironParser) Get(key string) string {
	return os.Getenv(key)
}
<commit_msg>wrong signature; should run tests first next time<commit_after>package goenvy

import (
	"errors"
	"fmt"
	"os"
	"strconv"
)

var (
	ErrStringEnvNotDefined = errors.New("empty string value")
)

// Interface that refers to the environment
// in which config values are stored
type Env interface {
	// return a string from the environment
	// returns an error if the the value is not available
	GetString(string) (string, error)

	// return an integer from the environment
	// returns an error if the value is not available or
	// not the correct type
	GetInt(string) (int, error)

	// return an integer from the environment
	// returns an error if the value is not available or
	// not the correct type
	GetBool(string) (bool, error)
}

type SimpleEnv interface {
	Get(string) string
}

type ParsingEnv struct {
	SimpleEnv
}

// returns a string from the underlying env
// TODO: determine what exactly the error means
func (p *ParsingEnv) GetString(key string) (string, error) {
	value := p.Get(key)
	if value == "" {
		return "", ErrStringEnvNotDefined
	}
	return value, nil
}

// returns an int from the underlying env
// error is returned when it is not a valid integer
func (p *ParsingEnv) GetInt(key string) (int, error) {
	value := p.Get(key)
	return strconv.Atoi(value)
}

// returns an int from the underlying env
// error is returned when it is not a valid bool
// Expects either "true" or "false"
func (p *ParsingEnv) GetBool(key string) (bool, error) {
	value := p.Get(key)
	switch value {
	case "true":
		return true, nil
	case "false":
		return false, nil
	default:
		return false, fmt.Errorf("bool string %q not parsable", value)
	}
}

// Simple type to wrap an existing Env implementation,
// and call all methods with a prefix
// Note: this is very useful for the tests
type PrefixEnv struct {
	prefix string
	Env
}

func (p *PrefixEnv) GetString(key string) (string, error) {
	return p.Env.GetString(p.prefix + key)
}

func (p *PrefixEnv) GetInt(key string) (int, error) {
	return p.Env.GetInt(p.prefix + key)
}

func (p *PrefixEnv) GetBool(key string) (bool, error) {
	return p.Env.GetBool(p.prefix + key)
}

type OsEnvironParser struct{}

func (o *OsEnvironParser) Get(key string) string {
	return os.Getenv(key)
}
<|endoftext|>