<commit_before>package dry

import (
	"bytes"
	"crypto/md5"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"
)

func StringPrettifyJSON(compactJSON string) string {
	var buf bytes.Buffer
	err := json.Indent(&buf, []byte(compactJSON), "", "\t")
	if err != nil {
		return err.Error()
	}
	return buf.String()
}

func StringEscapeJSON(jsonString string) string {
	jsonString = strings.Replace(jsonString, `\`, `\\`, -1)
	jsonString = strings.Replace(jsonString, `"`, `\"`, -1)
	return jsonString
}

func StringStripHTMLTags(text string) (plainText string) {
	chars := []byte(text)
	tagStart := -1
	for i := 0; i < len(chars); i++ {
		if chars[i] == '<' {
			tagStart = i
		} else if chars[i] == '>' && tagStart != -1 {
			chars = append(chars[:tagStart], chars[i+1:]...)
			i, tagStart = tagStart-1, -1
		}
	}
	return string(chars)
}

func StringMD5(data string) string {
	hash := md5.New()
	hash.Write([]byte(data))
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func StringAddURLParam(url, name, value string) string {
	var separator string
	if strings.IndexRune(url, '?') == -1 {
		separator = "?"
	} else {
		separator = "&"
	}
	return url + separator + name + "=" + value
}

func StringConvertTime(timeString, formatIn, formatOut string) (resultTime string, err error) {
	if timeString == "" {
		return "", nil
	}
	t, err := time.Parse(formatIn, timeString)
	if err != nil {
		return "", err
	}
	return t.Format(formatOut), nil
}

func StringCSV(records [][]string) string {
	var buf bytes.Buffer
	writer := csv.NewWriter(&buf)
	err := writer.WriteAll(records)
	if err != nil {
		return ""
	}
	return buf.String()
}

func StringToInt(s string) int {
	i, _ := strconv.ParseInt(s, 10, 64)
	return int(i)
}

func StringToFloat(s string) float64 {
	f, _ := strconv.ParseFloat(s, 64)
	return f
}

func StringToBool(s string) bool {
	b, _ := strconv.ParseBool(s)
	return b
}

func StringInSlice(s string, slice []string) bool {
	for i := range slice {
		if slice[i] == s {
			return true
		}
	}
	return false
}

// StringJoinFormat formats every value in values with format
// and joins the result with sep as separator.
// values must be a slice of a formatable type
func StringJoinFormat(format string, values interface{}, sep string) string {
	v := reflect.ValueOf(values)
	if v.Kind() != reflect.Slice {
		panic("values is not a slice")
	}
	var buffer bytes.Buffer
	for i := 0; i < v.Len(); i++ {
		if i > 0 {
			buffer.WriteString(sep)
		}
		buffer.WriteString(fmt.Sprintf(format, v.Index(i).Interface()))
	}
	return buffer.String()
}

// StringJoin formats every value in values according to its default formatting
// and joins the result with sep as separator.
// values must be a slice of a formatable type
func StringJoin(values interface{}, sep string) string {
	v := reflect.ValueOf(values)
	if v.Kind() != reflect.Slice {
		panic("values is not a slice")
	}
	var buffer bytes.Buffer
	for i := 0; i < v.Len(); i++ {
		if i > 0 {
			buffer.WriteString(sep)
		}
		buffer.WriteString(fmt.Sprint(v.Index(i).Interface()))
	}
	return buffer.String()
}

func StringFormatMemory(mem uint64) string {
	switch {
	case mem >= 10e12:
		return fmt.Sprintf("%dTB", mem/1e12)
	case mem >= 1e12:
		return fmt.Sprintf("%.1fTB", float64(mem)/1e12)

	case mem >= 10e9:
		return fmt.Sprintf("%dGB", mem/1e9)
	case mem >= 1e9:
		return fmt.Sprintf("%.1fGB", float64(mem)/1e9)

	case mem >= 10e6:
		return fmt.Sprintf("%dMB", mem/1e6)
	case mem >= 1e6:
		return fmt.Sprintf("%.1fMB", float64(mem)/1e6)

	case mem >= 10e3:
		return fmt.Sprintf("%dkB", mem/1e3)
	case mem >= 1e3:
		return fmt.Sprintf("%.1fkB", float64(mem)/1e3)
	}
	return fmt.Sprintf("%dB", mem)
}
<commit_msg>added StringFormatBigInt<commit_after>package dry

import (
	"bytes"
	"crypto/md5"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"
)

func StringPrettifyJSON(compactJSON string) string {
	var buf bytes.Buffer
	err := json.Indent(&buf, []byte(compactJSON), "", "\t")
	if err != nil {
		return err.Error()
	}
	return buf.String()
}

func StringEscapeJSON(jsonString string) string {
	jsonString = strings.Replace(jsonString, `\`, `\\`, -1)
	jsonString = strings.Replace(jsonString, `"`, `\"`, -1)
	return jsonString
}

func StringStripHTMLTags(text string) (plainText string) {
	chars := []byte(text)
	tagStart := -1
	for i := 0; i < len(chars); i++ {
		if chars[i] == '<' {
			tagStart = i
		} else if chars[i] == '>' && tagStart != -1 {
			chars = append(chars[:tagStart], chars[i+1:]...)
			i, tagStart = tagStart-1, -1
		}
	}
	return string(chars)
}

func StringMD5(data string) string {
	hash := md5.New()
	hash.Write([]byte(data))
	return fmt.Sprintf("%x", hash.Sum(nil))
}

func StringAddURLParam(url, name, value string) string {
	var separator string
	if strings.IndexRune(url, '?') == -1 {
		separator = "?"
	} else {
		separator = "&"
	}
	return url + separator + name + "=" + value
}

func StringConvertTime(timeString, formatIn, formatOut string) (resultTime string, err error) {
	if timeString == "" {
		return "", nil
	}
	t, err := time.Parse(formatIn, timeString)
	if err != nil {
		return "", err
	}
	return t.Format(formatOut), nil
}

func StringCSV(records [][]string) string {
	var buf bytes.Buffer
	writer := csv.NewWriter(&buf)
	err := writer.WriteAll(records)
	if err != nil {
		return ""
	}
	return buf.String()
}

func StringToInt(s string) int {
	i, _ := strconv.ParseInt(s, 10, 64)
	return int(i)
}

func StringToFloat(s string) float64 {
	f, _ := strconv.ParseFloat(s, 64)
	return f
}

func StringToBool(s string) bool {
	b, _ := strconv.ParseBool(s)
	return b
}

func StringInSlice(s string, slice []string) bool {
	for i := range slice {
		if slice[i] == s {
			return true
		}
	}
	return false
}

// StringJoinFormat formats every value in values with format
// and joins the result with sep as separator.
// values must be a slice of a formatable type
func StringJoinFormat(format string, values interface{}, sep string) string {
	v := reflect.ValueOf(values)
	if v.Kind() != reflect.Slice {
		panic("values is not a slice")
	}
	var buffer bytes.Buffer
	for i := 0; i < v.Len(); i++ {
		if i > 0 {
			buffer.WriteString(sep)
		}
		buffer.WriteString(fmt.Sprintf(format, v.Index(i).Interface()))
	}
	return buffer.String()
}

// StringJoin formats every value in values according to its default formatting
// and joins the result with sep as separator.
// values must be a slice of a formatable type
func StringJoin(values interface{}, sep string) string {
	v := reflect.ValueOf(values)
	if v.Kind() != reflect.Slice {
		panic("values is not a slice")
	}
	var buffer bytes.Buffer
	for i := 0; i < v.Len(); i++ {
		if i > 0 {
			buffer.WriteString(sep)
		}
		buffer.WriteString(fmt.Sprint(v.Index(i).Interface()))
	}
	return buffer.String()
}

func StringFormatBigInt(mem uint64) string {
	switch {
	case mem >= 10e12:
		return fmt.Sprintf("%dT", mem/1e12)
	case mem >= 1e12:
		return fmt.Sprintf("%.1fT", float64(mem)/1e12)

	case mem >= 10e9:
		return fmt.Sprintf("%dG", mem/1e9)
	case mem >= 1e9:
		return fmt.Sprintf("%.1fG", float64(mem)/1e9)

	case mem >= 10e6:
		return fmt.Sprintf("%dM", mem/1e6)
	case mem >= 1e6:
		return fmt.Sprintf("%.1fM", float64(mem)/1e6)

	case mem >= 10e3:
		return fmt.Sprintf("%dk", mem/1e3)
	case mem >= 1e3:
		return fmt.Sprintf("%.1fk", float64(mem)/1e3)
	}
	return fmt.Sprintf("%d", mem)
}

func StringFormatMemory(mem uint64) string {
	return StringFormatBigInt(mem) + "B"
}
<|endoftext|>


<commit_before>package maxminddb

import (
	"errors"
	"fmt"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"math/big"
	"math/rand"
	"net"
	"testing"
	"time"
)

func TestMaxMindDbReader(t *testing.T) { TestingT(t) }

type MySuite struct{}

var _ = Suite(&MySuite{})

func (s *MySuite) TestReader(c *C) {
	for _, recordSize := range []uint{24, 28, 32} {
		for _, ipVersion := range []uint{4, 6} {
			fileName := fmt.Sprintf("test-data/test-data/MaxMind-DB-test-ipv%d-%d.mmdb", ipVersion, recordSize)
			reader, _ := Open(fileName)

			checkMetadata(c, reader, ipVersion, recordSize)

			if ipVersion == 4 {
				checkIpv4(c, reader)
			} else {
				checkIpv6(c, reader)
			}
		}
	}
}

func (s *MySuite) TestReaderBytes(c *C) {
	for _, recordSize := range []uint{24, 28, 32} {
		for _, ipVersion := range []uint{4, 6} {
			fileName := fmt.Sprintf("test-data/test-data/MaxMind-DB-test-ipv%d-%d.mmdb", ipVersion, recordSize)
			bytes, _ := ioutil.ReadFile(fileName)
			reader, _ := FromBytes(bytes)

			checkMetadata(c, reader, ipVersion, recordSize)

			if ipVersion == 4 {
				checkIpv4(c, reader)
			} else {
				checkIpv6(c, reader)
			}
		}
	}
}

func (s *MySuite) TestDecodingToInterface(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-decoder.mmdb")

	var recordInterface interface{}
	reader.Lookup(net.ParseIP("::1.1.1.0"), &recordInterface)
	record := recordInterface.(map[string]interface{})

	c.Assert(record["array"], DeepEquals, []interface{}{uint64(1), uint64(2), uint64(3)})
	c.Assert(record["boolean"], Equals, true)
	c.Assert(record["bytes"], DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
	c.Assert(record["double"], Equals, 42.123456)
	c.Assert(record["float"], Equals, float32(1.1))
	c.Assert(record["int32"], Equals, -268435456)
	c.Assert(record["map"], DeepEquals,
		map[string]interface{}{
			"mapX": map[string]interface{}{
				"arrayX":       []interface{}{uint64(7), uint64(8), uint64(9)},
				"utf8_stringX": "hello",
			}})

	c.Assert(record["uint16"], Equals, uint64(100))
	c.Assert(record["uint32"], Equals, uint64(268435456))
	c.Assert(record["uint64"], Equals, uint64(1152921504606846976))
	c.Assert(record["utf8_string"], Equals, "unicode! ☯ - ♫")
	bigInt := new(big.Int)
	bigInt.SetString("1329227995784915872903807060280344576", 10)
	c.Assert(record["uint128"], DeepEquals, bigInt)
}

type TestType struct {
	Array      []uint                 `maxminddb:"array"`
	Boolean    bool                   `maxminddb:"boolean"`
	Bytes      []byte                 `maxminddb:"bytes"`
	Double     float64                `maxminddb:"double"`
	Float      float32                `maxminddb:"float"`
	Int32      int32                  `maxminddb:"int32"`
	Map        map[string]interface{} `maxminddb:"map"`
	Uint16     uint16                 `maxminddb:"uint16"`
	Uint32     uint32                 `maxminddb:"uint32"`
	Uint64     uint64                 `maxminddb:"uint64"`
	Uint128    big.Int                `maxminddb:"uint128"`
	Utf8String string                 `maxminddb:"utf8_string"`
}

func (s *MySuite) TestDecoder(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-decoder.mmdb")

	var result TestType
	err := reader.Lookup(net.ParseIP("::1.1.1.0"), &result)
	if err != nil {
		c.Log(err)
		c.Fail()
	}

	c.Assert(result.Array, DeepEquals, []uint{uint(1), uint(2), uint(3)})
	c.Assert(result.Boolean, Equals, true)
	c.Assert(result.Bytes, DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
	c.Assert(result.Double, Equals, 42.123456)
	c.Assert(result.Float, Equals, float32(1.1))
	c.Assert(result.Int32, Equals, int32(-268435456))

	c.Assert(result.Map, DeepEquals,
		map[string]interface{}{
			"mapX": map[string]interface{}{
				"arrayX":       []interface{}{uint64(7), uint64(8), uint64(9)},
				"utf8_stringX": "hello",
			}})

	c.Assert(result.Uint16, Equals, uint16(100))
	c.Assert(result.Uint32, Equals, uint32(268435456))
	c.Assert(result.Uint64, Equals, uint64(1152921504606846976))
	c.Assert(result.Utf8String, Equals, "unicode! ☯ - ♫")
	bigInt := new(big.Int)
	bigInt.SetString("1329227995784915872903807060280344576", 10)
	c.Assert(&result.Uint128, DeepEquals, bigInt)

	reader.Close()
}

func (s *MySuite) TestIpv6inIpv4(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-ipv4-24.mmdb")

	var result TestType
	err := reader.Lookup(net.ParseIP("2001::"), &result)
	// if result != nil {
	// 	c.Log("nil record from lookup expected")
	// 	c.Fail()
	// }
	expected := errors.New("error looking up '2001::': you attempted to look up an IPv6 address in an IPv4-only database")
	c.Assert(err, DeepEquals, expected)
	reader.Close()

}

func (s *MySuite) TestBrokenDatabase(c *C) {
	c.Skip("NO TYPE DECODING VALIDATION DONE")
	reader, _ := Open("test-data/test-data/GeoIP2-City-Test-Broken-Double-Format.mmdb")
	// // Should return an error like: "The MaxMind DB file's data "
	// //                              "section contains bad data (unknown data "
	// //                              "type or corrupt data)"

	var result TestType
	reader.Lookup(net.ParseIP("2001:220::"), &result)
	reader.Close()

}

func (s *MySuite) TestMissingDatabase(c *C) {
	reader, err := Open("file-does-not-exist.mmdb")
	if reader != nil {
		c.Log("received reader when doing lookups on DB that doesn't exist")
		c.Fail()
	}
	c.Assert(err.Error(), Equals, "open file-does-not-exist.mmdb: no such file or directory")
}

func (s *MySuite) TestNonDatabase(c *C) {
	reader, err := Open("README.md")
	if reader != nil {
		c.Log("received reader when doing lookups on DB that doesn't exist")
		c.Fail()
	}
	c.Assert(err.Error(), Equals, "error opening database file: invalid MaxMind DB file")
}

func checkMetadata(c *C, reader *Reader, ipVersion uint, recordSize uint) {
	metadata := reader.Metadata

	c.Assert(metadata.BinaryFormatMajorVersion, Equals, uint(2))

	c.Assert(metadata.BinaryFormatMinorVersion, Equals, uint(0))
	c.Assert(metadata.BuildEpoch, FitsTypeOf, uint(0))
	c.Assert(metadata.DatabaseType, Equals, "Test")

	c.Assert(metadata.Description, DeepEquals,
		map[string]string{
			"en": "Test Database",
			"zh": "Test Database Chinese",
		})
	c.Assert(metadata.IPVersion, Equals, ipVersion)
	c.Assert(metadata.Languages, DeepEquals, []string{"en", "zh"})

	if ipVersion == 4 {
		c.Assert(metadata.NodeCount, Equals, uint(37))
	} else {
		c.Assert(metadata.NodeCount, Equals, uint(188))
	}

	c.Assert(metadata.RecordSize, Equals, recordSize)
}

func checkIpv4(c *C, reader *Reader) {

	for i := uint(0); i < 6; i++ {
		address := fmt.Sprintf("1.1.1.%d", uint(1)<<i)
		ip := net.ParseIP(address)

		// XXX - Figure out why ParseIP always returns 16 byte address.
		// Maybe update reader to accommodated
		var result map[string]string
		reader.Lookup(ip[12:], &result)
		c.Assert(result, DeepEquals, map[string]string{
			"ip": address})
	}
	pairs := map[string]string{
		"1.1.1.3":  "1.1.1.2",
		"1.1.1.5":  "1.1.1.4",
		"1.1.1.7":  "1.1.1.4",
		"1.1.1.9":  "1.1.1.8",
		"1.1.1.15": "1.1.1.8",
		"1.1.1.17": "1.1.1.16",
		"1.1.1.31": "1.1.1.16",
	}

	for keyAddress, valueAddress := range pairs {
		data := map[string]string{"ip": valueAddress}

		ip := net.ParseIP(keyAddress)

		var result map[string]string
		reader.Lookup(ip[12:], &result)
		c.Assert(result, DeepEquals, data)
	}

	for _, address := range []string{"1.1.1.33", "255.254.253.123"} {
		ip := net.ParseIP(address)

		var result map[string]string
		reader.Lookup(ip[12:], &result)
		c.Assert(result, IsNil)
	}
}

func checkIpv6(c *C, reader *Reader) {

	subnets := []string{"::1:ffff:ffff", "::2:0:0",
		"::2:0:40", "::2:0:50", "::2:0:58"}

	for _, address := range subnets {
		var result map[string]string
		reader.Lookup(net.ParseIP(address), &result)
		c.Assert(result, DeepEquals, map[string]string{"ip": address})
	}

	pairs := map[string]string{
		"::2:0:1":  "::2:0:0",
		"::2:0:33": "::2:0:0",
		"::2:0:39": "::2:0:0",
		"::2:0:41": "::2:0:40",
		"::2:0:49": "::2:0:40",
		"::2:0:52": "::2:0:50",
		"::2:0:57": "::2:0:50",
		"::2:0:59": "::2:0:58",
	}

	for keyAddress, valueAddress := range pairs {
		data := map[string]string{"ip": valueAddress}
		var result map[string]string
		reader.Lookup(net.ParseIP(keyAddress), &result)
		c.Assert(result, DeepEquals, data)
	}

	for _, address := range []string{"1.1.1.33", "255.254.253.123", "89fa::"} {
		var result map[string]string
		reader.Lookup(net.ParseIP(address), &result)
		c.Assert(result, IsNil)
	}
}

func BenchmarkMaxMindDB(b *testing.B) {
	db, err := Open("GeoLite2-City.mmdb")
	if err != nil {
		b.Fatal(err)
	}

	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	var result interface{}

	for i := 0; i < b.N; i++ {
		num := r.Uint32()
		ip := net.ParseIP(fmt.Sprintf("%d.%d.%d.%d", (0xFF000000&num)>>24,
			(0x00FF0000&num)>>16, (0x0000FF00&num)>>8, 0x000000F&num))
		err := db.Lookup(ip, &result)
		if err != nil {
			b.Fatal(err)
		}
	}
	db.Close()
}
<commit_msg>Changed node count for latest test databases<commit_after>package maxminddb

import (
	"errors"
	"fmt"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"math/big"
	"math/rand"
	"net"
	"testing"
	"time"
)

func TestMaxMindDbReader(t *testing.T) { TestingT(t) }

type MySuite struct{}

var _ = Suite(&MySuite{})

func (s *MySuite) TestReader(c *C) {
	for _, recordSize := range []uint{24, 28, 32} {
		for _, ipVersion := range []uint{4, 6} {
			fileName := fmt.Sprintf("test-data/test-data/MaxMind-DB-test-ipv%d-%d.mmdb", ipVersion, recordSize)
			reader, _ := Open(fileName)

			checkMetadata(c, reader, ipVersion, recordSize)

			if ipVersion == 4 {
				checkIpv4(c, reader)
			} else {
				checkIpv6(c, reader)
			}
		}
	}
}

func (s *MySuite) TestReaderBytes(c *C) {
	for _, recordSize := range []uint{24, 28, 32} {
		for _, ipVersion := range []uint{4, 6} {
			fileName := fmt.Sprintf("test-data/test-data/MaxMind-DB-test-ipv%d-%d.mmdb", ipVersion, recordSize)
			bytes, _ := ioutil.ReadFile(fileName)
			reader, _ := FromBytes(bytes)

			checkMetadata(c, reader, ipVersion, recordSize)

			if ipVersion == 4 {
				checkIpv4(c, reader)
			} else {
				checkIpv6(c, reader)
			}
		}
	}
}

func (s *MySuite) TestDecodingToInterface(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-decoder.mmdb")

	var recordInterface interface{}
	reader.Lookup(net.ParseIP("::1.1.1.0"), &recordInterface)
	record := recordInterface.(map[string]interface{})

	c.Assert(record["array"], DeepEquals, []interface{}{uint64(1), uint64(2), uint64(3)})
	c.Assert(record["boolean"], Equals, true)
	c.Assert(record["bytes"], DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
	c.Assert(record["double"], Equals, 42.123456)
	c.Assert(record["float"], Equals, float32(1.1))
	c.Assert(record["int32"], Equals, -268435456)
	c.Assert(record["map"], DeepEquals,
		map[string]interface{}{
			"mapX": map[string]interface{}{
				"arrayX":       []interface{}{uint64(7), uint64(8), uint64(9)},
				"utf8_stringX": "hello",
			}})

	c.Assert(record["uint16"], Equals, uint64(100))
	c.Assert(record["uint32"], Equals, uint64(268435456))
	c.Assert(record["uint64"], Equals, uint64(1152921504606846976))
	c.Assert(record["utf8_string"], Equals, "unicode! ☯ - ♫")
	bigInt := new(big.Int)
	bigInt.SetString("1329227995784915872903807060280344576", 10)
	c.Assert(record["uint128"], DeepEquals, bigInt)
}

type TestType struct {
	Array      []uint                 `maxminddb:"array"`
	Boolean    bool                   `maxminddb:"boolean"`
	Bytes      []byte                 `maxminddb:"bytes"`
	Double     float64                `maxminddb:"double"`
	Float      float32                `maxminddb:"float"`
	Int32      int32                  `maxminddb:"int32"`
	Map        map[string]interface{} `maxminddb:"map"`
	Uint16     uint16                 `maxminddb:"uint16"`
	Uint32     uint32                 `maxminddb:"uint32"`
	Uint64     uint64                 `maxminddb:"uint64"`
	Uint128    big.Int                `maxminddb:"uint128"`
	Utf8String string                 `maxminddb:"utf8_string"`
}

func (s *MySuite) TestDecoder(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-decoder.mmdb")

	var result TestType
	err := reader.Lookup(net.ParseIP("::1.1.1.0"), &result)
	if err != nil {
		c.Log(err)
		c.Fail()
	}

	c.Assert(result.Array, DeepEquals, []uint{uint(1), uint(2), uint(3)})
	c.Assert(result.Boolean, Equals, true)
	c.Assert(result.Bytes, DeepEquals, []byte{0x00, 0x00, 0x00, 0x2a})
	c.Assert(result.Double, Equals, 42.123456)
	c.Assert(result.Float, Equals, float32(1.1))
	c.Assert(result.Int32, Equals, int32(-268435456))

	c.Assert(result.Map, DeepEquals,
		map[string]interface{}{
			"mapX": map[string]interface{}{
				"arrayX":       []interface{}{uint64(7), uint64(8), uint64(9)},
				"utf8_stringX": "hello",
			}})

	c.Assert(result.Uint16, Equals, uint16(100))
	c.Assert(result.Uint32, Equals, uint32(268435456))
	c.Assert(result.Uint64, Equals, uint64(1152921504606846976))
	c.Assert(result.Utf8String, Equals, "unicode! ☯ - ♫")
	bigInt := new(big.Int)
	bigInt.SetString("1329227995784915872903807060280344576", 10)
	c.Assert(&result.Uint128, DeepEquals, bigInt)

	reader.Close()
}

func (s *MySuite) TestIpv6inIpv4(c *C) {
	reader, _ := Open("test-data/test-data/MaxMind-DB-test-ipv4-24.mmdb")

	var result TestType
	err := reader.Lookup(net.ParseIP("2001::"), &result)
	// if result != nil {
	// 	c.Log("nil record from lookup expected")
	// 	c.Fail()
	// }
	expected := errors.New("error looking up '2001::': you attempted to look up an IPv6 address in an IPv4-only database")
	c.Assert(err, DeepEquals, expected)
	reader.Close()

}

func (s *MySuite) TestBrokenDatabase(c *C) {
	c.Skip("NO TYPE DECODING VALIDATION DONE")
	reader, _ := Open("test-data/test-data/GeoIP2-City-Test-Broken-Double-Format.mmdb")
	// // Should return an error like: "The MaxMind DB file's data "
	// //                              "section contains bad data (unknown data "
	// //                              "type or corrupt data)"

	var result TestType
	reader.Lookup(net.ParseIP("2001:220::"), &result)
	reader.Close()

}

func (s *MySuite) TestMissingDatabase(c *C) {
	reader, err := Open("file-does-not-exist.mmdb")
	if reader != nil {
		c.Log("received reader when doing lookups on DB that doesn't exist")
		c.Fail()
	}
	c.Assert(err.Error(), Equals, "open file-does-not-exist.mmdb: no such file or directory")
}

func (s *MySuite) TestNonDatabase(c *C) {
	reader, err := Open("README.md")
	if reader != nil {
		c.Log("received reader when doing lookups on DB that doesn't exist")
		c.Fail()
	}
	c.Assert(err.Error(), Equals, "error opening database file: invalid MaxMind DB file")
}

func checkMetadata(c *C, reader *Reader, ipVersion uint, recordSize uint) {
	metadata := reader.Metadata

	c.Assert(metadata.BinaryFormatMajorVersion, Equals, uint(2))

	c.Assert(metadata.BinaryFormatMinorVersion, Equals, uint(0))
	c.Assert(metadata.BuildEpoch, FitsTypeOf, uint(0))
	c.Assert(metadata.DatabaseType, Equals, "Test")

	c.Assert(metadata.Description, DeepEquals,
		map[string]string{
			"en": "Test Database",
			"zh": "Test Database Chinese",
		})
	c.Assert(metadata.IPVersion, Equals, ipVersion)
	c.Assert(metadata.Languages, DeepEquals, []string{"en", "zh"})

	if ipVersion == 4 {
		c.Assert(metadata.NodeCount, Equals, uint(37))
	} else {
		c.Assert(metadata.NodeCount, Equals, uint(160))
	}

	c.Assert(metadata.RecordSize, Equals, recordSize)
}

func checkIpv4(c *C, reader *Reader) {

	for i := uint(0); i < 6; i++ {
		address := fmt.Sprintf("1.1.1.%d", uint(1)<<i)
		ip := net.ParseIP(address)

		// XXX - Figure out why ParseIP always returns 16 byte address.
		// Maybe update reader to accommodated
		var result map[string]string
		reader.Lookup(ip[12:], &result)
		c.Assert(result, DeepEquals, map[string]string{
			"ip": address})
	}
	pairs := map[string]string{
		"1.1.1.3":  "1.1.1.2",
		"1.1.1.5":  "1.1.1.4",
		"1.1.1.7":  "1.1.1.4",
		"1.1.1.9":  "1.1.1.8",
		"1.1.1.15": "1.1.1.8",
		"1.1.1.17": "1.1.1.16",
		"1.1.1.31": "1.1.1.16",
	}

	for keyAddress, valueAddress := range pairs {
		data := map[string]string{"ip": valueAddress}

		ip := net.ParseIP(keyAddress)

		var result map[string]string
		reader.Lookup(ip[12:], &result)
		c.Assert(result, DeepEquals, data)
	}

	for _, address := range []string{"1.1.1.33", "255.254.253.123"} {
		ip := net.ParseIP(address)

		var result map[string]string
		reader.Lookup(ip[12:], &result)
		c.Assert(result, IsNil)
	}
}

func checkIpv6(c *C, reader *Reader) {

	subnets := []string{"::1:ffff:ffff", "::2:0:0",
		"::2:0:40", "::2:0:50", "::2:0:58"}

	for _, address := range subnets {
		var result map[string]string
		reader.Lookup(net.ParseIP(address), &result)
		c.Assert(result, DeepEquals, map[string]string{"ip": address})
	}

	pairs := map[string]string{
		"::2:0:1":  "::2:0:0",
		"::2:0:33": "::2:0:0",
		"::2:0:39": "::2:0:0",
		"::2:0:41": "::2:0:40",
		"::2:0:49": "::2:0:40",
		"::2:0:52": "::2:0:50",
		"::2:0:57": "::2:0:50",
		"::2:0:59": "::2:0:58",
	}

	for keyAddress, valueAddress := range pairs {
		data := map[string]string{"ip": valueAddress}
		var result map[string]string
		reader.Lookup(net.ParseIP(keyAddress), &result)
		c.Assert(result, DeepEquals, data)
	}

	for _, address := range []string{"1.1.1.33", "255.254.253.123", "89fa::"} {
		var result map[string]string
		reader.Lookup(net.ParseIP(address), &result)
		c.Assert(result, IsNil)
	}
}

func BenchmarkMaxMindDB(b *testing.B) {
	db, err := Open("GeoLite2-City.mmdb")
	if err != nil {
		b.Fatal(err)
	}

	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	var result interface{}

	for i := 0; i < b.N; i++ {
		num := r.Uint32()
		ip := net.ParseIP(fmt.Sprintf("%d.%d.%d.%d", (0xFF000000&num)>>24,
			(0x00FF0000&num)>>16, (0x0000FF00&num)>>8, 0x000000F&num))
		err := db.Lookup(ip, &result)
		if err != nil {
			b.Fatal(err)
		}
	}
	db.Close()
}
<|endoftext|>


<commit_before>package reaper

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"

	reaperaws "github.com/mozilla-services/reaper/aws"
	reaperevents "github.com/mozilla-services/reaper/events"
	"github.com/mozilla-services/reaper/reapable"
	log "github.com/mozilla-services/reaper/reaperlog"
	"github.com/mozilla-services/reaper/state"
	"github.com/mozilla-services/reaper/token"
)

type HTTPApi struct {
	conf   reaperevents.HTTPConfig
	server *http.Server
	ln     net.Listener
}

// Serve should be run in a goroutine
func (h *HTTPApi) Serve() (e error) {
	h.ln, e = net.Listen("tcp", h.conf.Listen)

	if e != nil {
		return
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/", processToken(h))
	mux.HandleFunc("/__heartbeat__", heartbeat(h))
	h.server = &http.Server{Handler: mux}

	log.Debug("Starting HTTP server: %s", h.conf.Listen)
	go h.server.Serve(h.ln)
	return nil
}

// Stop will close the listener, it waits for nothing
func (h *HTTPApi) Stop() (e error) {
	return h.ln.Close()
}

func NewHTTPApi(c reaperevents.HTTPConfig) *HTTPApi {
	return &HTTPApi{conf: c}
}

func writeResponse(w http.ResponseWriter, code int, body string) {
	w.Header().Set("Content-Type", "text/html")
	io.WriteString(w, fmt.Sprintf(`<DOCTYPE html>

				<title>Reaper API</title>

				%s

		`,
		body))
}

func heartbeat(h *HTTPApi) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		writeResponse(w, http.StatusOK, "Heart's a beatin'")
		return
	}
}

func processToken(h *HTTPApi) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		if err := req.ParseForm(); err != nil {
			writeResponse(w, http.StatusBadRequest, "Bad query string")
			return
		}

		userToken := req.Form.Get(h.conf.Token)
		if userToken == "" {
			writeResponse(w, http.StatusBadRequest, "Token Missing")
			return
		}

		if u, err := url.QueryUnescape(userToken); err == nil {
			userToken = u
		} else {
			writeResponse(w,
				http.StatusBadRequest, "Invalid Token, could not decode data")
			return
		}

		job, err := token.Untokenize(h.conf.TokenSecret, userToken)
		if err != nil {
			writeResponse(w,
				http.StatusBadRequest, "Invalid Token, Could not untokenize")
			return
		}

		if job.Expired() == true {
			writeResponse(w, http.StatusBadRequest, "Token expired")
			return
		}

		// find reapable associated with the job
		r, err := reapables.Get(reapable.Region(job.Region), reapable.ID(job.ID))
		if err != nil {
			writeResponse(w, http.StatusInternalServerError, err.Error())
			return
		}

		switch job.Action {
		case token.J_DELAY:
			log.Debug("Delay request received for %s in region %s until %s", job.ID, job.Region, job.IgnoreUntil.String())
			s := r.ReaperState()
			ok, err := r.Save(state.NewStateWithUntilAndState(s.Until.Add(job.IgnoreUntil), s.State))
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("Delay failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_TERMINATE:
			log.Debug("Terminate request received for %s in region %s.", job.ID, job.Region)
			ok, err := r.Terminate()
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("Terminate failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_WHITELIST:
			log.Debug("Whitelist request received for %s in region %s", job.ID, job.Region)
			ok, err := r.Whitelist()
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("Whitelist failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_STOP:
			log.Debug("Stop request received for %s in region %s", job.ID, job.Region)
			ok, err := r.Stop()
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("Stop failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_FORCESTOP:
			log.Debug("Force Stop request received for %s in region %s", job.ID, job.Region)
			ok, err := r.ForceStop()
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("ForceStop failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_SCHEDULE:
			log.Debug("Schedule request received for %s in region %s", job.ID, job.Region)
			if scaler, ok := r.(reaperaws.Scaler); ok {
				scaler.SetScaleDownString(job.ScaleDownString)
				scaler.SetScaleUpString(job.ScaleUpString)
				scaler.SaveSchedule()
			}
		default:
			log.Error("Unrecognized job token received.")
			writeResponse(w, http.StatusInternalServerError, "Unrecognized job token.")
			return
		}

		var consoleURL *url.URL
		switch t := r.(type) {
		case *reaperaws.Instance:
			consoleURL = t.AWSConsoleURL()
		case *reaperaws.AutoScalingGroup:
			consoleURL = t.AWSConsoleURL()
		default:
			log.Error("No AWSConsoleURL")
		}
		writeResponse(w, http.StatusOK, fmt.Sprintf("Success. Check %s out on the AWS Console.", r.ReapableDescriptionTiny(), consoleURL))
	}
}
<commit_msg>added /__lbheartbeat__ route<commit_after>package reaper

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"

	reaperaws "github.com/mozilla-services/reaper/aws"
	reaperevents "github.com/mozilla-services/reaper/events"
	"github.com/mozilla-services/reaper/reapable"
	log "github.com/mozilla-services/reaper/reaperlog"
	"github.com/mozilla-services/reaper/state"
	"github.com/mozilla-services/reaper/token"
)

type HTTPApi struct {
	conf   reaperevents.HTTPConfig
	server *http.Server
	ln     net.Listener
}

// Serve should be run in a goroutine
func (h *HTTPApi) Serve() (e error) {
	h.ln, e = net.Listen("tcp", h.conf.Listen)

	if e != nil {
		return
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/", processToken(h))
	mux.HandleFunc("/__heartbeat__", heartbeat(h))
	mux.HandleFunc("/__lbheartbeat__", heartbeat(h))
	h.server = &http.Server{Handler: mux}

	log.Debug("Starting HTTP server: %s", h.conf.Listen)
	go h.server.Serve(h.ln)
	return nil
}

// Stop will close the listener, it waits for nothing
func (h *HTTPApi) Stop() (e error) {
	return h.ln.Close()
}

func NewHTTPApi(c reaperevents.HTTPConfig) *HTTPApi {
	return &HTTPApi{conf: c}
}

func writeResponse(w http.ResponseWriter, code int, body string) {
	w.Header().Set("Content-Type", "text/html")
	io.WriteString(w, fmt.Sprintf(`<DOCTYPE html>

				<title>Reaper API</title>

				%s

		`,
		body))
}

func heartbeat(h *HTTPApi) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		writeResponse(w, http.StatusOK, "Heart's a beatin'")
		return
	}
}

func processToken(h *HTTPApi) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		if err := req.ParseForm(); err != nil {
			writeResponse(w, http.StatusBadRequest, "Bad query string")
			return
		}

		userToken := req.Form.Get(h.conf.Token)
		if userToken == "" {
			writeResponse(w, http.StatusBadRequest, "Token Missing")
			return
		}

		if u, err := url.QueryUnescape(userToken); err == nil {
			userToken = u
		} else {
			writeResponse(w,
				http.StatusBadRequest, "Invalid Token, could not decode data")
			return
		}

		job, err := token.Untokenize(h.conf.TokenSecret, userToken)
		if err != nil {
			writeResponse(w,
				http.StatusBadRequest, "Invalid Token, Could not untokenize")
			return
		}

		if job.Expired() == true {
			writeResponse(w, http.StatusBadRequest, "Token expired")
			return
		}

		// find reapable associated with the job
		r, err := reapables.Get(reapable.Region(job.Region), reapable.ID(job.ID))
		if err != nil {
			writeResponse(w, http.StatusInternalServerError, err.Error())
			return
		}

		switch job.Action {
		case token.J_DELAY:
			log.Debug("Delay request received for %s in region %s until %s", job.ID, job.Region, job.IgnoreUntil.String())
			s := r.ReaperState()
			ok, err := r.Save(state.NewStateWithUntilAndState(s.Until.Add(job.IgnoreUntil), s.State))
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("Delay failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_TERMINATE:
			log.Debug("Terminate request received for %s in region %s.", job.ID, job.Region)
			ok, err := r.Terminate()
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("Terminate failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_WHITELIST:
			log.Debug("Whitelist request received for %s in region %s", job.ID, job.Region)
			ok, err := r.Whitelist()
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("Whitelist failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_STOP:
			log.Debug("Stop request received for %s in region %s", job.ID, job.Region)
			ok, err := r.Stop()
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("Stop failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_FORCESTOP:
			log.Debug("Force Stop request received for %s in region %s", job.ID, job.Region)
			ok, err := r.ForceStop()
			if err != nil {
				writeResponse(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !ok {
				writeResponse(w, http.StatusInternalServerError, fmt.Sprintf("ForceStop failed for %s.", r.ReapableDescriptionTiny()))
				return
			}
		case token.J_SCHEDULE:
			log.Debug("Schedule request received for %s in region %s", job.ID, job.Region)
			if scaler, ok := r.(reaperaws.Scaler); ok {
				scaler.SetScaleDownString(job.ScaleDownString)
				scaler.SetScaleUpString(job.ScaleUpString)
				scaler.SaveSchedule()
			}
		default:
			log.Error("Unrecognized job token received.")
			writeResponse(w, http.StatusInternalServerError, "Unrecognized job token.")
			return
		}

		var consoleURL *url.URL
		switch t := r.(type) {
		case *reaperaws.Instance:
			consoleURL = t.AWSConsoleURL()
		case *reaperaws.AutoScalingGroup:
			consoleURL = t.AWSConsoleURL()
		default:
			log.Error("No AWSConsoleURL")
		}
		writeResponse(w, http.StatusOK, fmt.Sprintf("Success. Check %s out on the AWS Console.", r.ReapableDescriptionTiny(), consoleURL))
	}
}
<|endoftext|>


<commit_before>package stardust

import (
	"os"
	"reflect"
	"testing"
	"time"
)

func TestLeft(t *testing.T) {
	var tests = []struct {
		record Record
		out    string
	}{
		{Record{left: 0, right: 1, Fields: []string{"0", "1"}}, "0"},
		{Record{left: 1, right: 0, Fields: []string{"0", "1"}}, "1"},
	}
	for _, tt := range tests {
		out := tt.record.Left()
		if out != tt.out {
			t.Errorf("%+v.Left() => %s, want: %s", tt.record, out, tt.out)
		}
	}
}

func TestRight(t *testing.T) {
	var tests = []struct {
		record Record
		out    string
	}{
		{Record{left: 0, right: 1, Fields: []string{"0", "1"}}, "1"},
		{Record{left: 1, right: 0, Fields: []string{"0", "1"}}, "0"},
	}
	for _, tt := range tests {
		out := tt.record.Right()
		if out != tt.out {
			t.Errorf("%+v.Right() => %s, want: %s", tt.record, out, tt.out)
		}
	}
}

func TestParseColumnSpec(t *testing.T) {
	var tests = []struct {
		s   string
		out *ColumnSpec
	}{
		{"1,2", &ColumnSpec{left: 0, right: 1}},
	}

	for _, tt := range tests {
		out, _ := ParseColumnSpec(tt.s)
		if !reflect.DeepEqual(out, tt.out) {
			t.Errorf("ParseColumnSpec(%s) => %+v, want: %+v", tt.s, out, tt.out)
		}
	}
}

func TestRecordGeneratorFile(t *testing.T) {
	fixture := "fixtures/strings.tsv"
	if _, err := os.Stat(fixture); os.IsNotExist(err) {
		t.Skipf("skipping since fixture is missing: %s\n", fixture)
		return
	}
	file, err := os.Open(fixture)
	if err != nil {
		t.Errorf("could not open fixture: %v\n", err)
	}
	columnSpec := &ColumnSpec{left: 0, right: 1}
	records := RecordGeneratorFile(file, columnSpec)
	select {
	case r := <-records:
		if len(r.Fields) != 2 {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected length of record.fields: %v", file, columnSpec, r)
		}
		if r.left != 0 {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected value in record.left: %v", file, columnSpec, r)
		}
		if r.right != 1 {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected value in record.field : %v", file, columnSpec, r)
		}
		if r.Left() != "Somewhere over the rainbow" {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected record.Left(): %v, got: %v", file, columnSpec, r, r.Left())
		}
		if r.Right() != "Somewhere near Japan" {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected record.Right(): %v, got: %v", file, columnSpec, r, r.Right())
		}
	case <-time.After(1 * time.Microsecond):
		t.Errorf("RecordGeneratorFile(%v, %v) => timeout! ", file, columnSpec)
	}
}
<commit_msg>increase timeout in test<commit_after>package stardust

import (
	"os"
	"reflect"
	"testing"
	"time"
)

func TestLeft(t *testing.T) {
	var tests = []struct {
		record Record
		out    string
	}{
		{Record{left: 0, right: 1, Fields: []string{"0", "1"}}, "0"},
		{Record{left: 1, right: 0, Fields: []string{"0", "1"}}, "1"},
	}
	for _, tt := range tests {
		out := tt.record.Left()
		if out != tt.out {
			t.Errorf("%+v.Left() => %s, want: %s", tt.record, out, tt.out)
		}
	}
}

func TestRight(t *testing.T) {
	var tests = []struct {
		record Record
		out    string
	}{
		{Record{left: 0, right: 1, Fields: []string{"0", "1"}}, "1"},
		{Record{left: 1, right: 0, Fields: []string{"0", "1"}}, "0"},
	}
	for _, tt := range tests {
		out := tt.record.Right()
		if out != tt.out {
			t.Errorf("%+v.Right() => %s, want: %s", tt.record, out, tt.out)
		}
	}
}

func TestParseColumnSpec(t *testing.T) {
	var tests = []struct {
		s   string
		out *ColumnSpec
	}{
		{"1,2", &ColumnSpec{left: 0, right: 1}},
	}

	for _, tt := range tests {
		out, _ := ParseColumnSpec(tt.s)
		if !reflect.DeepEqual(out, tt.out) {
			t.Errorf("ParseColumnSpec(%s) => %+v, want: %+v", tt.s, out, tt.out)
		}
	}
}

func TestRecordGeneratorFile(t *testing.T) {
	fixture := "fixtures/strings.tsv"
	if _, err := os.Stat(fixture); os.IsNotExist(err) {
		t.Skipf("skipping since fixture is missing: %s\n", fixture)
		return
	}
	file, err := os.Open(fixture)
	if err != nil {
		t.Errorf("could not open fixture: %v\n", err)
	}
	columnSpec := &ColumnSpec{left: 0, right: 1}
	records := RecordGeneratorFile(file, columnSpec)
	select {
	case r := <-records:
		if len(r.Fields) != 2 {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected length of record.fields: %v", file, columnSpec, r)
		}
		if r.left != 0 {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected value in record.left: %v", file, columnSpec, r)
		}
		if r.right != 1 {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected value in record.field : %v", file, columnSpec, r)
		}
		if r.Left() != "Somewhere over the rainbow" {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected record.Left(): %v, got: %v", file, columnSpec, r, r.Left())
		}
		if r.Right() != "Somewhere near Japan" {
			t.Errorf("RecordGeneratorFile(%v, %v) => unexpected record.Right(): %v, got: %v", file, columnSpec, r, r.Right())
		}
	case <-time.After(1 * time.Millisecond):
		t.Errorf("RecordGeneratorFile(%v, %v) => timeout! ", file, columnSpec)
	}
}
<|endoftext|>


<commit_before>// Cache object to redis automatically by registering
// callback to gorm

package goal

import (
	"encoding/json"
	"fmt"

	"github.com/garyburd/redigo/redis"
	"github.com/jinzhu/gorm"
)

var pool *redis.Pool

// InitRedisPool initialize connection pool to redis
func InitRedisPool(p *redis.Pool) {
	pool = p

	// Register Gorm callbacks
	if db != nil {
		db.Callback().Create().After("gorm:after_create").Register("cacheToRedis", cacheToRedis)
		db.Callback().Update().After("gorm:after_update").Register("cacheToRedis", cacheToRedis)
		db.Callback().Query().After("gorm:after_query").Register("cacheToRedis", cacheToRedis)
		db.Callback().Delete().Before("gorm:before_delete").Register("uncacheBeforeDelete", uncacheFromRedis)
	}
}

// Pool return global connection pool to redis
func Pool() *redis.Pool {
	return pool
}

func uncacheFromRedis(scope *gorm.Scope) {
	// Reload object before delete
	scope.DB().New().First(scope.Value)

	// Delete from redis
	key := redisKeyFromScope(scope)
	RedisUnset(key)
}

func cacheToRedis(scope *gorm.Scope) {
	key := redisKeyFromScope(scope)
	RedisSet(key, scope.Value)
}

func redisKeyFromScope(scope *gorm.Scope) string {
	name := scope.TableName()
	id := scope.PrimaryKeyValue()
	key := DefaultRedisKey(name, id)
	return key
}

// RedisKey defines by the struct or fallback
// to name:id format
func RedisKey(resource interface{}) string {
	scope := db.NewScope(resource)
	return redisKeyFromScope(scope)
}

// DefaultRedisKey returns default format for redis key
func DefaultRedisKey(name string, id interface{}) string {
	return fmt.Sprintf("%s:%d", name, id)
}

// RedisSet saves object to redis
func RedisSet(key string, resource interface{}) error {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return err
	}

	defer conn.Close()

	var val []byte
	val, err = json.Marshal(resource)
	if err != nil {
		return err
	}

	_, err = conn.Do("SET", key, val)
	return err
}

// RedisUnset delete object from redis
func RedisUnset(key string) error {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return err
	}

	defer conn.Close()

	fmt.Printf("UNSET %s", key)
	_, err = conn.Do("DEL", key)
	return err
}

func RedisExists(key string) (bool, error) {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return false, err
	}

	defer conn.Close()

	var reply bool
	reply, err = redis.Bool(conn.Do("EXISTS", key))

	return reply, err
}

// RedisGet returns object from database
func RedisGet(key string, resource interface{}) error {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return err
	}

	defer conn.Close()

	if key == "" {
		return nil
	}

	var reply []byte
	reply, err = redis.Bytes(conn.Do("GET", key))
	if err != nil {
		return err
	}

	// Populate resource
	json.Unmarshal(reply, resource)

	return nil
}

// RedisClearAll clear all data from connection's CURRENT database
func RedisClearAll() error {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return err
	}

	defer conn.Close()

	_, err = conn.Do("FLUSHDB")

	if err != nil {
		fmt.Println("Error clear redis ", err)
	}

	return err
}
<commit_msg>minor fix<commit_after>// Cache object to redis automatically by registering
// callback to gorm

package goal

import (
	"encoding/json"
	"fmt"

	"github.com/garyburd/redigo/redis"
	"github.com/jinzhu/gorm"
)

var pool *redis.Pool

// InitRedisPool initialize connection pool to redis
func InitRedisPool(p *redis.Pool) {
	pool = p

	// Register Gorm callbacks
	if db != nil {
		db.Callback().Create().After("gorm:after_create").Register("goal:cache_after_create", cacheToRedis)
		db.Callback().Update().After("gorm:after_update").Register("goal:cache_after_update", cacheToRedis)
		db.Callback().Query().After("gorm:after_query").Register("goal:cache_after_query", cacheToRedis)
		db.Callback().Delete().Before("gorm:before_delete").Register("goal:uncache_after_delete", uncacheFromRedis)
	}
}

// Pool return global connection pool to redis
func Pool() *redis.Pool {
	return pool
}

func uncacheFromRedis(scope *gorm.Scope) {
	// Reload object before delete
	scope.DB().New().First(scope.Value)

	// Delete from redis
	key := redisKeyFromScope(scope)
	RedisUnset(key)
}

func cacheToRedis(scope *gorm.Scope) {
	key := redisKeyFromScope(scope)
	RedisSet(key, scope.Value)
}

func redisKeyFromScope(scope *gorm.Scope) string {
	name := scope.TableName()
	id := scope.PrimaryKeyValue()
	key := DefaultRedisKey(name, id)
	return key
}

// RedisKey defines by the struct or fallback
// to name:id format
func RedisKey(resource interface{}) string {
	scope := db.NewScope(resource)
	return redisKeyFromScope(scope)
}

// DefaultRedisKey returns default format for redis key
func DefaultRedisKey(name string, id interface{}) string {
	return fmt.Sprintf("%s:%d", name, id)
}

// RedisSet saves object to redis
func RedisSet(key string, resource interface{}) error {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return err
	}

	defer conn.Close()

	var val []byte
	val, err = json.Marshal(resource)
	if err != nil {
		return err
	}

	_, err = conn.Do("SET", key, val)
	return err
}

// RedisUnset delete object from redis
func RedisUnset(key string) error {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return err
	}

	defer conn.Close()

	fmt.Printf("UNSET %s", key)
	_, err = conn.Do("DEL", key)
	return err
}

func RedisExists(key string) (bool, error) {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return false, err
	}

	defer conn.Close()

	var reply bool
	reply, err = redis.Bool(conn.Do("EXISTS", key))

	return reply, err
}

// RedisGet returns object from database
func RedisGet(key string, resource interface{}) error {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return err
	}

	defer conn.Close()

	if key == "" {
		return nil
	}

	var reply []byte
	reply, err = redis.Bytes(conn.Do("GET", key))
	if err != nil {
		return err
	}

	// Populate resource
	json.Unmarshal(reply, resource)

	return nil
}

// RedisClearAll clear all data from connection's CURRENT database
func RedisClearAll() error {
	conn, err := pool.Dial()
	if err != nil {
		fmt.Println(err)
		return err
	}

	defer conn.Close()

	_, err = conn.Do("FLUSHDB")

	if err != nil {
		fmt.Println("Error clear redis ", err)
	}

	return err
}
<|endoftext|>


<commit_before>package parse

import (
	"bytes"
	"github.com/kortschak/utter"
	"reflect"
	"strings"
	"testing"
	"fmt"
	"github.com/rickb777/sqlgen2/sqlgen/parse/exit"
)

func TestStructWith3FieldsAndTags(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg1", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Id",
					Type: Type{"", "int64", Int64},
					Tags: &Tag{Primary: true, Auto: true},
				}, {
					Name: "Number",
					Type: Type{"", "int", Int},
					Tags: &Tag{},
				}, {
					Name: "Title",
					Type: Type{"", "string", String},
					Tags: &Tag{},
				}, {
					Name: "Description",
					Type: Type{"", "string", String},
					Tags: &Tag{},
				}, {
					Name: "Owner",
					Type: Type{"", "string", String},
					Tags: &Tag{},
				},
			},
		},
		"pkg1", "Struct",
		`package pkg1

		type Struct struct {
			Id       int64 |sql:"pk: true, auto: true"|
			Number   int
			Title, Description, Owner    string // must find all three fields
		}`,
	)
}

func TestStructWith1BoolFieldAndIgnoreTag(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg2", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Flag",
					Type: Type{"", "bool", Bool},
					Tags: &Tag{Skip: true},
				},
			},
		},
		"pkg2", "Struct",
		`package pkg2

		type Struct struct {
			Flag  bool     |sql:"-"|
		}`,
	)
}

func TestStructWith1SliceFieldAndJsonTag(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg3", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Labels",
					Type: Type{"", "[]string", Slice},
					Tags: &Tag{Encode: "json"},
				},
			},
		},
		"pkg3", "Struct",
		`package pkg3

		type Struct struct {
			Labels   []string  |sql:"encode: json"|
		}`,
	)
}

func TestStructWith1MapFieldAndJsonTag(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg4", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Table",
					Type: Type{"", "map[string]int", Map},
					Tags: &Tag{Encode: "json"},
				},
			},
		},
		"pkg4", "Struct",
		`package pkg4

		type Struct struct {
			Table    map[string]int  |sql:"encode: json"|
		}`,
	)
}

func TestStructWithNestedStructType(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg5", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Author",
					Type: Type{"", "Author", Ptr},
					Tags: &Tag{},
					Nodes: []*Node{
						{
							Name: "Name",
							Type: Type{"", "string", String},
							Tags: &Tag{},
						},
					},
				},
			},
		},
		"pkg5", "Struct",
		`package pkg5

		type Struct struct {
			Author    *Author
		}
		type Author struct {
			Name     string
		}`,
	)
}

func TestStructWithNestedSimpleType(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg6", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Cat",
					Type: Type{"pkg6", "Category", Int32},
					Tags: &Tag{},
				},
			},
		},
		"pkg6", "Struct",
		`package pkg6

		type Category int32

		type Struct struct {
			Cat      Category
		}`,
	)
}

func TestStructWithNestedSimpleTypeInOtherPackageOrder1(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg6", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Cat",
					Type: Type{"other", "Category", Int32},
					Tags: &Tag{},
				},
			},
		},
		"pkg6", "Struct",
		`package pkg6

		type Struct struct {
			Cat      other.Category
		}`,
		`package other

		type Category int32
		`,
	)
}

func TestStructWithNestedSimpleTypeInOtherPackageOrder2(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg6", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Cat",
					Type: Type{"other", "Category", Int32},
					Tags: &Tag{},
				},
			},
		},
		"pkg6", "Struct",
		`package other

		type Category int32
		`,
		`package pkg6

		type Struct struct {
			Cat      other.Category
		}`,
	)
}

func TestStructWithNestingAcross2Packages(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct7",
			Type: Type{"pkg7a", "Struct7", Struct},
			Nodes: []*Node{
				{
					Name: "Id",
					Type: Type{"", "uint32", Uint32},
					Tags: &Tag{},
				},
				{
					Name: "Wibble",
					Type: Type{"stringy", "Thingy", String},
					Tags: &Tag{},
				},
				{
					Name: "Bobble",
					Type: Type{"", "string", String},
					Tags: &Tag{},
				},
			},
		},
		"pkg7a", "Struct7",
		`package stringy

		type Thingy string
		`,
		`package pkg7b

		type Inner1 struct {
			Wibble stringy.Thingy
		}
		`,

		`package pkg7a

		type Struct7 struct {
			Id uint32
			pkg7b.Inner1
			Bobble string
		}
		`,
	)
}

func TestStructWithNestingAcross3Packages(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct7",
			Type: Type{"pkg7a", "Struct7", Struct},
			Nodes: []*Node{
				{
					Name: "Id",
					Type: Type{"", "uint32", Uint32},
					Tags: &Tag{},
				},
				{
					Name: "Uid",
					Type: Type{"", "uint32", Uint32},
					Tags: &Tag{},
				},
				{
					Name: "Name",
					Type: Type{"userindex", "Username", String},
					Tags: &Tag{},
				},
				{
					Name: "Wibble",
					Type: Type{"stringy", "Thingy", String},
					Tags: &Tag{},
				},
				{
					Name: "Bobble",
					Type: Type{"userindex", "Username", String},
					Tags: &Tag{},
				},
			},
		},
		"pkg7a", "Struct7",
		`package stringy

		type Thingy string
		`,

		`package userindex

		type Username string
		`,

		`package userindex

		type User struct {
			Uid  uint32
			Name Username
		}

		type UserWithThingy struct {
			User
			Wibble stringy.Thingy
		}
		`,

		`package pkg7a

		type Struct7 struct {
			Id uint32
			userindex.UserWithThingy
			Bobble userindex.Username
		}
		`,
	)
}

func doTestParseOK(t *testing.T, want *Node, pkg, name string, isource ...string) {
	t.Helper()
	exit.TestableExit()
	Debug = true

	// fix edges missing in the literal values
	for _, n0 := range want.Nodes {
		n0.Parent = want
		for _, n1 := range n0.Nodes {
			n1.Parent = n0
		}
	}

	files := make([]Source, len(isource))

	for i, s := range isource {
		// allow nested back-ticks
		source := strings.Replace(s, "|", "`", -1)
		files[i] = Source{fmt.Sprintf("issue%d.go", i), bytes.NewBufferString(source)}
	}

	err := parseAllFiles(files)
	if err != nil {
		t.Errorf("Error parsing: %s", err)
	}

	got, err := findMatchingNodes(pkg, name)
	if err != nil {
		t.Errorf("Error parsing: %s", err)
	}

	if !reflect.DeepEqual(got, want) {
		t.Errorf("Wanted %s\nGot %s", utter.Sdump(want), utter.Sdump(got))
	}
}
<commit_msg>revised tests<commit_after>package parse

import (
	"bytes"
	"github.com/kortschak/utter"
	"reflect"
	"strings"
	"testing"
	"fmt"
	"github.com/rickb777/sqlgen2/sqlgen/parse/exit"
	"os"
)

func TestStructWith3FieldsAndTags(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg1", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Id",
					Type: Type{"", "int64", Int64},
					Tags: &Tag{Primary: true, Auto: true},
				}, {
					Name: "Number",
					Type: Type{"", "int", Int},
					Tags: &Tag{},
				}, {
					Name: "Title",
					Type: Type{"", "string", String},
					Tags: &Tag{},
				}, {
					Name: "Description",
					Type: Type{"", "string", String},
					Tags: &Tag{},
				}, {
					Name: "Owner",
					Type: Type{"", "string", String},
					Tags: &Tag{},
				},
			},
		},
		"pkg1", "Struct",
		`package pkg1

		type Struct struct {
			Id       int64 |sql:"pk: true, auto: true"|
			Number   int
			Title, Description, Owner    string // must find all three fields
		}`,
	)
}

func TestStructWith1BoolFieldAndIgnoreTag(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg2", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Flag",
					Type: Type{"", "bool", Bool},
					Tags: &Tag{Skip: true},
				},
			},
		},
		"pkg2", "Struct",
		`package pkg2

		type Struct struct {
			Flag  bool     |sql:"-"|
		}`,
	)
}

func TestStructWith1SliceFieldAndJsonTag(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg3", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Labels",
					Type: Type{"", "[]string", Slice},
					Tags: &Tag{Encode: "json"},
				},
			},
		},
		"pkg3", "Struct",
		`package pkg3

		type Struct struct {
			Labels   []string  |sql:"encode: json"|
		}`,
	)
}

func TestStructWith1MapFieldAndJsonTag(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg4", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Table",
					Type: Type{"", "map[string]int", Map},
					Tags: &Tag{Encode: "json"},
				},
			},
		},
		"pkg4", "Struct",
		`package pkg4

		type Struct struct {
			Table    map[string]int  |sql:"encode: json"|
		}`,
	)
}

func TestStructWithNestedStructType(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg5", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Author",
					Type: Type{"", "Author", Ptr},
					Tags: &Tag{},
					Nodes: []*Node{
						{
							Name: "Name",
							Type: Type{"", "string", String},
							Tags: &Tag{},
						},
					},
				},
			},
		},
		"pkg5", "Struct",
		`package pkg5

		type Struct struct {
			Author    *Author
		}
		type Author struct {
			Name     string
		}`,
	)
}

func TestStructWithNestedSimpleType(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg6", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Cat",
					Type: Type{"pkg6", "Category", Int32},
					Tags: &Tag{},
				},
			},
		},
		"pkg6", "Struct",
		`package pkg6

		type Category int32

		type Struct struct {
			Cat      Category
		}`,
	)
}

func TestStructWithNestedSimpleTypeInOtherPackageOrder1(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg7", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Cat",
					Type: Type{"other", "Category", Int32},
					Tags: &Tag{},
				},
			},
		},
		"pkg7", "Struct",
		`package pkg7

		type Struct struct {
			Cat      other.Category
		}`,
		`package other

		type Category int32
		`,
	)
}

func TestStructWithNestedSimpleTypeInOtherPackageOrder2(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg8", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Cat",
					Type: Type{"other", "Category", Int32},
					Tags: &Tag{},
				},
			},
		},
		"pkg8", "Struct",
		`package other

		type Category int32
		`,
		`package pkg8

		type Struct struct {
			Cat      other.Category
		}`,
	)
}

func TestStructWithNestingAcross2Packages(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg9", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Id",
					Type: Type{"", "uint32", Uint32},
					Tags: &Tag{},
				},
				{
					Name: "Wibble",
					Type: Type{"stringy", "Thingy", String},
					Tags: &Tag{},
				},
				{
					Name: "Bobble",
					Type: Type{"", "string", String},
					Tags: &Tag{},
				},
			},
		},
		"pkg9", "Struct",
		`package stringy

		type Thingy string
		`,
		`package froob

		type Inner1 struct {
			Wibble stringy.Thingy
		}
		`,

		`package pkg9

		type Struct struct {
			Id uint32
			froob.Inner1
			Bobble string
		}
		`,
	)
}

func TestStructWithNestingInTheSamePackage(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg10", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Id",
					Type: Type{"", "uint32", Uint32},
					Tags: &Tag{},
				},
				{
					Name: "Uid",
					Type: Type{"", "uint32", Uint32},
					Tags: &Tag{},
				},
				{
					Name: "Name",
					Type: Type{"pkg10", "Username", String},
					Tags: &Tag{},
				},
				{
					Name: "Wibble",
					Type: Type{"pkg10", "Thingy", String},
					Tags: &Tag{},
				},
				{
					Name: "Bobble",
					Type: Type{"pkg10", "Username", String},
					Tags: &Tag{},
				},
			},
		},
		"pkg10", "Struct",
		`package pkg10

		type Thingy string

		type Username string

		type User struct {
			Uid  uint32
			Name Username
		}

		type UserWithThingy struct {
			User
			Wibble Thingy
		}

		type Struct struct {
			Id uint32
			UserWithThingy
			Bobble Username
		}
		`,
	)
}

func TestStructWithNestingAcross3Packages(t *testing.T) {
	doTestParseOK(t,
		&Node{
			Name: "Struct",
			Type: Type{"pkg11", "Struct", Struct},
			Nodes: []*Node{
				{
					Name: "Id",
					Type: Type{"", "uint32", Uint32},
					Tags: &Tag{},
				},
				{
					Name: "Uid",
					Type: Type{"", "uint32", Uint32},
					Tags: &Tag{},
				},
				{
					Name: "Name",
					Type: Type{"userindex", "Username", String},
					Tags: &Tag{},
				},
				{
					Name: "Wibble",
					Type: Type{"stringy", "Thingy", String},
					Tags: &Tag{},
				},
				{
					Name: "Bobble",
					Type: Type{"userindex", "Username", String},
					Tags: &Tag{},
				},
			},
		},
		"pkg11", "Struct",
		`package stringy

		type Thingy string
		`,

		`package userindex

		type Username string
		`,

		`package userindex

		type User struct {
			Uid  uint32
			Name Username
		}

		type UserWithThingy struct {
			User
			Wibble stringy.Thingy
		}
		`,

		`package pkg11

		type Struct struct {
			Id uint32
			userindex.UserWithThingy
			Bobble userindex.Username
		}
		`,
	)
}

func doTestParseOK(t *testing.T, want *Node, pkg, name string, isource ...string) {
	t.Helper()
	exit.TestableExit()
	Debug = true

	// fix edges missing in the literal values
	for _, n0 := range want.Nodes {
		n0.Parent = want
		for _, n1 := range n0.Nodes {
			n1.Parent = n0
		}
	}

	files := make([]Source, len(isource))

	for i, s := range isource {
		// allow nested back-ticks
		source := strings.Replace(s, "|", "`", -1)
		files[i] = Source{fmt.Sprintf("issue%d.go", i), bytes.NewBufferString(source)}
	}

	err := parseAllFiles(files)
	if err != nil {
		t.Errorf("Error parsing: %s", err)
	}

	got, err := findMatchingNodes(pkg, name)
	if err != nil {
		t.Errorf("Error parsing: %s", err)
	}

	if !reflect.DeepEqual(got, want) {
		ex := utter.Sdump(want)
		ac := utter.Sdump(got)
		outputDiff(ex, pkg+"-expected.txt")
		outputDiff(ac, pkg+"-got.txt")
		t.Errorf("Wanted %s\nGot %s", ex, ac)
	}
}

//-------------------------------------------------------------------------------------------------

func outputDiff(a, name string) {
	f, err := os.Create(name)
	if err != nil {
		panic(err)
	}
	f.WriteString(a)
	f.WriteString("\n")
	err = f.Close()
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package e2e

import (
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	image          = "gcr.io/google_containers/serve_hostname:1.1"
	simulationTime = 20 * time.Minute
	smallRCSize    = 5
	mediumRCSize   = 30
	bigRCSize      = 250
)

// This test suite can take a long time to run, so by default it is added to
// the ginkgo.skip list (see driver.go).
// To run this suite you must explicitly ask for it by setting the
// -t/--test flag or ginkgo.focus flag.
var _ = Describe("Load", func() {
	var c *client.Client
	var nodeCount int
	var ns string

	BeforeEach(func() {
		var err error
		c, err = loadClient()
		expectNoError(err)
		nodes, err := c.Nodes().List(labels.Everything(), fields.Everything())
		expectNoError(err)
		nodeCount = len(nodes.Items)
		Expect(nodeCount).NotTo(BeZero())
		nsForTesting, err := createTestingNS("load", c)
		ns = nsForTesting.Name
		expectNoError(err)
	})

	// TODO add flag that allows to skip cleanup on failure
	AfterEach(func() {
		By(fmt.Sprintf("Destroying namespace for this suite %v", ns))
		if err := c.Namespaces().Delete(ns); err != nil {
			Failf("Couldn't delete ns %s", err)
		}
	})

	type Load struct {
		skip        bool
		podsPerNode int
	}

	loadTests := []Load{
		{podsPerNode: 30, skip: true},
	}

	for _, testArg := range loadTests {
		name := fmt.Sprintf("should be able to handle %v pods per node", testArg.podsPerNode)
		if testArg.skip {
			name = "[Skipped] " + name
		}

		It(name, func() {
			totalPods := testArg.podsPerNode * nodeCount
			smallRCCount, mediumRCCount, bigRCCount := computeRCCounts(totalPods)
			threads := smallRCCount + mediumRCCount + bigRCCount

			var wg sync.WaitGroup
			wg.Add(threads)

			// Run RC load for all kinds of RC.
			runRCLoad(c, &wg, ns, smallRCSize, smallRCCount)
			runRCLoad(c, &wg, ns, mediumRCSize, mediumRCCount)
			runRCLoad(c, &wg, ns, bigRCSize, bigRCCount)

			// Wait for all the pods from all the RC's to return.
			wg.Wait()
			// TODO verify latency metrics
		})
	}
})

func computeRCCounts(total int) (int, int, int) {
	// Small RCs owns ~0.5 of total number of pods, medium and big RCs ~0.25 each.
	// For example for 3000 pods (100 nodes, 30 pods per node) there are:
	//  - 500 small RCs each 5 pods
	//  - 25 medium RCs each 30 pods
	//  - 3 big RCs each 250 pods
	bigRCCount := total / 4 / bigRCSize
	mediumRCCount := (total - bigRCCount*bigRCSize) / 3 / mediumRCSize
	smallRCCount := (total - bigRCCount*bigRCSize - mediumRCCount*mediumRCSize) / smallRCSize
	return smallRCCount, mediumRCCount, bigRCCount
}

// The function creates a RC and then every few second resize it and with 0.1 probability deletes it.
func playWithRC(c *client.Client, wg *sync.WaitGroup, ns string, size int) {
	defer GinkgoRecover()
	defer wg.Done()
	rcExist := false
	var name string
	// Once every 1-2 minutes perform resize of RC.
	for start := time.Now(); time.Since(start) < simulationTime; time.Sleep(time.Duration(60+rand.Intn(60)) * time.Second) {
		if !rcExist {
			name = "load-test-" + string(util.NewUUID())
			expectNoError(RunRC(c, name, ns, image, size))
			rcExist = true
		}
		// Resize RC to a random size between 0.5x and 1.5x of the original size.
		newSize := uint(rand.Intn(size+1) + size/2)
		expectNoError(ResizeRC(c, ns, name, newSize))
		// With probability 0.1 remove this RC.
		if rand.Intn(10) == 0 {
			expectNoError(DeleteRC(c, ns, name))
			rcExist = false
		}
	}
	if rcExist {
		expectNoError(DeleteRC(c, ns, name))
	}
}

func runRCLoad(c *client.Client, wg *sync.WaitGroup, ns string, size, count int) {
	By(fmt.Sprintf("Running %v Replication Controllers with size %v and playing with them", count, size))
	for i := 0; i < count; i++ {
		go playWithRC(c, wg, ns, size)
	}
}
<commit_msg>Enabled load e2e test for gce-scalability<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package e2e

import (
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	image          = "gcr.io/google_containers/serve_hostname:1.1"
	simulationTime = 20 * time.Minute
	smallRCSize    = 5
	mediumRCSize   = 30
	bigRCSize      = 250
)

// This test suite can take a long time to run, so by default it is added to
// the ginkgo.skip list (see driver.go).
// To run this suite you must explicitly ask for it by setting the
// -t/--test flag or ginkgo.focus flag.
var _ = Describe("Load", func() {
	var c *client.Client
	var nodeCount int
	var ns string

	BeforeEach(func() {
		var err error
		c, err = loadClient()
		expectNoError(err)
		nodes, err := c.Nodes().List(labels.Everything(), fields.Everything())
		expectNoError(err)
		nodeCount = len(nodes.Items)
		Expect(nodeCount).NotTo(BeZero())
		nsForTesting, err := createTestingNS("load", c)
		ns = nsForTesting.Name
		expectNoError(err)
	})

	// TODO add flag that allows to skip cleanup on failure
	AfterEach(func() {
		By(fmt.Sprintf("Destroying namespace for this suite %v", ns))
		if err := c.Namespaces().Delete(ns); err != nil {
			Failf("Couldn't delete ns %s", err)
		}
	})

	type Load struct {
		podsPerNode int
	}

	loadTests := []Load{
		{podsPerNode: 30},
	}

	for _, testArg := range loadTests {
		name := fmt.Sprintf("[Performance suite] [Skipped] should be able to handle %v pods per node", testArg.podsPerNode)

		It(name, func() {
			totalPods := testArg.podsPerNode * nodeCount
			smallRCCount, mediumRCCount, bigRCCount := computeRCCounts(totalPods)
			threads := smallRCCount + mediumRCCount + bigRCCount

			var wg sync.WaitGroup
			wg.Add(threads)

			// Run RC load for all kinds of RC.
			runRCLoad(c, &wg, ns, smallRCSize, smallRCCount)
			runRCLoad(c, &wg, ns, mediumRCSize, mediumRCCount)
			runRCLoad(c, &wg, ns, bigRCSize, bigRCCount)

			// Wait for all the pods from all the RC's to return.
			wg.Wait()
			// TODO verify latency metrics
		})
	}
})

func computeRCCounts(total int) (int, int, int) {
	// Small RCs owns ~0.5 of total number of pods, medium and big RCs ~0.25 each.
	// For example for 3000 pods (100 nodes, 30 pods per node) there are:
	//  - 500 small RCs each 5 pods
	//  - 25 medium RCs each 30 pods
	//  - 3 big RCs each 250 pods
	bigRCCount := total / 4 / bigRCSize
	mediumRCCount := (total - bigRCCount*bigRCSize) / 3 / mediumRCSize
	smallRCCount := (total - bigRCCount*bigRCSize - mediumRCCount*mediumRCSize) / smallRCSize
	return smallRCCount, mediumRCCount, bigRCCount
}

// The function creates a RC and then every few second resize it and with 0.1 probability deletes it.
func playWithRC(c *client.Client, wg *sync.WaitGroup, ns string, size int) {
	defer GinkgoRecover()
	defer wg.Done()
	rcExist := false
	var name string
	// Once every 1-2 minutes perform resize of RC.
	for start := time.Now(); time.Since(start) < simulationTime; time.Sleep(time.Duration(60+rand.Intn(60)) * time.Second) {
		if !rcExist {
			name = "load-test-" + string(util.NewUUID())
			expectNoError(RunRC(c, name, ns, image, size))
			rcExist = true
		}
		// Resize RC to a random size between 0.5x and 1.5x of the original size.
		newSize := uint(rand.Intn(size+1) + size/2)
		expectNoError(ResizeRC(c, ns, name, newSize))
		// With probability 0.1 remove this RC.
		if rand.Intn(10) == 0 {
			expectNoError(DeleteRC(c, ns, name))
			rcExist = false
		}
	}
	if rcExist {
		expectNoError(DeleteRC(c, ns, name))
	}
}

func runRCLoad(c *client.Client, wg *sync.WaitGroup, ns string, size, count int) {
	By(fmt.Sprintf("Running %v Replication Controllers with size %v and playing with them", count, size))
	for i := 0; i < count; i++ {
		go playWithRC(c, wg, ns, size)
	}
}
<|endoftext|>


<commit_before>package auctionrunnerdelegate_test

import (
	"errors"

	"github.com/cloudfoundry-incubator/auction/auctiontypes"
	"github.com/cloudfoundry-incubator/bbs/fake_bbs"
	"github.com/cloudfoundry-incubator/bbs/models"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/cloudfoundry-incubator/rep/repfakes"
	"github.com/cloudfoundry/dropsonde/metric_sender/fake"
	"github.com/cloudfoundry/dropsonde/metrics"

	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/lager/lagertest"

	"github.com/cloudfoundry-incubator/auctioneer/auctionrunnerdelegate"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Auction Runner Delegate", func() {
	var (
		delegate         *auctionrunnerdelegate.AuctionRunnerDelegate
		bbsClient        *fake_bbs.FakeClient
		metricSender     *fake.FakeMetricSender
		repClientFactory *repfakes.FakeClientFactory
		repClient        *repfakes.FakeClient
		logger           lager.Logger
	)

	BeforeEach(func() {
		metricSender = fake.NewFakeMetricSender()
		metrics.Initialize(metricSender, nil)

		bbsClient = &fake_bbs.FakeClient{}
		repClientFactory = &repfakes.FakeClientFactory{}
		repClient = &repfakes.FakeClient{}
		repClientFactory.CreateClientReturns(repClient)
		logger = lagertest.NewTestLogger("delegate")

		delegate = auctionrunnerdelegate.New(repClientFactory, bbsClient, logger)
	})

	Describe("fetching cell reps", func() {
		Context("when the BSS succeeds", func() {
			BeforeEach(func() {
				cellPresence1 := models.NewCellPresence("cell-A", "cell-a.url", "zone-1", models.NewCellCapacity(123, 456, 789), []string{}, []string{})
				cellPresence2 := models.NewCellPresence("cell-B", "cell-b.url", "zone-1", models.NewCellCapacity(123, 456, 789), []string{}, []string{})
				cells := []*models.CellPresence{&cellPresence1, &cellPresence2}

				bbsClient.CellsReturns(cells, nil)
			})

			It("creates rep clients with the correct addresses", func() {
				_, err := delegate.FetchCellReps()
				Expect(err).NotTo(HaveOccurred())
				Expect(repClientFactory.CreateClientCallCount()).To(Equal(2))
				urls := []string{
					repClientFactory.CreateClientArgsForCall(0),
					repClientFactory.CreateClientArgsForCall(1),
				}
				Expect(urls).To(ConsistOf("cell-a.url", "cell-b.url"))
			})

			It("returns correctly configured auction_http_clients", func() {
				reps, err := delegate.FetchCellReps()
				Expect(err).NotTo(HaveOccurred())
				Expect(reps).To(HaveLen(2))
				Expect(reps).To(HaveKey("cell-A"))
				Expect(reps).To(HaveKey("cell-B"))

				Expect(reps["cell-A"]).To(Equal(repClient))
				Expect(reps["cell-B"]).To(Equal(repClient))
			})
		})

		Context("when the BBS errors", func() {
			BeforeEach(func() {
				bbsClient.CellsReturns(nil, errors.New("boom"))
			})

			It("should error", func() {
				cells, err := delegate.FetchCellReps()
				Expect(err).To(MatchError(errors.New("boom")))
				Expect(cells).To(BeEmpty())
			})
		})
	})

	Describe("when batches are distributed", func() {
		var results auctiontypes.AuctionResults

		BeforeEach(func() {
			resource := rep.NewResource(10, 10, "linux", []string{})

			results = auctiontypes.AuctionResults{
				SuccessfulLRPs: []auctiontypes.LRPAuction{
					{
						LRP: rep.NewLRP(models.NewActualLRPKey("successful-start", 0, "domain"), resource),
					},
				},
				SuccessfulTasks: []auctiontypes.TaskAuction{
					{
						Task: rep.NewTask("successful-task", "domain", resource),
					},
				},
				FailedLRPs: []auctiontypes.LRPAuction{
					{
						LRP:           rep.NewLRP(models.NewActualLRPKey("insufficient-capacity", 0, "domain"), resource),
						AuctionRecord: auctiontypes.AuctionRecord{PlacementError: rep.ErrorInsufficientResources.Error()},
					},
					{
						LRP:           rep.NewLRP(models.NewActualLRPKey("incompatible-stacks", 0, "domain"), resource),
						AuctionRecord: auctiontypes.AuctionRecord{PlacementError: auctiontypes.ErrorCellMismatch.Error()},
					},
				},
				FailedTasks: []auctiontypes.TaskAuction{
					{
						Task:          rep.NewTask("failed-task", "domain", resource),
						AuctionRecord: auctiontypes.AuctionRecord{PlacementError: rep.ErrorInsufficientResources.Error()},
					},
				},
			}

			delegate.AuctionCompleted(results)
		})

		It("should mark all failed tasks as COMPLETE with the appropriate failure reason", func() {
			Expect(bbsClient.FailTaskCallCount()).To(Equal(1))
			taskGuid, failureReason := bbsClient.FailTaskArgsForCall(0)
			Expect(taskGuid).To(Equal("failed-task"))
			Expect(failureReason).To(Equal(rep.ErrorInsufficientResources.Error()))
		})

		It("should mark all failed LRPs as UNCLAIMED with the appropriate placement error", func() {
			Expect(bbsClient.FailActualLRPCallCount()).To(Equal(2))
			lrpKey, errorMessage := bbsClient.FailActualLRPArgsForCall(0)
			Expect(*lrpKey).To(Equal(models.NewActualLRPKey("insufficient-capacity", 0, "domain")))
			Expect(errorMessage).To(Equal(rep.ErrorInsufficientResources.Error()))

			lrpKey1, errorMessage1 := bbsClient.FailActualLRPArgsForCall(1)
			Expect(*lrpKey1).To(Equal(models.NewActualLRPKey("incompatible-stacks", 0, "domain")))
			Expect(errorMessage1).To(Equal(auctiontypes.ErrorCellMismatch.Error()))
		})
	})
})
<commit_msg>Rename fake_bbs.FakeClient->fake_bbs.FakeInternalClient<commit_after>package auctionrunnerdelegate_test

import (
	"errors"

	"github.com/cloudfoundry-incubator/auction/auctiontypes"
	"github.com/cloudfoundry-incubator/bbs/fake_bbs"
	"github.com/cloudfoundry-incubator/bbs/models"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/cloudfoundry-incubator/rep/repfakes"
	"github.com/cloudfoundry/dropsonde/metric_sender/fake"
	"github.com/cloudfoundry/dropsonde/metrics"

	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/lager/lagertest"

	"github.com/cloudfoundry-incubator/auctioneer/auctionrunnerdelegate"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Auction Runner Delegate", func() {
	var (
		delegate         *auctionrunnerdelegate.AuctionRunnerDelegate
		bbsClient        *fake_bbs.FakeInternalClient
		metricSender     *fake.FakeMetricSender
		repClientFactory *repfakes.FakeClientFactory
		repClient        *repfakes.FakeClient
		logger           lager.Logger
	)

	BeforeEach(func() {
		metricSender = fake.NewFakeMetricSender()
		metrics.Initialize(metricSender, nil)

		bbsClient = &fake_bbs.FakeInternalClient{}
		repClientFactory = &repfakes.FakeClientFactory{}
		repClient = &repfakes.FakeClient{}
		repClientFactory.CreateClientReturns(repClient)
		logger = lagertest.NewTestLogger("delegate")

		delegate = auctionrunnerdelegate.New(repClientFactory, bbsClient, logger)
	})

	Describe("fetching cell reps", func() {
		Context("when the BSS succeeds", func() {
			BeforeEach(func() {
				cellPresence1 := models.NewCellPresence("cell-A", "cell-a.url", "zone-1", models.NewCellCapacity(123, 456, 789), []string{}, []string{})
				cellPresence2 := models.NewCellPresence("cell-B", "cell-b.url", "zone-1", models.NewCellCapacity(123, 456, 789), []string{}, []string{})
				cells := []*models.CellPresence{&cellPresence1, &cellPresence2}

				bbsClient.CellsReturns(cells, nil)
			})

			It("creates rep clients with the correct addresses", func() {
				_, err := delegate.FetchCellReps()
				Expect(err).NotTo(HaveOccurred())
				Expect(repClientFactory.CreateClientCallCount()).To(Equal(2))
				urls := []string{
					repClientFactory.CreateClientArgsForCall(0),
					repClientFactory.CreateClientArgsForCall(1),
				}
				Expect(urls).To(ConsistOf("cell-a.url", "cell-b.url"))
			})

			It("returns correctly configured auction_http_clients", func() {
				reps, err := delegate.FetchCellReps()
				Expect(err).NotTo(HaveOccurred())
				Expect(reps).To(HaveLen(2))
				Expect(reps).To(HaveKey("cell-A"))
				Expect(reps).To(HaveKey("cell-B"))

				Expect(reps["cell-A"]).To(Equal(repClient))
				Expect(reps["cell-B"]).To(Equal(repClient))
			})
		})

		Context("when the BBS errors", func() {
			BeforeEach(func() {
				bbsClient.CellsReturns(nil, errors.New("boom"))
			})

			It("should error", func() {
				cells, err := delegate.FetchCellReps()
				Expect(err).To(MatchError(errors.New("boom")))
				Expect(cells).To(BeEmpty())
			})
		})
	})

	Describe("when batches are distributed", func() {
		var results auctiontypes.AuctionResults

		BeforeEach(func() {
			resource := rep.NewResource(10, 10, "linux", []string{})

			results = auctiontypes.AuctionResults{
				SuccessfulLRPs: []auctiontypes.LRPAuction{
					{
						LRP: rep.NewLRP(models.NewActualLRPKey("successful-start", 0, "domain"), resource),
					},
				},
				SuccessfulTasks: []auctiontypes.TaskAuction{
					{
						Task: rep.NewTask("successful-task", "domain", resource),
					},
				},
				FailedLRPs: []auctiontypes.LRPAuction{
					{
						LRP:           rep.NewLRP(models.NewActualLRPKey("insufficient-capacity", 0, "domain"), resource),
						AuctionRecord: auctiontypes.AuctionRecord{PlacementError: rep.ErrorInsufficientResources.Error()},
					},
					{
						LRP:           rep.NewLRP(models.NewActualLRPKey("incompatible-stacks", 0, "domain"), resource),
						AuctionRecord: auctiontypes.AuctionRecord{PlacementError: auctiontypes.ErrorCellMismatch.Error()},
					},
				},
				FailedTasks: []auctiontypes.TaskAuction{
					{
						Task:          rep.NewTask("failed-task", "domain", resource),
						AuctionRecord: auctiontypes.AuctionRecord{PlacementError: rep.ErrorInsufficientResources.Error()},
					},
				},
			}

			delegate.AuctionCompleted(results)
		})

		It("should mark all failed tasks as COMPLETE with the appropriate failure reason", func() {
			Expect(bbsClient.FailTaskCallCount()).To(Equal(1))
			taskGuid, failureReason := bbsClient.FailTaskArgsForCall(0)
			Expect(taskGuid).To(Equal("failed-task"))
			Expect(failureReason).To(Equal(rep.ErrorInsufficientResources.Error()))
		})

		It("should mark all failed LRPs as UNCLAIMED with the appropriate placement error", func() {
			Expect(bbsClient.FailActualLRPCallCount()).To(Equal(2))
			lrpKey, errorMessage := bbsClient.FailActualLRPArgsForCall(0)
			Expect(*lrpKey).To(Equal(models.NewActualLRPKey("insufficient-capacity", 0, "domain")))
			Expect(errorMessage).To(Equal(rep.ErrorInsufficientResources.Error()))

			lrpKey1, errorMessage1 := bbsClient.FailActualLRPArgsForCall(1)
			Expect(*lrpKey1).To(Equal(models.NewActualLRPKey("incompatible-stacks", 0, "domain")))
			Expect(errorMessage1).To(Equal(auctiontypes.ErrorCellMismatch.Error()))
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2012, The gohg Authors. All rights reserved.
// Use of this source code is governed by a BSD style license
// that can be found in the LICENSE.txt file.

// This program is an example of how to use the gohg library.

package main

import (
	"bitbucket.org/gohg/gohg/gohg_lib"
	"fmt"
	"log"
)

func main() {
	// Set var hgexe to whatever is appropriate for your situation.
	// You can also change it to test with different versions of Mercurial.
	// hgexe := "M:/DEV/hg-stable/hg"
	hgexe := "hg"
	// repo := "C:/DEV/thg-qt-work"
	// repo = "C:/DEV/go/src/bitbucket.org/gohg/gohg"
	repo := "."

	fmt.Println("========== Begin of example1 ==========")

	var err error
	fmt.Printf("Using Mercurial repo at: %s\n", repo)
	fmt.Println("--------------------")

	hc := gohg_lib.NewHgClient()
	var cfg []string
	if err = hc.Connect(hgexe, repo, cfg); err != nil {
		log.Fatal(err)
	}
	defer hc.Close()
	defer func() { fmt.Println("========== End of example1 ==========") }()

	var v string
	if v, err = hc.Version(); err != nil {
		// fmt.Println(err)
		// return
		log.Println(err)
	}
	fmt.Printf("[[version]]: %s\n", v)

	fmt.Println("--------------------")

	var s string
	if s, err = hc.Summary(); err != nil {
		fmt.Println(err)
		return
	}
	// use 'go run example1.go | less' (or more) to view big results (such as a full log)
	fmt.Printf("[[Summary]]:\n%s", s)

	// give time to see the Hg CS session live and die from Process Explorer
	// fmt.Print("waiting...")
	// time.Sleep(3 * time.Second)
}
<commit_msg>example1: added command with options<commit_after>// Copyright 2012, The gohg Authors. All rights reserved.
// Use of this source code is governed by a BSD style license
// that can be found in the LICENSE.txt file.

// This program is an example of how to use the gohg library.

package main

import (
	"bitbucket.org/gohg/gohg/gohg_lib"
	"fmt"
	"log"
)

func main() {
	// Set var hgexe to whatever is appropriate for your situation.
	// You can also change it to test with different versions of Mercurial.
	// hgexe := "M:/DEV/hg-stable/hg"
	hgexe := "hg"
	// repo := "C:/DEV/thg-qt-work"
	// repo = "C:/DEV/go/src/bitbucket.org/gohg/gohg"
	repo := "."

	fmt.Println("========== Begin of example1 ==========")

	var err error
	fmt.Printf("Using Mercurial repo at: %s\n", repo)
	fmt.Println("--------------------")

	hc := gohg_lib.NewHgClient()
	var cfg []string
	if err = hc.Connect(hgexe, repo, cfg); err != nil {
		log.Fatal(err)
	}
	defer hc.Close()
	defer func() { fmt.Println("========== End of example1 ==========") }()

	var v string
	if v, err = hc.Version(); err != nil {
		// fmt.Println(err)
		// return
		log.Println(err)
	}
	fmt.Printf("[[version]]: %s\n", v)

	fmt.Println("--------------------")

	var s string
	if s, err = hc.Summary([]string{}); err != nil {
		fmt.Println(err)
		return
	}
	fmt.Printf("[[summary]]:\n%s", s)

	fmt.Println("--------------------")

	var l []byte
	if l, err = hc.Log([]string{"-l", "2"}); err != nil {
		fmt.Println(err)
		return
	}
	// use 'go run example1.go | less' (or more) to view big results (such as a full log)
	fmt.Printf("[[log -l 2]]:\n%s", l)

	// give time to see the Hg CS session live and die from Process Explorer
	// fmt.Print("waiting...")
	// time.Sleep(3 * time.Second)
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Command surfer scrapes the signal status page of a SB6121 or SB6183 cable
// modem and exports values as prometheus metrics.
package main

import (
	"context"
	"flag"
	"net/http"
	_ "net/http/pprof"
	"strconv"
	"time"

	"github.com/golang/glog"
	"github.com/golang/groupcache/singleflight"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/wathiede/surfer/modem"
	_ "github.com/wathiede/surfer/modem/sb6121"
	_ "github.com/wathiede/surfer/modem/sb6183"
)

var (
	port         = flag.Int("port", 6666, "port to listen on when serving prometheus metrics")
	timeout      = flag.Duration("timeout", 1*time.Second, "timeout for HTTP get to cable modem")
	fakeDataPath = flag.String("fake", "", "path to fake HTML data.  (default) fetch over HTTP")

	downstreamSNRMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "downstream_snr",
		Help: "Downstream signal-to-noise ratio in dB",
	},
		[]string{"channel", "frequency_hz", "modulation"},
	)
	downstreamPowerLevelMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "downstream_power_level",
		Help: "Downstream power level reading in dBmV",
	},
		[]string{"channel", "frequency_hz", "modulation"},
	)

	codewordsUnerroredMetric = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "codewords_unerrored",
		Help: "Unerrored codeword count",
	},
		[]string{"channel"},
	)
	codewordsCorrectableMetric = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "codewords_correctable",
		Help: "Correctable codeword count",
	},
		[]string{"channel"},
	)
	codewordsUncorrectableMetric = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "codewords_uncorrectable",
		Help: "Uncorrectable codeword count",
	},
		[]string{"channel"},
	)

	upstreamSymbolRateMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "upstream_symbol_rate",
		Help: "Upstream symbol rate in sym/sec",
	},
		[]string{"channel", "frequency_hz", "modulation", "ranging_status"},
	)
	upstreamPowerLevelMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "upstream_power_level",
		Help: "Upstream power level reading in dBmV",
	},
		[]string{"channel", "frequency_hz", "modulation", "ranging_status"},
	)

	fetchErrorsMetric = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "fetch_errors",
		Help: "Count of errors when fetching metrics from modem.",
	})

	fetchSuccessesMetric = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "fetch_successes",
		Help: "Count of successes when fetching metrics from modem.",
	})
)

func init() {
	prometheus.MustRegister(downstreamSNRMetric)
	prometheus.MustRegister(downstreamPowerLevelMetric)
	prometheus.MustRegister(upstreamSymbolRateMetric)
	prometheus.MustRegister(upstreamPowerLevelMetric)
	prometheus.MustRegister(codewordsUnerroredMetric)
	prometheus.MustRegister(codewordsCorrectableMetric)
	prometheus.MustRegister(codewordsUncorrectableMetric)
	prometheus.MustRegister(fetchErrorsMetric)
	prometheus.MustRegister(fetchSuccessesMetric)
}

func main() {
	flag.Parse()
	defer glog.Flush()

	ctx := context.Background()
	var m modem.Modem
	for {
		ctx, cancel := context.WithTimeout(ctx, *timeout)
		m = modem.New(ctx, *fakeDataPath)
		cancel()
		if m != nil {
			break
		}
		glog.Infof("Failed to find modem, sleeping")
		time.Sleep(5 * time.Second)
	}
	glog.Infof("Found modem %q", m.Name())

	g := &singleflight.Group{}
	ph := prometheus.Handler()
	// Refresh data every prometheus poll.
	http.Handle("/metrics", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Only make one query to the cable modem if concurrent requests come in.
		if _, err := g.Do("get", func() (interface{}, error) {
			ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
			defer cancel()
			s, err := m.Status(ctx)
			if err != nil {
				fetchErrorsMetric.Inc()
				return nil, err
			}
			for ch, d := range s.Downstream {
				downstreamSNRMetric.WithLabelValues(string(ch), d.Frequency, d.Modulation).Set(d.SNR)
				downstreamPowerLevelMetric.WithLabelValues(string(ch), d.Frequency, d.Modulation).Set(d.PowerLevel)
				codewordsUnerroredMetric.WithLabelValues(string(ch)).Set(d.Unerrored)
				codewordsCorrectableMetric.WithLabelValues(string(ch)).Set(d.Correctable)
				codewordsUncorrectableMetric.WithLabelValues(string(ch)).Set(d.Uncorrectable)
			}

			for ch, u := range s.Upstream {
				upstreamSymbolRateMetric.WithLabelValues(string(ch), u.Frequency, u.Modulation, u.Status).Set(u.SymbolRate)
				upstreamPowerLevelMetric.WithLabelValues(string(ch), u.Frequency, u.Modulation, u.Status).Set(u.PowerLevel)
			}
			fetchSuccessesMetric.Inc()
			return nil, nil
		}); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		ph.ServeHTTP(w, r)
	}))
	glog.Fatalf("Listener returned: %v", http.ListenAndServe(":"+strconv.Itoa(*port), nil))
}
<commit_msg>correct the help text for the timeout flag<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Command surfer scrapes the signal status page of a SB6121 or SB6183 cable
// modem and exports values as prometheus metrics.
package main

import (
	"context"
	"flag"
	"net/http"
	_ "net/http/pprof"
	"strconv"
	"time"

	"github.com/golang/glog"
	"github.com/golang/groupcache/singleflight"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/wathiede/surfer/modem"
	_ "github.com/wathiede/surfer/modem/sb6121"
	_ "github.com/wathiede/surfer/modem/sb6183"
)

var (
	port         = flag.Int("port", 6666, "port to listen on when serving prometheus metrics")
	timeout      = flag.Duration("timeout", 1*time.Second, "timeout for the HTTP GET to cable modem")
	fakeDataPath = flag.String("fake", "", "path to fake HTML data.  (default) fetch over HTTP")

	downstreamSNRMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "downstream_snr",
		Help: "Downstream signal-to-noise ratio in dB",
	},
		[]string{"channel", "frequency_hz", "modulation"},
	)
	downstreamPowerLevelMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "downstream_power_level",
		Help: "Downstream power level reading in dBmV",
	},
		[]string{"channel", "frequency_hz", "modulation"},
	)

	codewordsUnerroredMetric = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "codewords_unerrored",
		Help: "Unerrored codeword count",
	},
		[]string{"channel"},
	)
	codewordsCorrectableMetric = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "codewords_correctable",
		Help: "Correctable codeword count",
	},
		[]string{"channel"},
	)
	codewordsUncorrectableMetric = prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "codewords_uncorrectable",
		Help: "Uncorrectable codeword count",
	},
		[]string{"channel"},
	)

	upstreamSymbolRateMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "upstream_symbol_rate",
		Help: "Upstream symbol rate in sym/sec",
	},
		[]string{"channel", "frequency_hz", "modulation", "ranging_status"},
	)
	upstreamPowerLevelMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "upstream_power_level",
		Help: "Upstream power level reading in dBmV",
	},
		[]string{"channel", "frequency_hz", "modulation", "ranging_status"},
	)

	fetchErrorsMetric = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "fetch_errors",
		Help: "Count of errors when fetching metrics from modem.",
	})

	fetchSuccessesMetric = prometheus.NewCounter(prometheus.CounterOpts{
		Name: "fetch_successes",
		Help: "Count of successes when fetching metrics from modem.",
	})
)

func init() {
	prometheus.MustRegister(downstreamSNRMetric)
	prometheus.MustRegister(downstreamPowerLevelMetric)
	prometheus.MustRegister(upstreamSymbolRateMetric)
	prometheus.MustRegister(upstreamPowerLevelMetric)
	prometheus.MustRegister(codewordsUnerroredMetric)
	prometheus.MustRegister(codewordsCorrectableMetric)
	prometheus.MustRegister(codewordsUncorrectableMetric)
	prometheus.MustRegister(fetchErrorsMetric)
	prometheus.MustRegister(fetchSuccessesMetric)
}

func main() {
	flag.Parse()
	defer glog.Flush()

	ctx := context.Background()
	var m modem.Modem
	for {
		ctx, cancel := context.WithTimeout(ctx, *timeout)
		m = modem.New(ctx, *fakeDataPath)
		cancel()
		if m != nil {
			break
		}
		glog.Infof("Failed to find modem, sleeping")
		time.Sleep(5 * time.Second)
	}
	glog.Infof("Found modem %q", m.Name())

	g := &singleflight.Group{}
	ph := prometheus.Handler()
	// Refresh data every prometheus poll.
	http.Handle("/metrics", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Only make one query to the cable modem if concurrent requests come in.
		if _, err := g.Do("get", func() (interface{}, error) {
			ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
			defer cancel()
			s, err := m.Status(ctx)
			if err != nil {
				fetchErrorsMetric.Inc()
				return nil, err
			}
			for ch, d := range s.Downstream {
				downstreamSNRMetric.WithLabelValues(string(ch), d.Frequency, d.Modulation).Set(d.SNR)
				downstreamPowerLevelMetric.WithLabelValues(string(ch), d.Frequency, d.Modulation).Set(d.PowerLevel)
				codewordsUnerroredMetric.WithLabelValues(string(ch)).Set(d.Unerrored)
				codewordsCorrectableMetric.WithLabelValues(string(ch)).Set(d.Correctable)
				codewordsUncorrectableMetric.WithLabelValues(string(ch)).Set(d.Uncorrectable)
			}

			for ch, u := range s.Upstream {
				upstreamSymbolRateMetric.WithLabelValues(string(ch), u.Frequency, u.Modulation, u.Status).Set(u.SymbolRate)
				upstreamPowerLevelMetric.WithLabelValues(string(ch), u.Frequency, u.Modulation, u.Status).Set(u.PowerLevel)
			}
			fetchSuccessesMetric.Inc()
			return nil, nil
		}); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		ph.ServeHTTP(w, r)
	}))
	glog.Fatalf("Listener returned: %v", http.ListenAndServe(":"+strconv.Itoa(*port), nil))
}
<|endoftext|>