<commit_before>package main

import (
	"bytes"
	"container/list"
	"errors"
	"fmt"
	"strings"

	"github.com/Jeffail/gabs"
)

// tokenize takes a string delimited by slahes and splits it up into tokens
// returns a linkedlist.
func tokenize(path string) *list.List {

	// Creates the list.
	l := list.New()
	for _, tok := range strings.Split(path, "/") {
		l.PushBack(tok)
	}
	return l
}

// Returns error if neither value present.
func getPrefix(c *gabs.Container) (string, error) {
	d := c.Path(expandKey).Data()
	if d != nil {
		s, oks := d.(string)
		i, oki := d.(float64)
		if oks {
			return fmt.Sprintf("%s", s), nil
		} else if oki {
			return fmt.Sprintf("%.f", i), nil
		} else {
			return "", errors.New(fmt.Sprintf("Unexpected type of expansion value, got %T instead of int or string.", d))
		}
	}
	q := c.Path(queryKey).Data()
	if q != nil {
		if s, ok := q.(string); ok {
			return s, nil
		} else {
			return "", errors.New(fmt.Sprintf("Casting query key to string failed for %T:%v", q, q))
		}
	}

	return "", errors.New(fmt.Sprintf("error in config, no expand or query key in %s\n", c.String()))

}
func expand(c *gabs.Container, token *list.Element, res *bytes.Buffer) {
	// base case
	if token == nil {
		return
	}
	res.WriteString("/")
	children, _ := c.ChildrenMap()
	for key, child := range children {
		if key == token.Value {
			p, err := getPrefix(child)
			if err != nil {
				fmt.Println(err.Error())
				return
			}
			res.WriteString(p)
			expand(child, token.Next(), res)
			return
		}
	}
	// handle base case if no keys matched
	res.WriteString(token.Value.(string))

	// if tokens left over, append the rest
	for e := token.Next(); e != nil; e = e.Next() {
		res.WriteString("/")
		res.WriteString(e.Value.(string))
	}
}
<commit_msg>clean up linter errors<commit_after>package main

import (
	"bytes"
	"container/list"
	"fmt"
	"strings"

	"github.com/Jeffail/gabs"
)

// tokenize takes a string delimited by slahes and splits it up into tokens
// returns a linkedlist.
func tokenize(path string) *list.List {

	// Creates the list.
	l := list.New()
	for _, tok := range strings.Split(path, "/") {
		l.PushBack(tok)
	}
	return l
}

// Returns error if neither value present.
func getPrefix(c *gabs.Container) (string, error) {
	d := c.Path(expandKey).Data()
	if d != nil {
		s, oks := d.(string)
		i, oki := d.(float64)
		if oks {
			return fmt.Sprintf("%s", s), nil
		} else if oki {
			return fmt.Sprintf("%.f", i), nil
		}
		return "", fmt.Errorf("unexpected type of expansion value, got %T instead of int or string", d)
	}
	q := c.Path(queryKey).Data()
	if q != nil {
		if s, ok := q.(string); ok {
			return s, nil
		}
		return "", fmt.Errorf("Casting query key to string failed for %T:%v", q, q)
	}
	return "", fmt.Errorf("error in config, no expand or query key in %s", c.String())
}
func expand(c *gabs.Container, token *list.Element, res *bytes.Buffer) {
	// base case
	if token == nil {
		return
	}
	res.WriteString("/")
	children, _ := c.ChildrenMap()
	for key, child := range children {
		if key == token.Value {
			p, err := getPrefix(child)
			if err != nil {
				fmt.Println(err.Error())
				return
			}
			res.WriteString(p)
			expand(child, token.Next(), res)
			return
		}
	}
	// handle base case if no keys matched
	res.WriteString(token.Value.(string))

	// if tokens left over, append the rest
	for e := token.Next(); e != nil; e = e.Next() {
		res.WriteString("/")
		res.WriteString(e.Value.(string))
	}
}
<|endoftext|>


<commit_before>package thor

import (
	"io"
	"strconv"

	"github.com/skelterjohn/gopp"
)

const grammar = `
ignore: /^#.*\n/
ignore: /^\s+/

Doc => {type=Node} {field=Key} {/} <List>
List => {field=Type} {list} {field=Kids} <<Node>>*
Node => {type=Node} {field=Key} <id> ':' <Value>
Node => {type=Node} <Value>
Value => {type=Node} '{' <List> '}'
Value => {type=Node} {field=Type} {num} {field=Val} <num>
Value => {type=Node} {field=Type} {bool} {field=Val} <bool>
Value => {type=Node} {field=Type} {str} {field=Val} <str>

num = /([-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?)/
bool = /(true|false)/
str = /"((?:[^"\\]|\\.)*)"/
id = /([\pL][\pL\pN\-_]*)/
`

var (
	decFact *gopp.DecoderFactory
)

func init() {
	var err error
	decFact, err = gopp.NewDecoderFactory(grammar, "Doc")
	if err != nil {
		panic(err)
	}
	decFact.RegisterType(Node{})
}

type Node struct {
	Key  string
	Val  string
	Type string
	Kids []Node
}

func (n Node) String() string {
	s := ""
	switch n.Type {
	case "list":
		for i := range n.Kids {
			s += n.Kids[i].String()
			if i < len(n.Kids)-1 {
				s += " "
			}
		}
	case "str":
		s = strconv.Quote(n.Val)
	default:
		s = n.Val
	}
	if n.Type == "list" && n.Key != "/" {
		s = "{" + s + "}"
	}
	if n.Key != "" && n.Key != "/" {
		s = n.Key + ":" + s
	}
	return s
}

func Parse(in io.Reader) (*Node, error) {
	dec := decFact.NewDecoder(in)
	tree := &Node{}
	err := dec.Decode(tree)
	return tree, err
}
<commit_msg>Changed comment regex to handle no EOL at EOF.<commit_after>package thor

import (
	"io"
	"strconv"

	"github.com/skelterjohn/gopp"
)

const grammar = `
ignore: /^#.*/
ignore: /^\s+/

Doc => {type=Node} {field=Key} {/} <List>
List => {field=Type} {list} {field=Kids} <<Node>>*
Node => {type=Node} {field=Key} <id> ':' <Value>
Node => {type=Node} <Value>
Value => {type=Node} '{' <List> '}'
Value => {type=Node} {field=Type} {num} {field=Val} <num>
Value => {type=Node} {field=Type} {bool} {field=Val} <bool>
Value => {type=Node} {field=Type} {str} {field=Val} <str>

num = /([-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?)/
bool = /(true|false)/
str = /"((?:[^"\\]|\\.)*)"/
id = /([\pL][\pL\pN\-_]*)/
`

var (
	decFact *gopp.DecoderFactory
)

func init() {
	var err error
	decFact, err = gopp.NewDecoderFactory(grammar, "Doc")
	if err != nil {
		panic(err)
	}
	decFact.RegisterType(Node{})
}

type Node struct {
	Key  string
	Val  string
	Type string
	Kids []Node
}

func (n Node) String() string {
	s := ""
	switch n.Type {
	case "list":
		for i := range n.Kids {
			s += n.Kids[i].String()
			if i < len(n.Kids)-1 {
				s += " "
			}
		}
	case "str":
		s = strconv.Quote(n.Val)
	default:
		s = n.Val
	}
	if n.Type == "list" && n.Key != "/" {
		s = "{" + s + "}"
	}
	if n.Key != "" && n.Key != "/" {
		s = n.Key + ":" + s
	}
	return s
}

func Parse(in io.Reader) (*Node, error) {
	dec := decFact.NewDecoder(in)
	tree := &Node{}
	err := dec.Decode(tree)
	return tree, err
}
<|endoftext|>


<commit_before>package pogs

import (
	"reflect"
	"testing"

	"zombiezen.com/go/capnproto2"
	air "zombiezen.com/go/capnproto2/internal/aircraftlib"
)

type VerVal struct {
	Val int16
}

type VerValNoTag struct {
	Val int16
}

type VerValTag1 struct {
	Val int16 `capnp:"val"`
}

type VerValTag2 struct {
	Val int16 `capnp:"val"`
}

type VerValTag3 struct {
	Val int16 `capnp:"val"`
}

type VerOneData struct {
	VerVal
}

type VerTwoData struct {
	*VerVal
	Duo int64
}

type VerTwoDataOmit struct {
	VerVal `capnp:"-"`
	Duo    int64
}

type VerOneDataTop struct {
	Val int16
	VerVal
}

type VerOneDataTopWithLowerCollision struct {
	Val int16
	VerVal
	VerValNoTag
}

type VerOneDataNoTags struct {
	VerVal
	VerValNoTag
}

type VerOneData1Tag struct {
	VerVal
	VerValTag1
}

type VerOneData1TagWithUntagged struct {
	VerVal
	VerValTag1
	VerValNoTag
}

type VerOneData2Tags struct {
	VerValTag1
	VerValTag2
}

type VerOneData3Tags struct {
	VerValTag1
	VerValTag2
	VerValTag3
}

type F16 struct {
	PlaneBase `capnp:"base"`
}

func TestExtract_Embed(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v1, err := air.NewRootVerOneData(seg)
	if err != nil {
		t.Fatalf("NewRootVerOneData: %v", err)
	}
	v1.SetVal(123)
	out := new(VerOneData)
	if err := Extract(out, air.VerOneData_TypeID, v1.Struct); err != nil {
		t.Errorf("Extract error: %v", err)
	}
	if out.Val != 123 {
		t.Errorf("Extract produced %s; want %s", zpretty.Sprint(out), zpretty.Sprint(&VerOneData{VerVal{123}}))
	}
}

func TestExtract_EmbedPtr(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	v2.SetVal(123)
	v2.SetDuo(456)
	out := new(VerTwoData)
	if err := Extract(out, air.VerTwoData_TypeID, v2.Struct); err != nil {
		t.Errorf("Extract error: %v", err)
	}
	if out.VerVal == nil || out.Val != 123 || out.Duo != 456 {
		t.Errorf("Extract produced %s; want %s", zpretty.Sprint(out), zpretty.Sprint(&VerTwoData{&VerVal{123}, 456}))
	}
}

func TestExtract_EmbedOmit(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	v2.SetVal(123)
	v2.SetDuo(456)
	out := new(VerTwoDataOmit)
	if err := Extract(out, air.VerTwoData_TypeID, v2.Struct); err != nil {
		t.Errorf("Extract error: %v", err)
	}
	if out.Val != 0 || out.Duo != 456 {
		t.Errorf("Extract produced %s; want %s", zpretty.Sprint(out), zpretty.Sprint(&VerTwoDataOmit{VerVal{}, 456}))
	}
}

func TestExtract_EmbedName(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	f16, err := air.NewRootF16(seg)
	if err != nil {
		t.Fatalf("NewRootF16: %v", err)
	}
	base, err := f16.NewBase()
	if err != nil {
		t.Fatalf("F16.NewBase: %v", err)
	}
	if err := base.SetName("ALL YOUR BASE"); err != nil {
		t.Fatalf("Planebase.SetName: %v", err)
	}
	base.SetRating(5)
	base.SetCanFly(true)

	out := new(F16)
	if err := Extract(out, air.F16_TypeID, f16.Struct); err != nil {
		t.Errorf("Extract error: %v", err)
	}
	if out.Name != "ALL YOUR BASE" || out.Rating != 5 || !out.CanFly {
		t.Errorf("Extract produced %s; want %s", zpretty.Sprint(out), zpretty.Sprint(&F16{PlaneBase{Name: "ALL YOUR BASE", Rating: 5, CanFly: true}}))
	}
}

func TestExtract_EmbedCollide(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v1, err := air.NewRootVerOneData(seg)
	if err != nil {
		t.Fatalf("NewRootVerOneData: %v", err)
	}
	v1.SetVal(123)

	tests := []struct {
		name string
		want interface{}
	}{
		{"top", &VerOneDataTop{Val: 123}},
		{"no tags", &VerOneDataNoTags{}},
		{"1 tag", &VerOneData1Tag{VerValTag1: VerValTag1{123}}},
		{"1 tag with untagged", &VerOneData1TagWithUntagged{VerValTag1: VerValTag1{123}}},
		{"2 tags", &VerOneData2Tags{}},
		{"3 tags", &VerOneData3Tags{}},
		{"top with lower collision", &VerOneDataTopWithLowerCollision{Val: 123}},
	}
	for _, test := range tests {
		out := reflect.New(reflect.TypeOf(test.want).Elem()).Interface()
		if err := Extract(out, air.VerOneData_TypeID, v1.Struct); err != nil {
			t.Errorf("%s: Extract error: %v", test.name, err)
		}
		if !reflect.DeepEqual(out, test.want) {
			t.Errorf("%s: Extract produced %s; want %s", test.name, zpretty.Sprint(out), zpretty.Sprint(test.want))
		}
	}
}

func TestInsert_Embed(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v1, err := air.NewRootVerOneData(seg)
	if err != nil {
		t.Fatalf("NewRootVerOneData: %v", err)
	}
	gv1 := &VerOneData{VerVal{123}}
	err = Insert(air.VerOneData_TypeID, v1.Struct, gv1)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(gv1), err)
	}
	if v1.Val() != 123 {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(gv1), v1)
	}
}

func TestInsert_EmbedPtr(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	gv2 := &VerTwoData{&VerVal{123}, 456}
	err = Insert(air.VerTwoData_TypeID, v2.Struct, gv2)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(gv2), err)
	}
	if v2.Val() != 123 || v2.Duo() != 456 {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(gv2), v2)
	}
}

func TestInsert_EmbedNilPtr(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	gv2 := &VerTwoData{nil, 456}
	err = Insert(air.VerTwoData_TypeID, v2.Struct, gv2)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(gv2), err)
	}
	if v2.Val() != 0 || v2.Duo() != 456 {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(gv2), v2)
	}
}

func TestInsert_EmbedOmit(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	in := &VerTwoDataOmit{VerVal{123}, 456}
	err = Insert(air.VerTwoData_TypeID, v2.Struct, in)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(in), err)
	}
	if v2.Val() != 0 || v2.Duo() != 456 {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(in), v2)
	}
}

func TestInsert_EmbedNamed(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	f16, err := air.NewRootF16(seg)
	if err != nil {
		t.Fatalf("NewRootF16: %v", err)
	}
	in := &F16{PlaneBase{Name: "ALL YOUR BASE", Rating: 5, CanFly: true}}
	err = Insert(air.F16_TypeID, f16.Struct, in)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(in), err)
	}
	base, err := f16.Base()
	if err != nil {
		t.Errorf("f16.base: %v", err)
	}
	name, err := base.Name()
	if err != nil {
		t.Errorf("f16.base.name: %v", err)
	}
	if name != "ALL YOUR BASE" || base.Rating() != 5 || !base.CanFly() {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(in), f16)
	}
}

func TestInsert_EmbedCollide(t *testing.T) {
	tests := []struct {
		name string
		in   interface{}
		val  int16
	}{
		{"top", &VerOneDataTop{Val: 123, VerVal: VerVal{456}}, 123},
		{"no tags", &VerOneDataNoTags{VerVal: VerVal{123}, VerValNoTag: VerValNoTag{456}}, 0},
		{"1 tag", &VerOneData1Tag{VerVal: VerVal{123}, VerValTag1: VerValTag1{456}}, 456},
		{
			"1 tag with untagged",
			&VerOneData1TagWithUntagged{
				VerVal:      VerVal{123},
				VerValTag1:  VerValTag1{456},
				VerValNoTag: VerValNoTag{789},
			},
			456,
		},
		{"2 tags", &VerOneData2Tags{VerValTag1: VerValTag1{123}, VerValTag2: VerValTag2{456}}, 0},
		{"3 tags", &VerOneData3Tags{VerValTag1: VerValTag1{123}, VerValTag2: VerValTag2{456}, VerValTag3: VerValTag3{789}}, 0},
		{
			"top with lower collision",
			&VerOneDataTopWithLowerCollision{
				Val:         123,
				VerVal:      VerVal{456},
				VerValNoTag: VerValNoTag{789},
			},
			123,
		},
	}
	for _, test := range tests {
		_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
		if err != nil {
			t.Errorf("%s: NewMessage: %v", test.name, err)
			continue
		}
		v1, err := air.NewRootVerOneData(seg)
		if err != nil {
			t.Errorf("%s: NewRootVerOneData: %v", test.name, err)
			continue
		}
		err = Insert(air.VerOneData_TypeID, v1.Struct, test.in)
		if err != nil {
			t.Errorf("%s: Insert(..., %s): %v", test.name, zpretty.Sprint(test.in))
		}
		if v1.Val() != test.val {
			t.Errorf("%s: Insert(..., %s) produced %v; want (val = %d)", test.name, zpretty.Sprint(test.in), v1, test.val)
		}
	}
}
<commit_msg>pogs: rearrange collision tests in file<commit_after>package pogs

import (
	"reflect"
	"testing"

	"zombiezen.com/go/capnproto2"
	air "zombiezen.com/go/capnproto2/internal/aircraftlib"
)

type VerVal struct {
	Val int16
}

type VerOneData struct {
	VerVal
}

type VerTwoData struct {
	*VerVal
	Duo int64
}

type F16 struct {
	PlaneBase `capnp:"base"`
}

func TestExtract_Embed(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v1, err := air.NewRootVerOneData(seg)
	if err != nil {
		t.Fatalf("NewRootVerOneData: %v", err)
	}
	v1.SetVal(123)
	out := new(VerOneData)
	if err := Extract(out, air.VerOneData_TypeID, v1.Struct); err != nil {
		t.Errorf("Extract error: %v", err)
	}
	if out.Val != 123 {
		t.Errorf("Extract produced %s; want %s", zpretty.Sprint(out), zpretty.Sprint(&VerOneData{VerVal{123}}))
	}
}

func TestExtract_EmbedPtr(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	v2.SetVal(123)
	v2.SetDuo(456)
	out := new(VerTwoData)
	if err := Extract(out, air.VerTwoData_TypeID, v2.Struct); err != nil {
		t.Errorf("Extract error: %v", err)
	}
	if out.VerVal == nil || out.Val != 123 || out.Duo != 456 {
		t.Errorf("Extract produced %s; want %s", zpretty.Sprint(out), zpretty.Sprint(&VerTwoData{&VerVal{123}, 456}))
	}
}

func TestExtract_EmbedOmit(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	v2.SetVal(123)
	v2.SetDuo(456)
	out := new(VerTwoDataOmit)
	if err := Extract(out, air.VerTwoData_TypeID, v2.Struct); err != nil {
		t.Errorf("Extract error: %v", err)
	}
	if out.Val != 0 || out.Duo != 456 {
		t.Errorf("Extract produced %s; want %s", zpretty.Sprint(out), zpretty.Sprint(&VerTwoDataOmit{VerVal{}, 456}))
	}
}

func TestExtract_EmbedName(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	f16, err := air.NewRootF16(seg)
	if err != nil {
		t.Fatalf("NewRootF16: %v", err)
	}
	base, err := f16.NewBase()
	if err != nil {
		t.Fatalf("F16.NewBase: %v", err)
	}
	if err := base.SetName("ALL YOUR BASE"); err != nil {
		t.Fatalf("Planebase.SetName: %v", err)
	}
	base.SetRating(5)
	base.SetCanFly(true)

	out := new(F16)
	if err := Extract(out, air.F16_TypeID, f16.Struct); err != nil {
		t.Errorf("Extract error: %v", err)
	}
	if out.Name != "ALL YOUR BASE" || out.Rating != 5 || !out.CanFly {
		t.Errorf("Extract produced %s; want %s", zpretty.Sprint(out), zpretty.Sprint(&F16{PlaneBase{Name: "ALL YOUR BASE", Rating: 5, CanFly: true}}))
	}
}

func TestInsert_Embed(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v1, err := air.NewRootVerOneData(seg)
	if err != nil {
		t.Fatalf("NewRootVerOneData: %v", err)
	}
	gv1 := &VerOneData{VerVal{123}}
	err = Insert(air.VerOneData_TypeID, v1.Struct, gv1)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(gv1), err)
	}
	if v1.Val() != 123 {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(gv1), v1)
	}
}

func TestInsert_EmbedPtr(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	gv2 := &VerTwoData{&VerVal{123}, 456}
	err = Insert(air.VerTwoData_TypeID, v2.Struct, gv2)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(gv2), err)
	}
	if v2.Val() != 123 || v2.Duo() != 456 {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(gv2), v2)
	}
}

func TestInsert_EmbedNilPtr(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	gv2 := &VerTwoData{nil, 456}
	err = Insert(air.VerTwoData_TypeID, v2.Struct, gv2)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(gv2), err)
	}
	if v2.Val() != 0 || v2.Duo() != 456 {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(gv2), v2)
	}
}

func TestInsert_EmbedOmit(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v2, err := air.NewRootVerTwoData(seg)
	if err != nil {
		t.Fatalf("NewRootVerTwoData: %v", err)
	}
	in := &VerTwoDataOmit{VerVal{123}, 456}
	err = Insert(air.VerTwoData_TypeID, v2.Struct, in)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(in), err)
	}
	if v2.Val() != 0 || v2.Duo() != 456 {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(in), v2)
	}
}

func TestInsert_EmbedNamed(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	f16, err := air.NewRootF16(seg)
	if err != nil {
		t.Fatalf("NewRootF16: %v", err)
	}
	in := &F16{PlaneBase{Name: "ALL YOUR BASE", Rating: 5, CanFly: true}}
	err = Insert(air.F16_TypeID, f16.Struct, in)
	if err != nil {
		t.Errorf("Insert(%s) error: %v", zpretty.Sprint(in), err)
	}
	base, err := f16.Base()
	if err != nil {
		t.Errorf("f16.base: %v", err)
	}
	name, err := base.Name()
	if err != nil {
		t.Errorf("f16.base.name: %v", err)
	}
	if name != "ALL YOUR BASE" || base.Rating() != 5 || !base.CanFly() {
		t.Errorf("Insert(%s) produced %v", zpretty.Sprint(in), f16)
	}
}

type VerValNoTag struct {
	Val int16
}

type VerValTag1 struct {
	Val int16 `capnp:"val"`
}

type VerValTag2 struct {
	Val int16 `capnp:"val"`
}

type VerValTag3 struct {
	Val int16 `capnp:"val"`
}

type VerTwoDataOmit struct {
	VerVal `capnp:"-"`
	Duo    int64
}

type VerOneDataTop struct {
	Val int16
	VerVal
}

type VerOneDataTopWithLowerCollision struct {
	Val int16
	VerVal
	VerValNoTag
}

type VerOneDataNoTags struct {
	VerVal
	VerValNoTag
}

type VerOneData1Tag struct {
	VerVal
	VerValTag1
}

type VerOneData1TagWithUntagged struct {
	VerVal
	VerValTag1
	VerValNoTag
}

type VerOneData2Tags struct {
	VerValTag1
	VerValTag2
}

type VerOneData3Tags struct {
	VerValTag1
	VerValTag2
	VerValTag3
}

func TestExtract_EmbedCollide(t *testing.T) {
	_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	v1, err := air.NewRootVerOneData(seg)
	if err != nil {
		t.Fatalf("NewRootVerOneData: %v", err)
	}
	v1.SetVal(123)

	tests := []struct {
		name string
		want interface{}
	}{
		{"top", &VerOneDataTop{Val: 123}},
		{"no tags", &VerOneDataNoTags{}},
		{"1 tag", &VerOneData1Tag{VerValTag1: VerValTag1{123}}},
		{"1 tag with untagged", &VerOneData1TagWithUntagged{VerValTag1: VerValTag1{123}}},
		{"2 tags", &VerOneData2Tags{}},
		{"3 tags", &VerOneData3Tags{}},
		{"top with lower collision", &VerOneDataTopWithLowerCollision{Val: 123}},
	}
	for _, test := range tests {
		out := reflect.New(reflect.TypeOf(test.want).Elem()).Interface()
		if err := Extract(out, air.VerOneData_TypeID, v1.Struct); err != nil {
			t.Errorf("%s: Extract error: %v", test.name, err)
		}
		if !reflect.DeepEqual(out, test.want) {
			t.Errorf("%s: Extract produced %s; want %s", test.name, zpretty.Sprint(out), zpretty.Sprint(test.want))
		}
	}
}

func TestInsert_EmbedCollide(t *testing.T) {
	tests := []struct {
		name string
		in   interface{}
		val  int16
	}{
		{"top", &VerOneDataTop{Val: 123, VerVal: VerVal{456}}, 123},
		{"no tags", &VerOneDataNoTags{VerVal: VerVal{123}, VerValNoTag: VerValNoTag{456}}, 0},
		{"1 tag", &VerOneData1Tag{VerVal: VerVal{123}, VerValTag1: VerValTag1{456}}, 456},
		{
			"1 tag with untagged",
			&VerOneData1TagWithUntagged{
				VerVal:      VerVal{123},
				VerValTag1:  VerValTag1{456},
				VerValNoTag: VerValNoTag{789},
			},
			456,
		},
		{"2 tags", &VerOneData2Tags{VerValTag1: VerValTag1{123}, VerValTag2: VerValTag2{456}}, 0},
		{"3 tags", &VerOneData3Tags{VerValTag1: VerValTag1{123}, VerValTag2: VerValTag2{456}, VerValTag3: VerValTag3{789}}, 0},
		{
			"top with lower collision",
			&VerOneDataTopWithLowerCollision{
				Val:         123,
				VerVal:      VerVal{456},
				VerValNoTag: VerValNoTag{789},
			},
			123,
		},
	}
	for _, test := range tests {
		_, seg, err := capnp.NewMessage(capnp.SingleSegment(nil))
		if err != nil {
			t.Errorf("%s: NewMessage: %v", test.name, err)
			continue
		}
		v1, err := air.NewRootVerOneData(seg)
		if err != nil {
			t.Errorf("%s: NewRootVerOneData: %v", test.name, err)
			continue
		}
		err = Insert(air.VerOneData_TypeID, v1.Struct, test.in)
		if err != nil {
			t.Errorf("%s: Insert(..., %s): %v", test.name, zpretty.Sprint(test.in))
		}
		if v1.Val() != test.val {
			t.Errorf("%s: Insert(..., %s) produced %v; want (val = %d)", test.name, zpretty.Sprint(test.in), v1, test.val)
		}
	}
}
<|endoftext|>


<commit_before>package yalzo

import (
	"bytes"
	"fmt"
	"github.com/mattn/go-runewidth"
	"os"
)

const (
	LABEL_TEXT_WIDTH = 16
)

type TodoList struct {
	todos  []Todo
	archs  []Todo
	labels []string
	file   *os.File
}

type Todo struct {
	label      string
	title      string
	isArchived bool
	no         int
}

type Tab int

const (
	TODO Tab = iota
	ARCHIVE
)

func (t Tab) String() string {
	switch t {
	case TODO:
		return "Todo"
	case ARCHIVE:
		return "Archive"
	}
	return "Unknown"
}

func NewTodoList(fp *os.File, ls []string) *TodoList {
	l, as, err := ReadCSV(fp)
	if err != nil {
		panic(err)
	}
	fp.Seek(0, 0)

	return &TodoList{
		todos:  l,
		archs:  as,
		labels: ls,
		file:   fp,
	}
}

func (tl *TodoList) GetList(width int, tab Tab) []string {
	lines := make([]string, 0, 100)
	lists := tl.getListInTab(tab)
	for i := 0; i < len(lists); i++ {
		lines = append(lines, lists[i].createTodoText(width))
	}
	return lines
}

func containsStr(ary []string, val string) (int, bool) {
	for i, v := range ary {
		if v == val {
			return i, true
		}
	}

	return 0, false
}

func (tl *TodoList) GetLabels() []string {
	return tl.labels
}

func (tl *TodoList) AddLabel(label string) {
	if _, is := containsStr(tl.labels, label); !is {
		tl.labels = append(tl.labels, label)
	}
}

func (tl *TodoList) RemoveLavel(label string) {
	if i, is := containsStr(tl.labels, label); is {
		tl.labels = append(tl.labels[:i], tl.labels[i-1:]...)
	}
}

func (tl *TodoList) Save() {
	tl.file.Truncate(0)
	SaveCSV(append(tl.todos, tl.archs...), tl.file)
}

func (tl *TodoList) ChangeTitle(i int, t string, tab Tab) {
	switch tab {
	case TODO:
		(*tl).todos[i].title = t
	case ARCHIVE:
		(*tl).archs[i].title = t
	}
}

func (tl *TodoList) ChangeLabelName(i int, l string, tab Tab) {
	switch tab {
	case TODO:
		(*tl).todos[i].label = l
	case ARCHIVE:
		(*tl).archs[i].label = l
	}
}

func (tl *TodoList) Delete(n int) {
	for i := n; i < len((*tl).todos); i++ {
		(*tl).todos[i].setNumber(i)
	}
	tl.todos = append(tl.todos[:n], tl.todos[n+1:]...)
}

func (tl *TodoList) AddTodo(t string) int {
	no := len(tl.todos) + 1
	tl.todos = append(tl.todos, Todo{
		no:    no,
		label: "",
		title: t,
	})
	return no - 1
}

func (tl *TodoList) MoveTodo(n int, t Tab) {
	from, to := &tl.todos, &tl.archs
	isArched := true

	switch t {
	case ARCHIVE:
		from, to = to, from
		isArched = !isArched
	}

	length := len(*to)
	(*from)[n].isArchived = isArched
	(*from)[n].setNumber(length + 1)
	*to = append(*to, (*from)[n])
	for i := n + 1; i < len(*from); i++ {
		(*from)[i].setNumber(i)
	}
	*from = append((*from)[:n], (*from)[n+1:]...)
}

func (tl *TodoList) Exchange(i1 int, i2 int, tab Tab) {
	switch tab {
	case TODO:
		tl.todos[i2].setNumber(i1 + 1)
		tl.todos[i1].setNumber(i2 + 1)
		tl.todos[i2], tl.todos[i1] = tl.todos[i1], tl.todos[i2]
	case ARCHIVE:
		tl.archs[i2].setNumber(i1 + 1)
		tl.archs[i1].setNumber(i2 + 1)
		tl.archs[i2], tl.archs[i1] = tl.archs[i1], tl.archs[i2]
	}
}

func (t *Todo) setNumber(n int) {
	(*t).no = n
}

func (tl *TodoList) getListInTab(tab Tab) []Todo {
	switch tab {
	case ARCHIVE:
		return tl.archs
	case TODO:
		return tl.todos
	}
	return []Todo{}
}

func (t *Todo) createTodoText(limit int) string {
	num_s := fmt.Sprintf("%3d", t.no)
	label_s := t.label
	label_len := runewidth.StringWidth(label_s)
	if label_len > LABEL_TEXT_WIDTH {
		label_s = runewidth.Truncate(label_s, LABEL_TEXT_WIDTH, "")
	} else {
		padding := repeatSpace((LABEL_TEXT_WIDTH - label_len) / 2)
		buf := bytes.NewBuffer(padding)
		buf.Write([]byte(label_s))
		buf.Write(padding)
		label_s = buf.String()

		// if label string lenght is odd and more than LABEL_TEXT_WIDTH
		if runewidth.StringWidth(label_s) == LABEL_TEXT_WIDTH+1 {
			label_s = runewidth.Truncate(label_s, LABEL_TEXT_WIDTH, "")
		}
	}

	str := num_s + " [ " + label_s + " ] " + t.title
	length := runewidth.StringWidth(str)
	if length > limit {
		return runewidth.Truncate(str, limit, "")
	} else {
		str_len := runewidth.StringWidth(str)
		buf := bytes.NewBufferString(str)
		for str_len < limit {
			str_len++
			buf.Write([]byte(" "))
		}
		return buf.String()
	}
}

func repeatSpace(cnt int) []byte {
	buf := bytes.NewBuffer(make([]byte, 0, cnt))
	for cnt-1 >= 0 {
		buf.Write([]byte(" "))
		cnt--
	}
	return buf.Bytes()
}
<commit_msg>todolistにvalidationの追加<commit_after>package yalzo

import (
	"bytes"
	"fmt"
	"github.com/mattn/go-runewidth"
	"os"
)

const (
	LABEL_TEXT_WIDTH = 16
)

type TodoList struct {
	todos  []Todo
	archs  []Todo
	labels []string
	file   *os.File
}

type Todo struct {
	label      string
	title      string
	isArchived bool
	no         int
}

type Tab int

const (
	TODO Tab = iota
	ARCHIVE
)

func (t Tab) String() string {
	switch t {
	case TODO:
		return "Todo"
	case ARCHIVE:
		return "Archive"
	}
	return "Unknown"
}

func NewTodoList(fp *os.File, ls []string) *TodoList {
	l, as, err := ReadCSV(fp)
	if err != nil {
		panic(err)
	}
	fp.Seek(0, 0)

	return &TodoList{
		todos:  l,
		archs:  as,
		labels: ls,
		file:   fp,
	}
}

func (tl *TodoList) GetList(width int, tab Tab) []string {
	lines := make([]string, 0, 100)
	lists := tl.getListInTab(tab)
	for i := 0; i < len(lists); i++ {
		lines = append(lines, lists[i].createTodoText(width))
	}
	return lines
}

func containsStr(ary []string, val string) (int, bool) {
	for i, v := range ary {
		if v == val {
			return i, true
		}
	}

	return 0, false
}

func (tl *TodoList) GetLabels() []string {
	return tl.labels
}

func (tl *TodoList) AddLabel(label string) {
	if _, is := containsStr(tl.labels, label); !is {
		tl.labels = append(tl.labels, label)
	}
}

func (tl *TodoList) RemoveLavel(label string) {
	if i, is := containsStr(tl.labels, label); is {
		tl.labels = append(tl.labels[:i], tl.labels[i-1:]...)
	}
}

func (tl *TodoList) Save() {
	tl.file.Truncate(0)
	SaveCSV(append(tl.todos, tl.archs...), tl.file)
}

func (tl *TodoList) ChangeTitle(i int, t string, tab Tab) {
	ls := tl.getListInTab(tab)
	if len(ls)-1 < i {
		return
	}

	ls[i].title = t
}

func (tl *TodoList) ChangeLabelName(i int, l string, tab Tab) {
	ls := tl.getListInTab(tab)
	if len(ls)-1 < i {
		return
	}

	ls[i].label = l
}

func (tl *TodoList) Delete(n int) {
	if len((*tl).todos)-1 < n {
		return
	}

	for i := n; i < len((*tl).todos); i++ {
		(*tl).todos[i].setNumber(i)
	}
	tl.todos = append(tl.todos[:n], tl.todos[n+1:]...)
}

func (tl *TodoList) AddTodo(t string) int {
	no := len(tl.todos) + 1
	tl.todos = append(tl.todos, Todo{
		no:    no,
		label: "",
		title: t,
	})
	return no - 1
}

func (tl *TodoList) MoveTodo(n int, t Tab) {
	from, to := &tl.todos, &tl.archs
	isArched := true

	switch t {
	case ARCHIVE:
		from, to = to, from
		isArched = !isArched
	}

	if len(*from)-1 < n {
		return
	}

	length := len(*to)
	(*from)[n].isArchived = isArched
	(*from)[n].setNumber(length + 1)
	*to = append(*to, (*from)[n])
	for i := n + 1; i < len(*from); i++ {
		(*from)[i].setNumber(i)
	}
	*from = append((*from)[:n], (*from)[n+1:]...)
}

func (tl *TodoList) Exchange(i1 int, i2 int, tab Tab) {
	ls := tl.getListInTab(tab)
	ls[i2].setNumber(i1 + 1)
	ls[i1].setNumber(i2 + 1)
	ls[i2], ls[i1] = ls[i1], ls[i2]
}

func (t *Todo) setNumber(n int) {
	(*t).no = n
}

func (tl *TodoList) getListInTab(tab Tab) []Todo {
	switch tab {
	case ARCHIVE:
		return tl.archs
	case TODO:
		return tl.todos
	}

	return []Todo{}
}

func (t *Todo) createTodoText(limit int) string {
	num_s := fmt.Sprintf("%3d", t.no)
	label_s := t.label
	label_len := runewidth.StringWidth(label_s)
	if label_len > LABEL_TEXT_WIDTH {
		label_s = runewidth.Truncate(label_s, LABEL_TEXT_WIDTH, "")
	} else {
		padding := repeatSpace((LABEL_TEXT_WIDTH - label_len) / 2)
		buf := bytes.NewBuffer(padding)
		buf.Write([]byte(label_s))
		buf.Write(padding)
		label_s = buf.String()

		// if label string lenght is odd and more than LABEL_TEXT_WIDTH
		if runewidth.StringWidth(label_s) == LABEL_TEXT_WIDTH+1 {
			label_s = runewidth.Truncate(label_s, LABEL_TEXT_WIDTH, "")
		}
	}

	str := num_s + " [ " + label_s + " ] " + t.title
	length := runewidth.StringWidth(str)
	if length > limit {
		return runewidth.Truncate(str, limit, "")
	} else {
		str_len := runewidth.StringWidth(str)
		buf := bytes.NewBufferString(str)
		for str_len < limit {
			str_len++
			buf.Write([]byte(" "))
		}
		return buf.String()
	}
}

func repeatSpace(cnt int) []byte {
	buf := bytes.NewBuffer(make([]byte, 0, cnt))
	for cnt-1 >= 0 {
		buf.Write([]byte(" "))
		cnt--
	}
	return buf.Bytes()
}
<|endoftext|>


<commit_before>package h2spec

import (
	"fmt"
	"golang.org/x/net/http2"
)

func DataTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.1", "DATA")

	tg.AddTestCase(NewTestCase(
		"Sends a DATA frame with 0x0 stream identifier",
		"The endpoint MUST respond with a connection error of type PROTOCOL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteData(0, true, []byte("test"))

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a DATA frame on the stream that is not opend",
		"The endpoint MUST respond with a stream error of type STREAM_CLOSED.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)
			hdrs[0].Value = "POST"
			hdrs = append(hdrs, pair("content-type", "4"))

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = true
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)
			http2Conn.fr.WriteData(1, true, []byte("test"))

			actualCodes := []http2.ErrCode{http2.ErrCodeStreamClosed}
			return TestStreamError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a DATA frame with invalid pad length",
		"The endpoint MUST treat this as a connection error of type PROTOCOL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)
			hdrs[0].Value = "POST"
			hdrs = append(hdrs, pair("content-type", "4"))

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			// Data length: 5, Pad length: 6
			fmt.Fprintf(http2Conn.conn, "\x00\x00\x05\x00\x0b\x00\x00\x00\x01")
			fmt.Fprintf(http2Conn.conn, "\x06\x54\x65\x73\x74")

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestStreamError(ctx, http2Conn, actualCodes)
		},
	))

	return tg
}
<commit_msg>Typo fix: opend -> opened<commit_after>package h2spec

import (
	"fmt"
	"golang.org/x/net/http2"
)

func DataTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.1", "DATA")

	tg.AddTestCase(NewTestCase(
		"Sends a DATA frame with 0x0 stream identifier",
		"The endpoint MUST respond with a connection error of type PROTOCOL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteData(0, true, []byte("test"))

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a DATA frame on the stream that is not opened",
		"The endpoint MUST respond with a stream error of type STREAM_CLOSED.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)
			hdrs[0].Value = "POST"
			hdrs = append(hdrs, pair("content-type", "4"))

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = true
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)
			http2Conn.fr.WriteData(1, true, []byte("test"))

			actualCodes := []http2.ErrCode{http2.ErrCodeStreamClosed}
			return TestStreamError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a DATA frame with invalid pad length",
		"The endpoint MUST treat this as a connection error of type PROTOCOL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)
			hdrs[0].Value = "POST"
			hdrs = append(hdrs, pair("content-type", "4"))

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			// Data length: 5, Pad length: 6
			fmt.Fprintf(http2Conn.conn, "\x00\x00\x05\x00\x0b\x00\x00\x00\x01")
			fmt.Fprintf(http2Conn.conn, "\x06\x54\x65\x73\x74")

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestStreamError(ctx, http2Conn, actualCodes)
		},
	))

	return tg
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"os"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/catatsuy/isucon6-final/portal/job"
)

var s *httptest.Server

func TestMain(m *testing.M) {
	*startsAtHour = -1
	*endsAtHour = -1

	flag.Parse()
	err := initWeb()
	if err != nil {
		log.Fatal(err)
	}

	s = httptest.NewServer(buildMux())
	n := m.Run()
	s.Close()
	os.Exit(n)
}

type testHTTPClient struct {
	*http.Client
	*testing.T
}

func (c *testHTTPClient) Must(resp *http.Response, err error) *http.Response {
	require.NoError(c.T, err)
	return resp
}

func newTestClient(t *testing.T) *testHTTPClient {
	jar, _ := cookiejar.New(nil)
	return &testHTTPClient{
		Client: &http.Client{Jar: jar},
		T:      t,
	}
}

func TestLogin(t *testing.T) {
	resp, err := http.Get(s.URL)
	require.NoError(t, err)
	require.Equal(t, "/login", resp.Request.URL.Path)

	jar, _ := cookiejar.New(nil)
	cli := &http.Client{Jar: jar}

	resp, err = cli.PostForm(s.URL+"/login", url.Values{"team_id": {"1026"}, "password": {"dummy-pass-26"}})
	require.NoError(t, err)
	require.Equal(t, "/", resp.Request.URL.Path)
}

func readAll(r io.Reader) string {
	b, _ := ioutil.ReadAll(r)
	return string(b)
}

func benchGetJob(bench *testHTTPClient) *job.Job {
	resp := bench.Must(bench.Post(s.URL+"/mBGWHqBVEjUSKpBF/job/new", "", nil))
	if !assert.Equal(bench.T, http.StatusOK, resp.StatusCode) {
		return nil
	}

	var j job.Job
	err := json.NewDecoder(resp.Body).Decode(&j)
	require.NoError(bench.T, err)

	return &j
}

func benchPostResult(bench *testHTTPClient, j *job.Job, output *job.Output) {
	time.Sleep(1 * time.Second)

	result := job.Result{
		Job:    j,
		Output: output,
		Stderr: "",
	}
	resultJSON, err := json.Marshal(result)
	require.NoError(bench.T, err)

	resp := bench.Must(bench.Post(s.URL+"/mBGWHqBVEjUSKpBF/job/result", "application/json", bytes.NewBuffer(resultJSON)))
	require.Equal(bench.T, http.StatusOK, resp.StatusCode)
}

func cliLogin(cli *testHTTPClient, teamID int) {
	resp := cli.Must(
		cli.PostForm(
			s.URL+"/login",
			url.Values{
				"team_id":  {fmt.Sprint(teamID)},
				"password": {fmt.Sprintf("dummy-pass-%d", teamID-1000)},
			},
		),
	)
	require.Equal(cli.T, "/", resp.Request.URL.Path)
}

func TestPostJob(t *testing.T) {
	var (
		cli   = newTestClient(t)
		cli2  = newTestClient(t)
		bench = newTestClient(t)
	)

	var resp *http.Response

	// cli: ログイン
	resp = cli.Must(cli.PostForm(s.URL+"/login", url.Values{"team_id": {"1026"}, "password": {"dummy-pass-26"}}))
	require.Equal(t, "/", resp.Request.URL.Path)

	// bench: ジョブ取る
	resp = bench.Must(bench.Post(s.URL+"/mBGWHqBVEjUSKpBF/job/new", "", nil))
	require.Equal(t, http.StatusNoContent, resp.StatusCode)

	// cli: ジョブいれる→まだIP入れてないのでエラー
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	require.Equal(t, http.StatusBadRequest, resp.StatusCode)

	// cli: IP入れる
	resp = cli.Must(cli.PostForm(s.URL+"/team", url.Values{"ip_address": {"127.0.0.1"}, "instance_name": {""}}))
	assert.Contains(t, readAll(resp.Body), `<input class="form-control" type="text" name="ip_address" value="127.0.0.1" autocomplete="off">`, "IP入った表示")

	// cli: ジョブ入れる
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	assert.Contains(t, readAll(resp.Body), `<span class="label label-default">1026*</span>`, "ジョブ入った表示")

	// cli2: ログイン
	resp = cli2.Must(cli2.PostForm(s.URL+"/login", url.Values{"team_id": {"1005"}, "password": {"dummy-pass-5"}}))
	require.Equal(t, "/", resp.Request.URL.Path)
	assert.Contains(t, readAll(resp.Body), `<span class="label label-default">1026</span>`, "他人のジョブ入った表示")

	// bench: ジョブ取る
	j := benchGetJob(bench)

	// cli: ジョブいれる (2) → 入らない
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	assert.Contains(t, readAll(resp.Body), `Job already queued`)

	// cli2: IP入れる
	resp = cli2.Must(cli2.PostForm(s.URL+"/team", url.Values{"ip_address": {"127.0.0.2"}, "instance_name": {""}}))
	assert.Contains(t, readAll(resp.Body), `<input class="form-control" type="text" name="ip_address" value="127.0.0.2" autocomplete="off">`, "IP入った表示")

	// cli2: ジョブ入れる → 入る
	resp = cli2.Must(cli2.PostForm(s.URL+"/queue", nil))
	assert.Contains(t, readAll(resp.Body), `<span class="label label-default">1005*</span>`, "ジョブ入った表示")

	// bench: ジョブ取る → 放置
	j2 := benchGetJob(bench)
	_ = j2

	// cli: トップリロード
	resp = cli.Must(cli.Get(s.URL + "/"))
	assert.Contains(t, readAll(resp.Body), `<span class="label label-success">1026*</span>`, "ジョブ実行中の表示")

	// bench: 結果入れる
	benchPostResult(bench, j, &job.Output{Pass: false, Score: 5000})

	// cli: トップリロード
	resp = cli.Must(cli.Get(s.URL + "/"))
	body := readAll(resp.Body)
	require.Contains(t, body, `<th>Status</th><td>FAIL</td>`)
	require.Contains(t, body, `<th>Score</th><td>5000</td>`)
	require.Contains(t, body, `<th>Best</th><td>-</td>`)

	// cli: ジョブいれる (3)
	resp = cli.Must(cli.PostForm(s.URL+"/queue", url.Values{"ip_addr": {"127.0.0.1"}}))
	assert.NotContains(t, readAll(resp.Body), `Job already queued`)

	// bench: ジョブ取る
	j = benchGetJob(bench)

	// bench: 結果入れる
	benchPostResult(bench, j, &job.Output{Pass: true, Score: 3000})

	// cli: トップリロード
	resp = cli.Must(cli.Get(s.URL + "/"))
	body = readAll(resp.Body)
	require.Contains(t, body, `<th>Status</th><td>PASS</td>`)
	require.Contains(t, body, `<th>Score</th><td>3000</td>`)
	require.Contains(t, body, `<th>Best</th><td>3000</td>`)
	require.Regexp(t, `<td>ダミーチーム26</td>\s*<td>3000</td>`, body)
	require.NotContains(t, body, "ダミーチーム5")

	// bench: 結果入れる
	benchPostResult(bench, j2, &job.Output{Pass: true, Score: 4500})

	resp = cli2.Must(cli2.Get(s.URL + "/"))
	body = readAll(resp.Body)
	require.Contains(t, body, `<th>Status</th><td>PASS</td>`)
	require.Contains(t, body, `<th>Score</th><td>4500</td>`)
	require.Contains(t, body, `<th>Best</th><td>4500</td>`)
	require.Regexp(t, `<td>ダミーチーム5</td>\s*<td>4500</td>(?s:.*)<td>ダミーチーム26</td>\s*<td>3000</td>`, body)
}

func TestPostJobNotWithinContestTime(t *testing.T) {
	cli := newTestClient(t)
	cliLogin(cli, 1010)

	var resp *http.Response

	*startsAtHour = 24
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	assert.Equal(t, http.StatusForbidden, resp.StatusCode)
	assert.Equal(t, "Final has not started yet\n", readAll(resp.Body))
	*startsAtHour = -1

	*endsAtHour = 0
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	assert.Equal(t, http.StatusForbidden, resp.StatusCode)
	assert.Equal(t, "Final has finished\n", readAll(resp.Body))
	*endsAtHour = -1
}

func TestUpdateTeam(t *testing.T) {
	cli := newTestClient(t)
	admin := newTestClient(t)
	cliLogin(cli, 1011)

	resp := cli.Must(cli.PostForm(s.URL+"/team", url.Values{"instance_name": {"xxxxxx"}, "ip_address": {"0.0.0.0"}}))
	assert.Equal(t, http.StatusOK, resp.StatusCode)
	body := readAll(resp.Body)
	assert.Contains(t, body, `value="xxxxxx"`)
	assert.Contains(t, body, `value="0.0.0.0"`)

	resp = cli.Must(cli.Get(s.URL + "/"))
	body = readAll(resp.Body)
	assert.Contains(t, body, `value="xxxxxx"`)
	assert.Contains(t, body, `value="0.0.0.0"`)

	resp = admin.Must(admin.Get(s.URL + "/mBGWHqBVEjUSKpBF/proxy/nginx.conf"))
	require.Equal(t, http.StatusOK, resp.StatusCode)
	body = readAll(resp.Body)
	assert.Contains(t, body, `# team1011`)
	assert.Contains(t, body, `listen 11011;`)
	assert.Contains(t, body, `proxy_pass 0.0.0.0;`)
}

func TestUpdateProxies(t *testing.T) {
	cli := newTestClient(t)
	bench := newTestClient(t)
	admin := newTestClient(t)
	cliLogin(cli, 1012)

	// proxyのIP一覧を入れる
	nodes := `[{"Name":"portal","Addr":"192.168.0.10"},{"Name":"isu-proxy-1","Addr":"192.168.0.11"},{"Name":"isu-proxy-2","Addr":"192.168.0.12"}]`
	resp := admin.Must(admin.Post(s.URL+"/mBGWHqBVEjUSKpBF/proxy/update", "application/json", bytes.NewBuffer([]byte(nodes))))
	body := readAll(resp.Body)
	assert.NotContains(t, body, `192.168.0.10`, "portalのIP")
	assert.Contains(t, body, `192.168.0.11`, "proxy-1のIP")
	assert.Contains(t, body, `192.168.0.12`, "proxy-2のIP")

	// cli: IP入れる
	resp = cli.Must(cli.PostForm(s.URL+"/team", url.Values{"ip_address": {"127.0.0.1"}, "instance_name": {""}}))

	// cli: ジョブ入れる
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))

	// bench: ジョブ取る
	j := benchGetJob(bench)
	require.Equal(t, 1012, j.TeamID)
	assert.Contains(t, j.URLs, `https://192.168.0.11:11012`, "proxy-1のIP")
	assert.Contains(t, j.URLs, `https://192.168.0.12:11012`, "proxy-2のIP")
}
<commit_msg>Fix test<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"os"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/catatsuy/isucon6-final/portal/job"
)

var s *httptest.Server

func TestMain(m *testing.M) {
	*startsAtHour = -1
	*endsAtHour = -1

	flag.Parse()
	err := initWeb()
	if err != nil {
		log.Fatal(err)
	}

	s = httptest.NewServer(buildMux())
	n := m.Run()
	s.Close()
	os.Exit(n)
}

type testHTTPClient struct {
	*http.Client
	*testing.T
}

func (c *testHTTPClient) Must(resp *http.Response, err error) *http.Response {
	require.NoError(c.T, err)
	return resp
}

func newTestClient(t *testing.T) *testHTTPClient {
	jar, _ := cookiejar.New(nil)
	return &testHTTPClient{
		Client: &http.Client{Jar: jar},
		T:      t,
	}
}

func TestLogin(t *testing.T) {
	resp, err := http.Get(s.URL)
	require.NoError(t, err)
	require.Equal(t, "/login", resp.Request.URL.Path)

	jar, _ := cookiejar.New(nil)
	cli := &http.Client{Jar: jar}

	resp, err = cli.PostForm(s.URL+"/login", url.Values{"team_id": {"26"}, "password": {"p6aYuUempoticryg"}})
	require.NoError(t, err)
	require.Equal(t, "/", resp.Request.URL.Path)
}

func readAll(r io.Reader) string {
	b, _ := ioutil.ReadAll(r)
	return string(b)
}

func benchGetJob(bench *testHTTPClient) *job.Job {
	resp := bench.Must(bench.Post(s.URL+"/mBGWHqBVEjUSKpBF/job/new", "", nil))
	if !assert.Equal(bench.T, http.StatusOK, resp.StatusCode) {
		return nil
	}

	var j job.Job
	err := json.NewDecoder(resp.Body).Decode(&j)
	require.NoError(bench.T, err)

	return &j
}

func benchPostResult(bench *testHTTPClient, j *job.Job, output *job.Output) {
	time.Sleep(1 * time.Second)

	result := job.Result{
		Job:    j,
		Output: output,
		Stderr: "",
	}
	resultJSON, err := json.Marshal(result)
	require.NoError(bench.T, err)

	resp := bench.Must(bench.Post(s.URL+"/mBGWHqBVEjUSKpBF/job/result", "application/json", bytes.NewBuffer(resultJSON)))
	require.Equal(bench.T, http.StatusOK, resp.StatusCode)
}

func cliLogin(cli *testHTTPClient, teamID int, password string) {
	resp := cli.Must(
		cli.PostForm(
			s.URL+"/login",
			url.Values{
				"team_id":  {fmt.Sprint(teamID)},
				"password": {password},
			},
		),
	)
	require.Equal(cli.T, "/", resp.Request.URL.Path)
}

func TestPostJob(t *testing.T) {
	var (
		cli   = newTestClient(t)
		cli2  = newTestClient(t)
		bench = newTestClient(t)
	)

	var resp *http.Response

	// cli: ログイン
	resp = cli.Must(cli.PostForm(s.URL+"/login", url.Values{"team_id": {"26"}, "password": {"p6aYuUempoticryg"}}))
	require.Equal(t, "/", resp.Request.URL.Path)

	// bench: ジョブ取る
	resp = bench.Must(bench.Post(s.URL+"/mBGWHqBVEjUSKpBF/job/new", "", nil))
	require.Equal(t, http.StatusNoContent, resp.StatusCode)

	// cli: ジョブいれる→まだIP入れてないのでエラー
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	require.Equal(t, http.StatusBadRequest, resp.StatusCode)

	// cli: IP入れる
	resp = cli.Must(cli.PostForm(s.URL+"/team", url.Values{"ip_address": {"127.0.0.1"}, "instance_name": {""}}))
	assert.Contains(t, readAll(resp.Body), `<input class="form-control" type="text" name="ip_address" value="127.0.0.1" autocomplete="off">`, "IP入った表示")

	// cli: ジョブ入れる
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	assert.Contains(t, readAll(resp.Body), `<span class="label label-default">26*</span>`, "ジョブ入った表示")

	// cli2: ログイン
	resp = cli2.Must(cli2.PostForm(s.URL+"/login", url.Values{"team_id": {"5"}, "password": {"Y7i06XOllyJI5ogn"}}))
	require.Equal(t, "/", resp.Request.URL.Path)
	assert.Contains(t, readAll(resp.Body), `<span class="label label-default">26</span>`, "他人のジョブ入った表示")

	// bench: ジョブ取る
	j := benchGetJob(bench)

	// cli: ジョブいれる (2) → 入らない
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	assert.Contains(t, readAll(resp.Body), `Job already queued`)

	// cli2: IP入れる
	resp = cli2.Must(cli2.PostForm(s.URL+"/team", url.Values{"ip_address": {"127.0.0.2"}, "instance_name": {""}}))
	assert.Contains(t, readAll(resp.Body), `<input class="form-control" type="text" name="ip_address" value="127.0.0.2" autocomplete="off">`, "IP入った表示")

	// cli2: ジョブ入れる → 入る
	resp = cli2.Must(cli2.PostForm(s.URL+"/queue", nil))
	assert.Contains(t, readAll(resp.Body), `<span class="label label-default">5*</span>`, "ジョブ入った表示")

	// bench: ジョブ取る → 放置
	j2 := benchGetJob(bench)
	_ = j2

	// cli: トップリロード
	resp = cli.Must(cli.Get(s.URL + "/"))
	assert.Contains(t, readAll(resp.Body), `<span class="label label-success">26*</span>`, "ジョブ実行中の表示")

	// bench: 結果入れる
	benchPostResult(bench, j, &job.Output{Pass: false, Score: 5000})

	// cli: トップリロード
	resp = cli.Must(cli.Get(s.URL + "/"))
	body := readAll(resp.Body)
	require.Contains(t, body, `<th>Status</th><td>FAIL</td>`)
	require.Contains(t, body, `<th>Score</th><td>5000</td>`)
	require.Contains(t, body, `<th>Best</th><td>-</td>`)

	// cli: ジョブいれる (3)
	resp = cli.Must(cli.PostForm(s.URL+"/queue", url.Values{"ip_addr": {"127.0.0.1"}}))
	assert.NotContains(t, readAll(resp.Body), `Job already queued`)

	// bench: ジョブ取る
	j = benchGetJob(bench)

	// bench: 結果入れる
	benchPostResult(bench, j, &job.Output{Pass: true, Score: 3000})

	// cli: トップリロード
	resp = cli.Must(cli.Get(s.URL + "/"))
	body = readAll(resp.Body)
	require.Contains(t, body, `<th>Status</th><td>PASS</td>`)
	require.Contains(t, body, `<th>Score</th><td>3000</td>`)
	require.Contains(t, body, `<th>Best</th><td>3000</td>`)
	require.Regexp(t, `<td>RUDT</td>\s*<td>3000</td>`, body)
	require.NotContains(t, body, "流れ弾")

	// bench: 結果入れる
	benchPostResult(bench, j2, &job.Output{Pass: true, Score: 4500})

	resp = cli2.Must(cli2.Get(s.URL + "/"))
	body = readAll(resp.Body)
	require.Contains(t, body, `<th>Status</th><td>PASS</td>`)
	require.Contains(t, body, `<th>Score</th><td>4500</td>`)
	require.Contains(t, body, `<th>Best</th><td>4500</td>`)
	require.Regexp(t, `<td>流れ弾</td>\s*<td>4500</td>(?s:.*)<td>RUDT</td>\s*<td>3000</td>`, body)
}

func TestPostJobNotWithinContestTime(t *testing.T) {
	cli := newTestClient(t)
	cliLogin(cli, 10, "3svumTo3amDShlFy")

	var resp *http.Response

	*startsAtHour = 24
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	assert.Equal(t, http.StatusForbidden, resp.StatusCode)
	assert.Equal(t, "Final has not started yet\n", readAll(resp.Body))
	*startsAtHour = -1

	*endsAtHour = 0
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))
	assert.Equal(t, http.StatusForbidden, resp.StatusCode)
	assert.Equal(t, "Final has finished\n", readAll(resp.Body))
	*endsAtHour = -1
}

func TestUpdateTeam(t *testing.T) {
	cli := newTestClient(t)
	admin := newTestClient(t)
	cliLogin(cli, 11, "L6KZ7UJyAEtpVr9G")

	resp := cli.Must(cli.PostForm(s.URL+"/team", url.Values{"instance_name": {"xxxxxx"}, "ip_address": {"0.0.0.0"}}))
	assert.Equal(t, http.StatusOK, resp.StatusCode)
	body := readAll(resp.Body)
	assert.Contains(t, body, `value="xxxxxx"`)
	assert.Contains(t, body, `value="0.0.0.0"`)

	resp = cli.Must(cli.Get(s.URL + "/"))
	body = readAll(resp.Body)
	assert.Contains(t, body, `value="xxxxxx"`)
	assert.Contains(t, body, `value="0.0.0.0"`)

	resp = admin.Must(admin.Get(s.URL + "/mBGWHqBVEjUSKpBF/proxy/nginx.conf"))
	require.Equal(t, http.StatusOK, resp.StatusCode)
	body = readAll(resp.Body)
	assert.Contains(t, body, `# team11`)
	assert.Contains(t, body, `listen 10011;`)
	assert.Contains(t, body, `proxy_pass 0.0.0.0;`)
}

func TestUpdateProxies(t *testing.T) {
	cli := newTestClient(t)
	bench := newTestClient(t)
	admin := newTestClient(t)
	cliLogin(cli, 12, "YJUaDANoex8Y6MB")

	// proxyのIP一覧を入れる
	nodes := `[{"Name":"portal","Addr":"192.168.0.10","Status":1},{"Name":"isu-proxy-1","Addr":"192.168.0.11","Status":1},{"Name":"isu-proxy-2","Addr":"192.168.0.12","Status":1},{"Name":"isu-proxy-3","Addr":"192.168.0.13","Status":0}]`
	resp := admin.Must(admin.Post(s.URL+"/mBGWHqBVEjUSKpBF/proxy/update", "application/json", bytes.NewBuffer([]byte(nodes))))
	body := readAll(resp.Body)
	assert.NotContains(t, body, `192.168.0.10`, "portalのIP")
	assert.Contains(t, body, `192.168.0.11`, "proxy-1のIP")
	assert.Contains(t, body, `192.168.0.12`, "proxy-2のIP")
	assert.NotContains(t, body, `192.168.0.13`, "proxy-3のIP")

	// cli: IP入れる
	resp = cli.Must(cli.PostForm(s.URL+"/team", url.Values{"ip_address": {"127.0.0.1"}, "instance_name": {""}}))

	// cli: ジョブ入れる
	resp = cli.Must(cli.PostForm(s.URL+"/queue", nil))

	// bench: ジョブ取る
	j := benchGetJob(bench)
	require.Equal(t, 12, j.TeamID)
	assert.Contains(t, j.URLs, `https://192.168.0.11:10012`, "proxy-1のIP")
	assert.Contains(t, j.URLs, `https://192.168.0.12:10012`, "proxy-2のIP")
	assert.NotContains(t, j.URLs, `https://192.168.0.13:10012`, "proxy-3のIP")
}
<|endoftext|>


<commit_before>package heka_clever_plugins

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/Clever/heka-clever-plugins/postgres"
	_ "github.com/lib/pq"
	"github.com/mozilla-services/heka/message"
	. "github.com/mozilla-services/heka/pipeline"
)

type PostgresOutput struct {
	db                        *postgres.PostgresDB
	insertSchema              string
	insertTable               string
	insertMessageFields       []string
	insertTableColumns        []string
	flushInterval             uint32
	flushCount                int // Max messages before flush
	allowMissingMessageFields bool
}

type PostgresOutputConfig struct {
	// Table name and colums. Message fields to write.
	InsertSchema        string `toml:"insert_schema"`
	InsertTable         string `toml:"insert_table"`
	InsertTableColumns  string `toml:"insert_table_columns"`
	InsertMessageFields string `toml:"insert_message_fields"`
	// If a field is missing in the Heka message, allow writing NULL
	AllowMissingMessageFields bool `toml:"allow_missing_message_fields"`

	// Database Connection
	DBHost               string `toml:"db_host"`
	DBPort               int    `toml:"db_port"`
	DBName               string `toml:"db_name"`
	DBUser               string `toml:"db_user"`
	DBPassword           string `toml:"db_password"`
	DBConnectionTimeout  int    `toml:"db_connection_timeout"`
	DBMaxOpenConnections int    `toml:"db_max_open_connections"`
	DBSSLMode            string `toml:"db_ssl_mode"`

	// Interval at which accumulated messages should be written to Postgres,
	// in milliseconds (default 1000, i.e. 1 second)
	FlushInterval uint32 `toml:"flush_interval"`
	// Number of messages that triggers a write to Postgres (default 10000)
	FlushCount int `toml:"flush_count"`
}

func (po *PostgresOutput) ConfigStruct() interface{} {
	return &PostgresOutputConfig{
		AllowMissingMessageFields: true,
		DBConnectionTimeout:       5,
		DBMaxOpenConnections:      10,
		DBSSLMode:                 "require",
		FlushInterval:             uint32(1000),
		FlushCount:                10000,
		InsertSchema:              "public",
	}
}

func (po *PostgresOutput) Init(rawConf interface{}) error {
	config := rawConf.(*PostgresOutputConfig)
	po.flushInterval = config.FlushInterval
	po.flushCount = config.FlushCount
	po.insertSchema = config.InsertSchema
	po.insertTable = config.InsertTable
	if config.InsertMessageFields == "" {
		return fmt.Errorf("config item 'insert_message_fields' cannot be empty string")
	}
	po.insertMessageFields = strings.Split(config.InsertMessageFields, " ")
	if config.InsertTableColumns == "" {
		return fmt.Errorf("config item 'insert_table_columns' cannot be empty string")
	}
	po.insertTableColumns = strings.Split(config.InsertTableColumns, " ")
	po.allowMissingMessageFields = config.AllowMissingMessageFields
	p := postgres.DBConnectionParams{
		Host:           config.DBHost,
		Port:           config.DBPort,
		DBName:         config.DBName,
		User:           config.DBUser,
		Password:       config.DBPassword,
		ConnectTimeout: config.DBConnectionTimeout,
		SSLMode:        config.DBSSLMode,
	}
	db, err := postgres.New(&p)
	if err != nil {
		return err
	}
	db.SetMaxOpenConns(config.DBMaxOpenConnections)
	po.db = db
	return nil
}

func (o *PostgresOutput) Run(or OutputRunner, h PluginHelper) (err error) {
	defer o.db.Close()
	var wg sync.WaitGroup
	wg.Add(1)

	committers := o.makeCommitters(or, 5, &wg)
	go o.receiver(or, committers, &wg)

	wg.Wait()
	return
}

// Runs in a separate goroutine, accepting incoming messages, buffering output
// data until the ticker triggers the buffered data should be put onto the
// committer channel.
func (o *PostgresOutput) receiver(
	or OutputRunner, committers chan<- [][]interface{}, wg *sync.WaitGroup,
) {
	var pack *PipelinePack

	ticker := time.Tick(time.Duration(o.flushInterval) * time.Millisecond)
	batch := [][]interface{}{}

	for ok := true; ok; {
		select {
		case pack, ok = <-or.InChan():
			if !ok {
				// Closed inChan => we're shutting down, flush data
				committers <- batch
				close(committers)
				break
			}
			// Read values from message fields
			vals, err := o.convertMessageToValues(pack.Message, o.insertMessageFields)
			pack.Recycle()
			if err != nil {
				or.LogError(err)
			} else {
				batch = append(batch, vals)
				if len(batch) >= o.flushCount {
					committers <- batch
					batch = [][]interface{}{}
				}
			}
		case <-ticker:
			committers <- batch
			batch = [][]interface{}{}
		}
	}
	wg.Done()
}

// Runs in a separate goroutine, waits for buffered data on the committer
// channel, bulk inserts it into Postgres, and puts the now empty buffer on the
// return channel for reuse.
func (o *PostgresOutput) makeCommitters(
	or OutputRunner,
	count int,
	wg *sync.WaitGroup,
) chan<- [][]interface{} {
	batches := make(chan [][]interface{})

	for i := 0; i < count; i++ {
		wg.Add(1)
		go func(i int) {
			for batch := range batches {
				if len(batch) <= 0 {
					continue
				}

				done := o.commit(or, batch)
				timeout := time.NewTimer(time.Minute)

				select {
				case <-done:
					timeout.Stop()
				case <-timeout.C:
					or.LogError(errors.New("Postgres insert took more than 60s."))
				}
			}
			wg.Done()
		}(i)
	}

	return batches
}

func (o *PostgresOutput) commit(or OutputRunner, batch [][]interface{}) <-chan struct{} {
	done := make(chan struct{})

	go func() {
		err := o.db.Insert(o.insertSchema, o.insertTable, o.insertTableColumns, batch)
		if err != nil {
			or.LogError(err)
		}
		done <- struct{}{}
	}()

	return done
}

// convertMessageToValue reads a Heka Message and returns a slice of field values
func (po *PostgresOutput) convertMessageToValues(m *message.Message, insertFields []string) (fieldValues []interface{}, err error) {
	fieldValues = []interface{}{}
	missingFields := []string{}
	for _, field := range insertFields {
		// Special case: get "Timestamp" from Heka message
		if field == "Timestamp" {
			// Convert Heka time (Unix timestamp in nanoseconds) to Golang time
			v := time.Unix(0, m.GetTimestamp())
			fieldValues = append(fieldValues, v)
		} else {
			v, ok := m.GetFieldValue(field)
			if !ok {
				// If configured to do so, write NULL when a FieldValue isn't found in the Heka message
				if po.allowMissingMessageFields {
					v = nil
				} else {
					missingFields = append(missingFields, field)
					continue
				}
			}
			fieldValues = append(fieldValues, v)
		}
	}

	if len(missingFields) > 0 {
		return []interface{}{}, fmt.Errorf("message is missing expected fields: %s", strings.Join(missingFields, ", "))
	}

	return fieldValues, nil
}

func init() {
	RegisterPlugin("PostgresOutput", func() interface{} {
		return new(PostgresOutput)
	})
}
<commit_msg>Attached OutputRunner and PluginHelper to PostgressOutput object to make them more accessible.<commit_after>package heka_clever_plugins

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/Clever/heka-clever-plugins/postgres"
	_ "github.com/lib/pq"
	"github.com/mozilla-services/heka/message"
	. "github.com/mozilla-services/heka/pipeline"
)

type PostgresOutput struct {
	db *postgres.PostgresDB

	helper           PluginHelper
	runner           OutputRunner

	insertSchema              string
	insertTable               string
	insertMessageFields       []string
	insertTableColumns        []string
	flushInterval             uint32
	flushCount                int // Max messages before flush
	allowMissingMessageFields bool
}

type PostgresOutputConfig struct {
	// Table name and colums. Message fields to write.
	InsertSchema        string `toml:"insert_schema"`
	InsertTable         string `toml:"insert_table"`
	InsertTableColumns  string `toml:"insert_table_columns"`
	InsertMessageFields string `toml:"insert_message_fields"`
	// If a field is missing in the Heka message, allow writing NULL
	AllowMissingMessageFields bool `toml:"allow_missing_message_fields"`

	// Database Connection
	DBHost               string `toml:"db_host"`
	DBPort               int    `toml:"db_port"`
	DBName               string `toml:"db_name"`
	DBUser               string `toml:"db_user"`
	DBPassword           string `toml:"db_password"`
	DBConnectionTimeout  int    `toml:"db_connection_timeout"`
	DBMaxOpenConnections int    `toml:"db_max_open_connections"`
	DBSSLMode            string `toml:"db_ssl_mode"`

	// Interval at which accumulated messages should be written to Postgres,
	// in milliseconds (default 1000, i.e. 1 second)
	FlushInterval uint32 `toml:"flush_interval"`
	// Number of messages that triggers a write to Postgres (default 10000)
	FlushCount int `toml:"flush_count"`
}

func (po *PostgresOutput) ConfigStruct() interface{} {
	return &PostgresOutputConfig{
		AllowMissingMessageFields: true,
		DBConnectionTimeout:       5,
		DBMaxOpenConnections:      10,
		DBSSLMode:                 "require",
		FlushInterval:             uint32(1000),
		FlushCount:                10000,
		InsertSchema:              "public",
	}
}

func (po *PostgresOutput) Init(rawConf interface{}) error {
	config := rawConf.(*PostgresOutputConfig)
	po.flushInterval = config.FlushInterval
	po.flushCount = config.FlushCount
	po.insertSchema = config.InsertSchema
	po.insertTable = config.InsertTable
	if config.InsertMessageFields == "" {
		return fmt.Errorf("config item 'insert_message_fields' cannot be empty string")
	}
	po.insertMessageFields = strings.Split(config.InsertMessageFields, " ")
	if config.InsertTableColumns == "" {
		return fmt.Errorf("config item 'insert_table_columns' cannot be empty string")
	}
	po.insertTableColumns = strings.Split(config.InsertTableColumns, " ")
	po.allowMissingMessageFields = config.AllowMissingMessageFields
	p := postgres.DBConnectionParams{
		Host:           config.DBHost,
		Port:           config.DBPort,
		DBName:         config.DBName,
		User:           config.DBUser,
		Password:       config.DBPassword,
		ConnectTimeout: config.DBConnectionTimeout,
		SSLMode:        config.DBSSLMode,
	}
	db, err := postgres.New(&p)
	if err != nil {
		return err
	}
	db.SetMaxOpenConns(config.DBMaxOpenConnections)
	po.db = db
	return nil
}

func (o *PostgresOutput) Run(or OutputRunner, h PluginHelper) (err error) {
	defer o.db.Close()

	o.runner = or
	o.helper = h

	var wg sync.WaitGroup
	wg.Add(1)

	committers := o.makeCommitters(5, &wg)
	go o.receiver(committers, &wg)

	wg.Wait()
	return
}

// Runs in a separate goroutine, accepting incoming messages, buffering output
// data until the ticker triggers the buffered data should be put onto the
// committer channel.
func (o *PostgresOutput) receiver(committers chan<- [][]interface{}, wg *sync.WaitGroup) {
	var pack *PipelinePack

	ticker := time.Tick(time.Duration(o.flushInterval) * time.Millisecond)
	batch := [][]interface{}{}

	for ok := true; ok; {
		select {
		case pack, ok = <-o.runner.InChan():
			if !ok {
				// Closed inChan => we're shutting down, flush data
				committers <- batch
				close(committers)
				break
			}

			// Read values from message fields
			vals, err := o.convertMessageToValues(pack.Message, o.insertMessageFields)
			pack.Recycle()

			if err != nil {
				o.runner.LogError(err)
			} else {
				batch = append(batch, vals)
				if len(batch) >= o.flushCount {
					committers <- batch
					batch = [][]interface{}{}
				}
			}
		case <-ticker:
			committers <- batch
			batch = [][]interface{}{}
		}
	}
	wg.Done()
}

// Runs in a separate goroutine, waits for buffered data on the committer
// channel, bulk inserts it into Postgres, and puts the now empty buffer on the
// return channel for reuse.
func (o *PostgresOutput) makeCommitters(count int, wg *sync.WaitGroup) chan<- [][]interface{} {
	batches := make(chan [][]interface{})

	for i := 0; i < count; i++ {
		wg.Add(1)
		go func(i int) {
			for batch := range batches {
				if len(batch) <= 0 {
					continue
				}

				done := o.commit(batch)
				timeout := time.NewTimer(time.Minute)

				select {
				case <-done:
					timeout.Stop()
				case <-timeout.C:
					o.runner.LogError(errors.New("Postgres insert took more than 60s."))
				}
			}
			wg.Done()
		}(i)
	}

	return batches
}

func (o *PostgresOutput) commit(batch [][]interface{}) <-chan struct{} {
	done := make(chan struct{})

	go func() {
		err := o.db.Insert(o.insertSchema, o.insertTable, o.insertTableColumns, batch)
		if err != nil {
			o.runner.LogError(err)
		}
		done <- struct{}{}
	}()

	return done
}

// convertMessageToValue reads a Heka Message and returns a slice of field values
func (po *PostgresOutput) convertMessageToValues(m *message.Message, insertFields []string) (fieldValues []interface{}, err error) {
	fieldValues = []interface{}{}
	missingFields := []string{}
	for _, field := range insertFields {
		// Special case: get "Timestamp" from Heka message
		if field == "Timestamp" {
			// Convert Heka time (Unix timestamp in nanoseconds) to Golang time
			v := time.Unix(0, m.GetTimestamp())
			fieldValues = append(fieldValues, v)
		} else {
			v, ok := m.GetFieldValue(field)
			if !ok {
				// If configured to do so, write NULL when a FieldValue isn't found in the Heka message
				if po.allowMissingMessageFields {
					v = nil
				} else {
					missingFields = append(missingFields, field)
					continue
				}
			}
			fieldValues = append(fieldValues, v)
		}
	}

	if len(missingFields) > 0 {
		return []interface{}{}, fmt.Errorf("message is missing expected fields: %s", strings.Join(missingFields, ", "))
	}

	return fieldValues, nil
}

func init() {
	RegisterPlugin("PostgresOutput", func() interface{} {
		return new(PostgresOutput)
	})
}
<|endoftext|>


<commit_before>package fsd

import (
	"errors"
	"fmt"
	"math/rand"
	"net"
	"time"
)

var (
	Instance *Fsd
)

type Fsd struct {
	outgoing chan string
	address  string
	conn     net.Conn
}

func init() {
	Start("127.0.0.1:8125")
}

func Start(address string) {
	Instance = &Fsd{address: address, outgoing: make(chan string, 100000)}
	Instance.connect()

	go Instance.processOutgoing()
}

func (fsd *Fsd) connect() error {
	conn, err := net.Dial("udp", fsd.address)
	if err != nil {
		return err
	}

	fsd.conn = conn
	return nil
}

func (fsd *Fsd) processOutgoing() {
	for outgoing := range fsd.outgoing {
		data := fmt.Sprintf("%s", outgoing)

		if _, err := fsd.conn.Write([]byte(data)); err != nil {
			Count("fsd.sending.failure", 1)
			fsd.connect()
		} else {
			CountL("fsd.sending.sucess", 1, 0.1)
		}
	}
}

// To read about the different semantics check out
// https://github.com/b/statsd_spec
// http://docs.datadoghq.com/guides/dogstatsd/

// Increment the page.views counter.
// page.views:1|c
func Count(name string, value float64) {
	CountL(name, value, 1.0)
}

func CountL(name string, value float64, rate float64) {
	payload := createPayload(name, value) + "|c"

	suffix, err := rateCheck(rate)
	if err != nil {
		return
	}

	payload = payload + suffix
	send(payload)
}

// Record the fuel tank is half-empty
// fuel.level:0.5|g
func Gauge(name string, value float64) {
	payload := createPayload(name, value) + "|g"
	send(payload)
}

// A request latency
// request.latency:320|ms
// Or a payload of a image
// image.size:2.3|ms
func Timer(name string, milliseconds float64) {
	TimerL(name, milliseconds, 1.0)
}

func TimerL(name string, milliseconds float64, rate float64) {
	payload := createPayload(name, milliseconds) + "|ms"

	suffix, err := rateCheck(rate)
	if err != nil {
		return
	}

	payload = payload + suffix
	send(payload)
}

func Time(name string, lambda func()) {
	TimeL(name, 1.0, lambda)
}

func TimeL(name string, rate float64, lambda func()) {
	start := time.Now()
	lambda()
	TimerL(name, float64(time.Now().Sub(start).Nanoseconds()/1000000), rate)
}

// Track a unique visitor id to the site.
// users.uniques:1234|s
func Set(name string, value float64) {
	payload := createPayload(name, value) + "|s"
	send(payload)
}

func createPayload(name string, value float64) (payload string) {
	payload = fmt.Sprintf("%s:%f", name, value)
	return payload
}

func rateCheck(rate float64) (suffix string, err error) {
	if rate < 1 {
		if rand.Float64() < rate {
			return fmt.Sprintf("|@%f", rate), nil
		}
	} else {
		return "", nil
	}

	return "", errors.New("Out of rate limit")
}

func send(payload string) {
	length := float64(len(Instance.outgoing))
	capacity := float64(cap(Instance.outgoing))

	Gauge("fsd.buffer.fillrate", capacity/length)

	if length < capacity*0.9 {
		Instance.outgoing <- payload
		CountL("fsd.buffer.success", 1, 0.1)
	} else {
		CountL("fsd.buffer.failure", 1, 0.1)
	}
}
<commit_msg>Added TimeSince functions<commit_after>package fsd

import (
	"errors"
	"fmt"
	"math/rand"
	"net"
	"time"
)

var (
	Instance *Fsd
)

type Fsd struct {
	outgoing chan string
	address  string
	conn     net.Conn
}

func init() {
	Start("127.0.0.1:8125")
}

func Start(address string) {
	Instance = &Fsd{address: address, outgoing: make(chan string, 100000)}
	Instance.connect()

	go Instance.processOutgoing()
}

func (fsd *Fsd) connect() error {
	conn, err := net.Dial("udp", fsd.address)
	if err != nil {
		return err
	}

	fsd.conn = conn
	return nil
}

func (fsd *Fsd) processOutgoing() {
	for outgoing := range fsd.outgoing {
		data := fmt.Sprintf("%s", outgoing)

		if _, err := fsd.conn.Write([]byte(data)); err != nil {
			Count("fsd.sending.failure", 1)
			fsd.connect()
		} else {
			CountL("fsd.sending.sucess", 1, 0.1)
		}
	}
}

// To read about the different semantics check out
// https://github.com/b/statsd_spec
// http://docs.datadoghq.com/guides/dogstatsd/

// Increment the page.views counter.
// page.views:1|c
func Count(name string, value float64) {
	CountL(name, value, 1.0)
}

func CountL(name string, value float64, rate float64) {
	payload := createPayload(name, value) + "|c"

	suffix, err := rateCheck(rate)
	if err != nil {
		return
	}

	payload = payload + suffix
	send(payload)
}

// Record the fuel tank is half-empty
// fuel.level:0.5|g
func Gauge(name string, value float64) {
	payload := createPayload(name, value) + "|g"
	send(payload)
}

// A request latency
// request.latency:320|ms
// Or a payload of a image
// image.size:2.3|ms
func Timer(name string, milliseconds float64) {
	TimerL(name, milliseconds, 1.0)
}

func TimerL(name string, milliseconds float64, rate float64) {
	payload := createPayload(name, milliseconds) + "|ms"

	suffix, err := rateCheck(rate)
	if err != nil {
		return
	}

	payload = payload + suffix
	send(payload)
}

// TimeSince records a named timer with the duration since start
func TimeSince(name string, start time.Time) {
	TimeSinceL(name, start, 1.0)
}

// TimeSince records a rated and named timer with the duration since start
func TimeSinceL(name string, start time.Time, rate float64) {
	TimerL(name, rate, float64(time.Now().Sub(start).Nanoseconds()/1000000))
}

func Time(name string, lambda func()) {
	TimeL(name, 1.0, lambda)
}

func TimeL(name string, rate float64, lambda func()) {
	start := time.Now()
	lambda()
	TimerL(name, float64(time.Now().Sub(start).Nanoseconds()/1000000), rate)
}

// Track a unique visitor id to the site.
// users.uniques:1234|s
func Set(name string, value float64) {
	payload := createPayload(name, value) + "|s"
	send(payload)
}

func createPayload(name string, value float64) (payload string) {
	payload = fmt.Sprintf("%s:%f", name, value)
	return payload
}

func rateCheck(rate float64) (suffix string, err error) {
	if rate < 1 {
		if rand.Float64() < rate {
			return fmt.Sprintf("|@%f", rate), nil
		}
	} else {
		return "", nil
	}

	return "", errors.New("Out of rate limit")
}

func send(payload string) {
	length := float64(len(Instance.outgoing))
	capacity := float64(cap(Instance.outgoing))

	Gauge("fsd.buffer.fillrate", capacity/length)

	if length < capacity*0.9 {
		Instance.outgoing <- payload
		CountL("fsd.buffer.success", 1, 0.1)
	} else {
		CountL("fsd.buffer.failure", 1, 0.1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"html"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/xrash/smetrics"
)

const baseURL = "https://www.thetvdb.com"
const posterBaseURL = "https://www.thetvdb.com/banners/posters/"

func chop(str string, start string, end string) string {
	indexA := strings.Index(str, start)
	if indexA < 0 {
		return ""
	}
	remainder := str[indexA+len(start):]
	indexB := strings.Index(remainder, end)
	if indexB < 0 {
		return ""
	}
	return remainder[:indexB]
}

/// Finds the last occurrence of 'start' and return the text after it.
func chopLast(str string, start string) string {
	index := strings.LastIndex(str, start)
	if index < 0 {
		return ""
	}
	return str[index+len(start):]
}

/// Finds the first occurrence of 'start' and return the text after it.
func chopFirst(str string, start string) string {
	index := strings.Index(str, start)
	if index < 0 {
		return ""
	}
	return str[index+len(start):]
}

func searchUrl(name string) string {
	u, _ := url.Parse(baseURL + "/search") // Assume it'll work.
	values := url.Values{}
	values.Add("q", name)
	u.RawQuery = values.Encode()
	return u.String()
}

func get(url string) string {
	resp, respErr := http.Get(url)
	if respErr != nil {
		return ""
	}

	defer resp.Body.Close()
	body, bodyErr := ioutil.ReadAll(resp.Body)
	if bodyErr != nil {
		return ""
	}

	return string(body)
}

/// Returns the show url slug / id. "" if can't find. eg 'i-dream-of-jeannie'.
func tvdbSearchForSeries(name string) string {
	resp := get(searchUrl(name))
	results := chop(resp, "
Search Results
", "
")
	regex := regexp.MustCompile(`<a href="/series/(.*?)">(.*?)</a>`)
	matches := regex.FindAllStringSubmatch(results, -1)
	closestDistance := 99999
	var closestSlug string
	for _, match := range matches {
		if len(match) >= 3 {
			slug := match[1]
			name := match[2]

			thisDistance := smetrics.WagnerFischer(
				strings.ToLower(name),
				strings.ToLower(name),
				1, 3, 2)

			if thisDistance < closestDistance {
				closestDistance = thisDistance
				closestSlug = slug
			}
		}
	}
	return closestSlug
}

func unescapeTrim(str string) string {
	return strings.TrimSpace(html.UnescapeString(str))
}

// Some of these fields below have json names for compatability with old TMDB metadata.

type TVDBSeries struct {
	TVDBID       int // Eg 55271 - thetvdb id
	Name         string
	Overview     string
	Art          string // The 16:9 background full url.
	Poster       string // The portrait dvd cover full url.
	FirstAirDate string `json:"first_air_date"`
	Seasons      []TVDBSeason
}

type TVDBSeason struct {
	TVDBID int    // Eg 1 - the season number (old tvdb used to make this different from season number)
	Season int    `json:"season_number"` // Eg 1,2,3, or 0 for specials
	Name   string // Eg 'Specials' or 'Season 1'
	/// Following are only filled in when requesting season details.
	Episodes []TVDBEpisode
	Image    string // Full url.
}

/// Finds all the seasons in an html for a tv series details page.
func seasonsForResponse(resp string) []TVDBSeason {
	section := chop(resp, "
Seasons
", "
")
	regex := regexp.MustCompile(`(?s)seasons/([0-9]+)">(.*?)</a>`) // s flag means . can span newlines.
	seasons := make([]TVDBSeason, 0)
	matches := regex.FindAllStringSubmatch(section, -1)
	for _, match := range matches {
		if len(match) >= 3 {
			number, _ := strconv.Atoi(match[1])
			name := unescapeTrim(match[2])

			seasons = append(seasons, TVDBSeason{
				TVDBID: number,
				Season: number,
				Name:   name,
			})
		}
	}
	return seasons
}

func tvdbSeriesDetails(id string) (TVDBSeries, error) {
	seriesUrl := baseURL + "/series/" + id
	resp := get(seriesUrl)
	// Name and details.
	nameRaw := chop(resp, "<title>", " @ TheTVDB</title>")
	name := unescapeTrim(nameRaw)
	detailsArea := chop(resp, `<div class="row">`, "
")
	detailsRaw := chop(detailsArea, `<p>`, `</p>`)
	details := unescapeTrim(detailsRaw)
	// The 16:9 background.
	artArea := chop(resp, "
Backgrounds (Fan Art)
", "
")
	artURL := chop(artArea, `src="`, `"`)
	// The portrait dvd cover.
	postersArea := chop(resp, "
Posters
", "
")
	posterURL := chop(postersArea, `src="`, `"`)
	firstAiredArea := chop(resp, `First Aired`, `</li>`)
	firstAirDate := chop(firstAiredArea, `<span>`, `</span>`)
	// The seasons.
	seasons := seasonsForResponse(resp)
	if name == "" || len(seasons) == 0 {
		return TVDBSeries{}, errors.New("Could not find series")
	}
	series := TVDBSeries{
		TVDBID:       id,
		Name:         name,
		Overview:     details,
		Art:          artURL,
		Poster:       posterURL,
		FirstAirDate: firstAirDate,
		Seasons:      seasons,
	}
	return series, nil
}

type TVDBEpisode struct {
	TVDBID       int
	SeasonNumber int
	Episode      int `json:"episode_number"`
	Name         string
	AirDate      string `json:"air_date"`
	// Following are empty until you ask for episode details.
	Overview string
	Image    string // full url.
}

func seasonUrl(seriesid int, seasonid int) string {
	u, _ := url.Parse(baseURL) // Assume it'll work.
	values := url.Values{}
	values.Add("seriesid", strconv.Itoa(seriesid))
	values.Add("seasonid", strconv.Itoa(seasonid))
	values.Add("tab", "season")
	u.RawQuery = values.Encode()
	return u.String()
}

/// Extracts the episodes from the season screen.
func episodesForSeasonDetails(resp string, seasonNumber int) []TVDBEpisode {
	section := chop(resp, `<td class="head">Episode Number</td>`, `</table>`)
	regex := regexp.MustCompile(`&id=([0-9]+).+?>([0-9]+)<.+?;lid=7">(.*)</a></td><td class=".+?">(.*?)</td>`)
	// regex := regexp.MustCompile(`&id=([0-9]+).+?>([0-9]+)<.+?;lid=7">(.+)</a></td><td class=".+?">(.+?)</td>`)
	episodes := make([]TVDBEpisode, 0)
	matches := regex.FindAllStringSubmatch(section, -1)
	for _, match := range matches {
		if len(match) >= 5 {
			id, _ := strconv.Atoi(match[1])
			number, _ := strconv.Atoi(match[2])
			name := unescapeTrim(match[3])
			date := unescapeTrim(match[4])

			if name != "" && number > 0 && id > 0 {
				episodes = append(episodes, TVDBEpisode{
					TVDBID:       id,
					SeasonNumber: seasonNumber,
					Episode:      number,
					Name:         name,
					AirDate:      date,
				})
			}
		}
	}
	return episodes
}

/// Get the episodes in a season.
func tvdbSeasonDetails(seriesid int, seasonId int, seasonNumber int) (TVDBSeason, error) {
	resp := get(seasonUrl(seriesid, seasonId))
	titleSection := chop(resp, `<div class="titlesection">`, `</table>`)
	titleRaw := chop(titleSection, `<h2>`, `<`)
	title := unescapeTrim(titleRaw) // Eg 'Season 2'
	bannersSection := chop(resp, `<h1>Season Banners</h1>`, `</div>`)
	banner := chop(bannersSection, `<td align=right><a href="`, `"`) // eg banners/seasons/90601-2-2.jpg
	episodes := episodesForSeasonDetails(resp, seasonNumber)
	if title == "" || len(episodes) == 0 {
		return TVDBSeason{}, errors.New("Could not find season")
	}
	season := TVDBSeason{
		TVDBID:   seasonId,
		Season:   seasonNumber,
		Name:     title,
		Image:    baseURL + `/` + banner,
		Episodes: episodes,
	}
	return season, nil
}

func episodeUrl(seriesid int, seasonid int, episodeid int) string {
	u, _ := url.Parse(baseURL) // Assume it'll work.
	values := url.Values{}
	values.Add("seriesid", strconv.Itoa(seriesid))
	values.Add("seasonid", strconv.Itoa(seasonid))
	values.Add("id", strconv.Itoa(episodeid))
	values.Add("tab", "episode")
	u.RawQuery = values.Encode()
	return u.String()
}

func tvdbEpisodeDetails(seriesid int, seasonId int, seasonNumber int, episodeid int) (TVDBEpisode, error) {
	resp := get(episodeUrl(seriesid, seasonId, episodeid))
	episodeNumberRaw := chop(resp, `<input type="text" name="EpisodeNumber" value="`, `"`)
	episodeNumber, _ := strconv.Atoi(episodeNumberRaw)
	titleSection := chop(resp, `<div class="titlesection">`, `</div>`)
	titleH3 := chop(titleSection, `<h3>`, `</h3>`)
	title := unescapeTrim(chopFirst(titleH3, `: `))
	airDate := chop(resp, `<input type="text" name="FirstAired" value="`, `"`)
	overview := unescapeTrim(chop(resp, `name="Overview_7" style="display: inline">`, `</textarea>`))
	images := chop(resp, `<h1>Episode Image</h1>`, `</div>`)
	image := chop(images, `<a href="`, `">`) // eg 'banners/episodes/90601/3038381.jpg'

	if title == "" || episodeNumber <= 0 {
		return TVDBEpisode{}, errors.New("Could not find episode")
	}

	episode := TVDBEpisode{
		TVDBID:       episodeid,
		SeasonNumber: seasonNumber,
		Episode:      episodeNumber,
		Name:         title,
		AirDate:      airDate,
		// Following are empty until you ask for episode details.
		Overview: overview,
		Image:    baseURL + `/` + image,
	}
	return episode, nil
}
<commit_msg>episodesForSeasonDetails<commit_after>package main

import (
	"errors"
	"html"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/xrash/smetrics"
)

const baseURL = "https://www.thetvdb.com"
const posterBaseURL = "https://www.thetvdb.com/banners/posters/"

func chop(str string, start string, end string) string {
	indexA := strings.Index(str, start)
	if indexA < 0 {
		return ""
	}
	remainder := str[indexA+len(start):]
	indexB := strings.Index(remainder, end)
	if indexB < 0 {
		return ""
	}
	return remainder[:indexB]
}

/// Finds the last occurrence of 'start' and return the text after it.
func chopLast(str string, start string) string {
	index := strings.LastIndex(str, start)
	if index < 0 {
		return ""
	}
	return str[index+len(start):]
}

/// Finds the first occurrence of 'start' and return the text after it.
func chopFirst(str string, start string) string {
	index := strings.Index(str, start)
	if index < 0 {
		return ""
	}
	return str[index+len(start):]
}

func searchUrl(name string) string {
	u, _ := url.Parse(baseURL + "/search") // Assume it'll work.
	values := url.Values{}
	values.Add("q", name)
	u.RawQuery = values.Encode()
	return u.String()
}

func get(url string) string {
	resp, respErr := http.Get(url)
	if respErr != nil {
		return ""
	}

	defer resp.Body.Close()
	body, bodyErr := ioutil.ReadAll(resp.Body)
	if bodyErr != nil {
		return ""
	}

	return string(body)
}

/// Returns the show url slug / id. "" if can't find. eg 'i-dream-of-jeannie'.
func tvdbSearchForSeries(name string) string {
	resp := get(searchUrl(name))
	results := chop(resp, "
Search Results
", "
")
	regex := regexp.MustCompile(`<a href="/series/(.*?)">(.*?)</a>`)
	matches := regex.FindAllStringSubmatch(results, -1)
	closestDistance := 99999
	var closestSlug string
	for _, match := range matches {
		if len(match) >= 3 {
			slug := match[1]
			name := match[2]

			thisDistance := smetrics.WagnerFischer(
				strings.ToLower(name),
				strings.ToLower(name),
				1, 3, 2)

			if thisDistance < closestDistance {
				closestDistance = thisDistance
				closestSlug = slug
			}
		}
	}
	return closestSlug
}

func unescapeTrim(str string) string {
	return strings.TrimSpace(html.UnescapeString(str))
}

// Some of these fields below have json names for compatability with old TMDB metadata.

type TVDBSeries struct {
	TVDBID       string // Eg 'i-dream-of-jeannie' - thetvdb url slug / id
	Name         string
	Overview     string
	Art          string // The 16:9 background full url.
	Poster       string // The portrait dvd cover full url.
	FirstAirDate string `json:"first_air_date"`
	Seasons      []TVDBSeason
}

type TVDBSeason struct {
	TVDBID int    // Eg 1 - the season number (old tvdb used to make this different from season number)
	Season int    `json:"season_number"` // Eg 1,2,3, or 0 for specials
	Name   string // Eg 'Specials' or 'Season 1'
	/// Following are only filled in when requesting season details.
	Episodes []TVDBEpisode
	Image    string // Full url.
}

/// Finds all the seasons in an html for a tv series details page.
func seasonsForResponse(resp string) []TVDBSeason {
	section := chop(resp, "
Seasons
", "
")
	regex := regexp.MustCompile(`(?s)seasons/([0-9]+)">(.*?)</a>`) // s flag means . can span newlines.
	seasons := make([]TVDBSeason, 0)
	matches := regex.FindAllStringSubmatch(section, -1)
	for _, match := range matches {
		if len(match) >= 3 {
			number, _ := strconv.Atoi(match[1])
			name := unescapeTrim(match[2])

			seasons = append(seasons, TVDBSeason{
				TVDBID: number,
				Season: number,
				Name:   name,
			})
		}
	}
	return seasons
}

func tvdbSeriesDetails(id string) (TVDBSeries, error) {
	seriesUrl := baseURL + "/series/" + id
	resp := get(seriesUrl)
	// Name and details.
	nameRaw := chop(resp, "<title>", " @ TheTVDB</title>")
	name := unescapeTrim(nameRaw)
	detailsArea := chop(resp, `<div class="row">`, "
")
	detailsRaw := chop(detailsArea, `<p>`, `</p>`)
	details := unescapeTrim(detailsRaw)
	// The 16:9 background.
	artArea := chop(resp, "
Backgrounds (Fan Art)
", "
")
	artURL := chop(artArea, `src="`, `"`)
	// The portrait dvd cover.
	postersArea := chop(resp, "
Posters
", "
")
	posterURL := chop(postersArea, `src="`, `"`)
	firstAiredArea := chop(resp, `First Aired`, `</li>`)
	firstAirDate := chop(firstAiredArea, `<span>`, `</span>`)
	// The seasons.
	seasons := seasonsForResponse(resp)
	if name == "" || len(seasons) == 0 {
		return TVDBSeries{}, errors.New("Could not find series")
	}
	series := TVDBSeries{
		TVDBID:       id,
		Name:         name,
		Overview:     details,
		Art:          artURL,
		Poster:       posterURL,
		FirstAirDate: firstAirDate,
		Seasons:      seasons,
	}
	return series, nil
}

type TVDBEpisode struct {
	TVDBID       int
	SeasonNumber int
	Episode      int `json:"episode_number"`
	Name         string
	AirDate      string `json:"air_date"`
	Image        string // full url.
	// Following are empty until you ask for episode details.
	Overview string
}

/// Extracts the episodes from the season screen.
func episodesForSeasonDetails(resp string, seasonNumber int) []TVDBEpisode {
	section := chop(resp, `<div class="col-xs-12 col-sm-8 episodes">`, `</table>`)
	regex := regexp.MustCompile(`(?s)<tr>.*?/episodes/([0-9]+)">.*?([0-9]+).*?<span.*?>(.*?)</span>.*?<td>.*?([0-9][0-9])/([0-9][0-9])/([0-9][0-9][0-9][0-9]).*?</td>.*?<td>.*?data-featherlight="(.*?)".*?</td>.*?</tr>`)
	episodes := make([]TVDBEpisode, 0)
	matches := regex.FindAllStringSubmatch(section, -1)
	for _, match := range matches {
		if len(match) >= 8 {
			id, _ := strconv.Atoi(match[1])
			number, _ := strconv.Atoi(match[2])
			name := unescapeTrim(match[3])
			mm := match[4]   // mm
			dd := match[5]   // dd
			yyyy := match[6] // yyyy
			date := yyyy + "-" + mm + "-" + dd
			image := match[7]

			if name != "" && number > 0 && id > 0 {
				episodes = append(episodes, TVDBEpisode{
					TVDBID:       id,
					SeasonNumber: seasonNumber,
					Episode:      number,
					Name:         name,
					AirDate:      date,
					Image:        image,
				})
			}
		}
	}
	return episodes
}

/// Get the episodes in a season.
func tvdbSeasonDetails(seriesid string, seasonId int, seasonNumber int) (TVDBSeason, error) {
	seasonUrl := baseURL + "/series/" + seriesid + "/seasons/" + strconv.Itoa(seasonId)
	resp := get(seasonUrl)

	titleSection := chop(resp, `<title>`, `</title>`)
	titleRaw := chop(titleSection, ` - `, ` @ TheTVDB`)
	title := unescapeTrim(titleRaw) // Eg 'Season 2'

	postersSection := chop(resp, `<h2>Posters</h2>`, `</div>`)
	posterURL := chop(postersSection, `src="`, `"`)

	episodes := episodesForSeasonDetails(resp, seasonNumber)

	if title == "" || len(episodes) == 0 {
		return TVDBSeason{}, errors.New("Could not find season")
	}

	season := TVDBSeason{
		TVDBID:   seasonId,
		Season:   seasonNumber,
		Name:     title,
		Image:    posterURL,
		Episodes: episodes,
	}
	return season, nil
}

func episodeUrl(seriesid int, seasonid int, episodeid int) string {
	u, _ := url.Parse(baseURL) // Assume it'll work.
	values := url.Values{}
	values.Add("seriesid", strconv.Itoa(seriesid))
	values.Add("seasonid", strconv.Itoa(seasonid))
	values.Add("id", strconv.Itoa(episodeid))
	values.Add("tab", "episode")
	u.RawQuery = values.Encode()
	return u.String()
}

func tvdbEpisodeDetails(seriesid int, seasonId int, seasonNumber int, episodeid int) (TVDBEpisode, error) {
	resp := get(episodeUrl(seriesid, seasonId, episodeid))
	episodeNumberRaw := chop(resp, `<input type="text" name="EpisodeNumber" value="`, `"`)
	episodeNumber, _ := strconv.Atoi(episodeNumberRaw)
	titleSection := chop(resp, `<div class="titlesection">`, `</div>`)
	titleH3 := chop(titleSection, `<h3>`, `</h3>`)
	title := unescapeTrim(chopFirst(titleH3, `: `))
	airDate := chop(resp, `<input type="text" name="FirstAired" value="`, `"`)
	overview := unescapeTrim(chop(resp, `name="Overview_7" style="display: inline">`, `</textarea>`))
	images := chop(resp, `<h1>Episode Image</h1>`, `</div>`)
	image := chop(images, `<a href="`, `">`) // eg 'banners/episodes/90601/3038381.jpg'

	if title == "" || episodeNumber <= 0 {
		return TVDBEpisode{}, errors.New("Could not find episode")
	}

	episode := TVDBEpisode{
		TVDBID:       episodeid,
		SeasonNumber: seasonNumber,
		Episode:      episodeNumber,
		Name:         title,
		AirDate:      airDate,
		// Following are empty until you ask for episode details.
		Overview: overview,
		Image:    baseURL + `/` + image,
	}
	return episode, nil
}
<|endoftext|>