<commit_before>// Copyright 2016 The LUCI Authors.
//
// 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 templateproto

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
	. "go.chromium.org/luci/common/testing/assertions"
)

func TestLoadFromConfig(t *testing.T) {
	t.Parallel()

	Convey("LoadFile", t, func() {
		templateContent := `
		template: <
			key: "hardcode"
			value: <
				doc: "it's hard-coded"
				body: <<EOF
					{"woot": ["sauce"]}
				EOF
			>
		>

		template: <
			key: "templ_1"
			value: <
				doc: << EOF
					This template is templ_1!
					It's pretty "exciting"!
				EOF
				body: << EOF
					{
						"json_key": ${json_key},
						"cmd": ["array", "of", ${thing}],
						"extra": ${extra}
					}
				EOF
				param: <
					key: "${json_key}"
					value: <
						doc: "it's a json key"
						schema: <int:<>>
					>
				>
				param: <
					key: "${thing}"
					value: <
						doc: << EOF
							${thing} represents a color or a fruit
						EOF
						schema: <enum:<
							entry: < doc: "fruit" token: "banana" >
							entry: < doc: "color" token: "white" >
							entry: < doc: "color" token: "purple" >
						>>
					>
				>
				param: <
					key: "${extra}"
					value: <
						nullable: true
						schema:<object:<>>
						default: <object: <<EOF
							{"yes": "please"}
						EOF
						>
					>
				>
			>
		>
		`

		Convey("basic load", func() {
			file, err := LoadFile(templateContent)
			So(err, ShouldBeNil)
			So(file, ShouldResemble, &File{Template: map[string]*File_Template{
				"hardcode": {
					Doc:  "it's hard-coded",
					Body: `{"woot": ["sauce"]}`,
				},

				"templ_1": {
					Doc:  "This template is templ_1!\nIt's pretty \"exciting\"!",
					Body: "{\n\t\"json_key\": ${json_key},\n\t\"cmd\": [\"array\", \"of\", ${thing}],\n\t\"extra\": ${extra}\n}",
					Param: map[string]*File_Template_Parameter{
						"${json_key}": {
							Doc:    "it's a json key",
							Schema: &Schema{Schema: &Schema_Int{&Schema_Atom{}}},
						},
						"${thing}": {
							Doc: "${thing} represents a color or a fruit",
							Schema: &Schema{
								Schema: &Schema_Enum{&Schema_Set{Entry: []*Schema_Set_Entry{
									{Doc: "fruit", Token: "banana"},
									{Doc: "color", Token: "white"},
									{Doc: "color", Token: "purple"},
								}}},
							},
						},
						"${extra}": {
							Nullable: true,
							Schema:   &Schema{Schema: &Schema_Object{&Schema_JSON{}}},
							Default:  MustNewValue(map[string]interface{}{"yes": "please"}),
						},
					},
				},
			}})

			Convey("basic render", func() {
				ret, err := file.RenderL("templ_1", LiteralMap{"${thing}": "white", "${json_key}": 20})
				So(err, ShouldBeNil)
				So(ret, ShouldEqual, `{
	"json_key": 20,
	"cmd": ["array", "of", "white"],
	"extra": {"yes":"please"}
}`)
			})

			Convey("null override", func() {
				ret, err := file.RenderL("templ_1", LiteralMap{"${thing}": "white", "${json_key}": 20, "${extra}": nil})
				So(err, ShouldBeNil)
				So(ret, ShouldEqual, `{
	"json_key": 20,
	"cmd": ["array", "of", "white"],
	"extra": null
}`)
			})

			Convey("bad render gets context", func() {
				_, err := file.RenderL("templ_1", LiteralMap{"${thing}": 10, "${json_key}": 20, "${extra}": nil})
				So(err, ShouldErrLike, "rendering \"templ_1\": param \"${thing}\": type is \"int\", expected \"enum\"")
			})

			Convey("hardcode", func() {
				ret, err := file.RenderL("hardcode", nil)
				So(err, ShouldBeNil)
				So(ret, ShouldEqual, `{"woot": ["sauce"]}`)
			})
		})

	})
}
<commit_msg>templateproto: use ShouldResembleProto<commit_after>// Copyright 2016 The LUCI Authors.
//
// 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 templateproto

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
	. "go.chromium.org/luci/common/testing/assertions"
)

func TestLoadFromConfig(t *testing.T) {
	t.Parallel()

	Convey("LoadFile", t, func() {
		templateContent := `
		template: <
			key: "hardcode"
			value: <
				doc: "it's hard-coded"
				body: <<EOF
					{"woot": ["sauce"]}
				EOF
			>
		>

		template: <
			key: "templ_1"
			value: <
				doc: << EOF
					This template is templ_1!
					It's pretty "exciting"!
				EOF
				body: << EOF
					{
						"json_key": ${json_key},
						"cmd": ["array", "of", ${thing}],
						"extra": ${extra}
					}
				EOF
				param: <
					key: "${json_key}"
					value: <
						doc: "it's a json key"
						schema: <int:<>>
					>
				>
				param: <
					key: "${thing}"
					value: <
						doc: << EOF
							${thing} represents a color or a fruit
						EOF
						schema: <enum:<
							entry: < doc: "fruit" token: "banana" >
							entry: < doc: "color" token: "white" >
							entry: < doc: "color" token: "purple" >
						>>
					>
				>
				param: <
					key: "${extra}"
					value: <
						nullable: true
						schema:<object:<>>
						default: <object: <<EOF
							{"yes": "please"}
						EOF
						>
					>
				>
			>
		>
		`

		Convey("basic load", func() {
			file, err := LoadFile(templateContent)
			So(err, ShouldBeNil)
			So(file, ShouldResembleProto, &File{Template: map[string]*File_Template{
				"hardcode": {
					Doc:  "it's hard-coded",
					Body: `{"woot": ["sauce"]}`,
				},

				"templ_1": {
					Doc:  "This template is templ_1!\nIt's pretty \"exciting\"!",
					Body: "{\n\t\"json_key\": ${json_key},\n\t\"cmd\": [\"array\", \"of\", ${thing}],\n\t\"extra\": ${extra}\n}",
					Param: map[string]*File_Template_Parameter{
						"${json_key}": {
							Doc:    "it's a json key",
							Schema: &Schema{Schema: &Schema_Int{&Schema_Atom{}}},
						},
						"${thing}": {
							Doc: "${thing} represents a color or a fruit",
							Schema: &Schema{
								Schema: &Schema_Enum{&Schema_Set{Entry: []*Schema_Set_Entry{
									{Doc: "fruit", Token: "banana"},
									{Doc: "color", Token: "white"},
									{Doc: "color", Token: "purple"},
								}}},
							},
						},
						"${extra}": {
							Nullable: true,
							Schema:   &Schema{Schema: &Schema_Object{&Schema_JSON{}}},
							Default:  MustNewValue(map[string]interface{}{"yes": "please"}),
						},
					},
				},
			}})

			Convey("basic render", func() {
				ret, err := file.RenderL("templ_1", LiteralMap{"${thing}": "white", "${json_key}": 20})
				So(err, ShouldBeNil)
				So(ret, ShouldEqual, `{
	"json_key": 20,
	"cmd": ["array", "of", "white"],
	"extra": {"yes":"please"}
}`)
			})

			Convey("null override", func() {
				ret, err := file.RenderL("templ_1", LiteralMap{"${thing}": "white", "${json_key}": 20, "${extra}": nil})
				So(err, ShouldBeNil)
				So(ret, ShouldEqual, `{
	"json_key": 20,
	"cmd": ["array", "of", "white"],
	"extra": null
}`)
			})

			Convey("bad render gets context", func() {
				_, err := file.RenderL("templ_1", LiteralMap{"${thing}": 10, "${json_key}": 20, "${extra}": nil})
				So(err, ShouldErrLike, "rendering \"templ_1\": param \"${thing}\": type is \"int\", expected \"enum\"")
			})

			Convey("hardcode", func() {
				ret, err := file.RenderL("hardcode", nil)
				So(err, ShouldBeNil)
				So(ret, ShouldEqual, `{"woot": ["sauce"]}`)
			})
		})

	})
}
<|endoftext|>


<commit_before>package swarm

import (
	"errors"
	"fmt"
	"sort"
	"sync"

	"github.com/docker/cli/cli/command/formatter"
	"github.com/docker/cli/cli/command/service"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/swarm"
	"github.com/moncho/dry/appui"
	"github.com/moncho/dry/docker"
	"github.com/moncho/dry/ui"
	"github.com/moncho/dry/ui/termui"

	gizaktermui "github.com/gizak/termui"
)

var defaultServiceTableHeader = serviceTableHeader()

//ServicesWidget shows information about services running on the Swarm
type ServicesWidget struct {
	swarmClient          docker.SwarmAPI
	services             []*ServiceRow
	header               *termui.TableHeader
	selectedIndex        int
	offset               int
	x, y                 int
	height, width        int
	startIndex, endIndex int
	mounted              bool
	sync.RWMutex
}

//NewServicesWidget creates a ServicesWidget
func NewServicesWidget(swarmClient docker.SwarmAPI, y int) *ServicesWidget {
	w := &ServicesWidget{
		swarmClient:   swarmClient,
		header:        defaultServiceTableHeader,
		selectedIndex: 0,
		offset:        0,
		x:             0,
		y:             y,
		height:        appui.MainScreenAvailableHeight(),
		width:         ui.ActiveScreen.Dimensions.Width}

	return w

}

//Mount prepares this widget for rendering
func (s *ServicesWidget) Mount() error {
	s.Lock()
	defer s.Unlock()
	if !s.mounted {
		s.mounted = true
		var rows []*ServiceRow
		if services, servicesInfo, err := getServiceInfo(s.swarmClient); err == nil {
			sort.SliceStable(services, func(i, j int) bool {
				return services[i].Spec.Name < services[j].Spec.Name
			})
			for _, service := range services {
				rows = append(rows, NewServiceRow(service, servicesInfo[service.ID], s.header))
			}
		}
		s.services = rows
	}
	s.align()
	return nil
}

func (s *ServicesWidget) Name() string {
	return ""
}

//Unmount marks this widget as unmounted
func (s *ServicesWidget) Unmount() error {
	s.Lock()
	defer s.Unlock()

	s.mounted = false
	return nil

}

//Align aligns rows
func (s *ServicesWidget) align() {
	x := s.x
	width := s.width

	s.header.SetWidth(width)
	s.header.SetX(x)

	for _, service := range s.services {
		service.SetX(x)
		service.SetWidth(width)
	}

}

//Buffer returns the content of this widget as a termui.Buffer
func (s *ServicesWidget) Buffer() gizaktermui.Buffer {
	s.Lock()
	defer s.Unlock()
	y := s.y
	buf := gizaktermui.NewBuffer()

	widgetHeader := appui.WidgetHeader("Service", s.RowCount(), "")
	widgetHeader.Y = y
	buf.Merge(widgetHeader.Buffer())
	y += widgetHeader.GetHeight()

	s.header.SetY(y)
	buf.Merge(s.header.Buffer())
	y += s.header.GetHeight()

	s.highlightSelectedRow()
	for _, service := range s.visibleRows() {
		service.SetY(y)
		service.Height = 1
		y += service.GetHeight()
		buf.Merge(service.Buffer())
	}

	return buf
}

//RowCount returns the number of rowns of this widget.
func (s *ServicesWidget) RowCount() int {
	return len(s.services)
}
func (s *ServicesWidget) highlightSelectedRow() {
	if s.RowCount() == 0 {
		return
	}
	index := ui.ActiveScreen.Cursor.Position()
	if index > s.RowCount() {
		index = s.RowCount() - 1
	}
	s.services[s.selectedIndex].NotHighlighted()
	s.selectedIndex = index
	s.services[s.selectedIndex].Highlighted()
}

//OnEvent runs the given command
func (s *ServicesWidget) OnEvent(event appui.EventCommand) error {
	if s.RowCount() > 0 {
		return event(s.services[s.selectedIndex].service.ID)
	}
	return errors.New("the service list is empty")
}

func (s *ServicesWidget) visibleRows() []*ServiceRow {

	//no screen
	if s.height < 0 {
		return nil
	}
	rows := s.services
	count := len(rows)
	cursor := ui.ActiveScreen.Cursor
	selected := cursor.Position()
	//everything fits
	if count <= s.height {
		return rows
	}
	//at the the start
	if selected == 0 {
		//internal state is reset
		s.startIndex = 0
		s.endIndex = s.height
		return rows[s.startIndex : s.endIndex+1]
	}

	if selected >= s.endIndex {
		if selected-s.height >= 0 {
			s.startIndex = selected - s.height
		}
		s.endIndex = selected
	}
	if selected <= s.startIndex {
		s.startIndex = s.startIndex - 1
		if selected+s.height < count {
			s.endIndex = s.startIndex + s.height
		}
	}
	start := s.startIndex
	end := s.endIndex + 1
	return rows[start:end]
}

func serviceTableHeader() *termui.TableHeader {
	fields := []string{
		"ID", "NAME", "MODE", "REPLICAS", "SERVICE PORT(S)", "IMAGE"}

	header := termui.NewHeader(appui.DryTheme)
	header.ColumnSpacing = appui.DefaultColumnSpacing
	header.AddColumn(fields[0])
	header.AddColumn(fields[1])
	header.AddFixedWidthColumn(fields[2], 12)
	header.AddFixedWidthColumn(fields[3], 10)
	header.AddColumn(fields[4])
	header.AddColumn(fields[5])

	return header
}

func getServiceInfo(swarmClient docker.SwarmAPI) ([]swarm.Service, map[string]formatter.ServiceListInfo, error) {

	serviceFilters := filters.NewArgs()
	serviceFilters.Add("runtime", string(swarm.RuntimeContainer))
	services, err := swarmClient.Services()
	if err != nil {
		return nil, nil, err
	}

	info := map[string]formatter.ServiceListInfo{}
	if len(services) > 0 {

		tasks, err := swarmClient.ServiceTasks(serviceIDs(services)...)
		if err != nil {
			return nil, nil, err
		}

		nodes, err := swarmClient.Nodes()
		if err != nil {
			return nil, nil, err
		}

		info = service.GetServicesStatus(services, nodes, tasks)
	}
	return services, info, nil
}

func serviceIDs(services []swarm.Service) []string {

	ids := make([]string, len(services))
	for i, service := range services {
		ids[i] = service.ID
	}

	return ids
}

// getServicesStatus returns a map of mode and replicas
func getServicesStatus(services []swarm.Service, nodes []swarm.Node, tasks []swarm.Task) map[string]formatter.ServiceListInfo {
	running := map[string]int{}
	tasksNoShutdown := map[string]int{}

	activeNodes := make(map[string]struct{})
	for _, n := range nodes {
		if n.Status.State != swarm.NodeStateDown {
			activeNodes[n.ID] = struct{}{}
		}
	}

	for _, task := range tasks {
		if task.DesiredState != swarm.TaskStateShutdown {
			tasksNoShutdown[task.ServiceID]++
		}

		if _, nodeActive := activeNodes[task.NodeID]; nodeActive && task.Status.State == swarm.TaskStateRunning {
			running[task.ServiceID]++
		}
	}

	info := map[string]formatter.ServiceListInfo{}
	for _, service := range services {
		info[service.ID] = formatter.ServiceListInfo{}
		if service.Spec.Mode.Replicated != nil && service.Spec.Mode.Replicated.Replicas != nil {
			info[service.ID] = formatter.ServiceListInfo{
				Mode:     "replicated",
				Replicas: fmt.Sprintf("%d/%d", running[service.ID], *service.Spec.Mode.Replicated.Replicas),
			}
		} else if service.Spec.Mode.Global != nil {
			info[service.ID] = formatter.ServiceListInfo{
				Mode:     "global",
				Replicas: fmt.Sprintf("%d/%d", running[service.ID], tasksNoShutdown[service.ID]),
			}
		}
	}
	return info
}
<commit_msg>Set the name of the services widget<commit_after>package swarm

import (
	"fmt"
	"sort"
	"sync"

	"github.com/docker/cli/cli/command/formatter"
	"github.com/docker/cli/cli/command/service"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/swarm"
	"github.com/moncho/dry/appui"
	"github.com/moncho/dry/docker"
	"github.com/moncho/dry/ui"
	"github.com/moncho/dry/ui/termui"

	gizaktermui "github.com/gizak/termui"
)

var defaultServiceTableHeader = serviceTableHeader()

//ServicesWidget shows information about services running on the Swarm
type ServicesWidget struct {
	swarmClient          docker.SwarmAPI
	services             []*ServiceRow
	header               *termui.TableHeader
	selectedIndex        int
	offset               int
	x, y                 int
	height, width        int
	startIndex, endIndex int
	mounted              bool
	sync.RWMutex
}

//NewServicesWidget creates a ServicesWidget
func NewServicesWidget(swarmClient docker.SwarmAPI, y int) *ServicesWidget {
	w := &ServicesWidget{
		swarmClient:   swarmClient,
		header:        defaultServiceTableHeader,
		selectedIndex: 0,
		offset:        0,
		x:             0,
		y:             y,
		height:        appui.MainScreenAvailableHeight(),
		width:         ui.ActiveScreen.Dimensions.Width}

	return w

}

//Mount prepares this widget for rendering
func (s *ServicesWidget) Mount() error {
	s.Lock()
	defer s.Unlock()
	if !s.mounted {
		s.mounted = true
		var rows []*ServiceRow
		if services, servicesInfo, err := getServiceInfo(s.swarmClient); err == nil {
			sort.SliceStable(services, func(i, j int) bool {
				return services[i].Spec.Name < services[j].Spec.Name
			})
			for _, service := range services {
				rows = append(rows, NewServiceRow(service, servicesInfo[service.ID], s.header))
			}
		}
		s.services = rows
	}
	s.align()
	return nil
}

//Name returns this widget name
func (s *ServicesWidget) Name() string {
	return "ServicesWidget"
}

//Unmount marks this widget as unmounted
func (s *ServicesWidget) Unmount() error {
	s.Lock()
	defer s.Unlock()

	s.mounted = false
	return nil

}

//Align aligns rows
func (s *ServicesWidget) align() {
	x := s.x
	width := s.width

	s.header.SetWidth(width)
	s.header.SetX(x)

	for _, service := range s.services {
		service.SetX(x)
		service.SetWidth(width)
	}

}

//Buffer returns the content of this widget as a termui.Buffer
func (s *ServicesWidget) Buffer() gizaktermui.Buffer {
	s.Lock()
	defer s.Unlock()
	y := s.y
	buf := gizaktermui.NewBuffer()

	widgetHeader := appui.WidgetHeader("Service", s.RowCount(), "")
	widgetHeader.Y = y
	buf.Merge(widgetHeader.Buffer())
	y += widgetHeader.GetHeight()

	s.header.SetY(y)
	buf.Merge(s.header.Buffer())
	y += s.header.GetHeight()

	s.highlightSelectedRow()
	for _, service := range s.visibleRows() {
		service.SetY(y)
		service.Height = 1
		y += service.GetHeight()
		buf.Merge(service.Buffer())
	}

	return buf
}

//RowCount returns the number of rowns of this widget.
func (s *ServicesWidget) RowCount() int {
	return len(s.services)
}
func (s *ServicesWidget) highlightSelectedRow() {
	if s.RowCount() == 0 {
		return
	}
	index := ui.ActiveScreen.Cursor.Position()
	if index > s.RowCount() {
		index = s.RowCount() - 1
	}
	s.services[s.selectedIndex].NotHighlighted()
	s.selectedIndex = index
	s.services[s.selectedIndex].Highlighted()
}

//OnEvent runs the given command
func (s *ServicesWidget) OnEvent(event appui.EventCommand) error {
	if s.RowCount() > 0 {
		return event(s.services[s.selectedIndex].service.ID)
	}
	return nil
}

func (s *ServicesWidget) visibleRows() []*ServiceRow {

	//no screen
	if s.height < 0 {
		return nil
	}
	rows := s.services
	count := len(rows)
	cursor := ui.ActiveScreen.Cursor
	selected := cursor.Position()
	//everything fits
	if count <= s.height {
		return rows
	}
	//at the the start
	if selected == 0 {
		//internal state is reset
		s.startIndex = 0
		s.endIndex = s.height
		return rows[s.startIndex : s.endIndex+1]
	}

	if selected >= s.endIndex {
		if selected-s.height >= 0 {
			s.startIndex = selected - s.height
		}
		s.endIndex = selected
	}
	if selected <= s.startIndex {
		s.startIndex = s.startIndex - 1
		if selected+s.height < count {
			s.endIndex = s.startIndex + s.height
		}
	}
	start := s.startIndex
	end := s.endIndex + 1
	return rows[start:end]
}

func serviceTableHeader() *termui.TableHeader {
	fields := []string{
		"ID", "NAME", "MODE", "REPLICAS", "SERVICE PORT(S)", "IMAGE"}

	header := termui.NewHeader(appui.DryTheme)
	header.ColumnSpacing = appui.DefaultColumnSpacing
	header.AddColumn(fields[0])
	header.AddColumn(fields[1])
	header.AddFixedWidthColumn(fields[2], 12)
	header.AddFixedWidthColumn(fields[3], 10)
	header.AddColumn(fields[4])
	header.AddColumn(fields[5])

	return header
}

func getServiceInfo(swarmClient docker.SwarmAPI) ([]swarm.Service, map[string]formatter.ServiceListInfo, error) {

	serviceFilters := filters.NewArgs()
	serviceFilters.Add("runtime", string(swarm.RuntimeContainer))
	services, err := swarmClient.Services()
	if err != nil {
		return nil, nil, err
	}

	info := map[string]formatter.ServiceListInfo{}
	if len(services) > 0 {

		tasks, err := swarmClient.ServiceTasks(serviceIDs(services)...)
		if err != nil {
			return nil, nil, err
		}

		nodes, err := swarmClient.Nodes()
		if err != nil {
			return nil, nil, err
		}

		info = service.GetServicesStatus(services, nodes, tasks)
	}
	return services, info, nil
}

func serviceIDs(services []swarm.Service) []string {

	ids := make([]string, len(services))
	for i, service := range services {
		ids[i] = service.ID
	}

	return ids
}

// getServicesStatus returns a map of mode and replicas
func getServicesStatus(services []swarm.Service, nodes []swarm.Node, tasks []swarm.Task) map[string]formatter.ServiceListInfo {
	running := map[string]int{}
	tasksNoShutdown := map[string]int{}

	activeNodes := make(map[string]struct{})
	for _, n := range nodes {
		if n.Status.State != swarm.NodeStateDown {
			activeNodes[n.ID] = struct{}{}
		}
	}

	for _, task := range tasks {
		if task.DesiredState != swarm.TaskStateShutdown {
			tasksNoShutdown[task.ServiceID]++
		}

		if _, nodeActive := activeNodes[task.NodeID]; nodeActive && task.Status.State == swarm.TaskStateRunning {
			running[task.ServiceID]++
		}
	}

	info := map[string]formatter.ServiceListInfo{}
	for _, service := range services {
		info[service.ID] = formatter.ServiceListInfo{}
		if service.Spec.Mode.Replicated != nil && service.Spec.Mode.Replicated.Replicas != nil {
			info[service.ID] = formatter.ServiceListInfo{
				Mode:     "replicated",
				Replicas: fmt.Sprintf("%d/%d", running[service.ID], *service.Spec.Mode.Replicated.Replicas),
			}
		} else if service.Spec.Mode.Global != nil {
			info[service.ID] = formatter.ServiceListInfo{
				Mode:     "global",
				Replicas: fmt.Sprintf("%d/%d", running[service.ID], tasksNoShutdown[service.ID]),
			}
		}
	}
	return info
}
<|endoftext|>


<commit_before>package ginmon

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

const TestMode string = "test"

const checkMark = "\u2713"
const ballotX = "\u2717"

const testpath = "/foo/bar"

func internalGinCtx() *gin.Context {
	return &gin.Context{
		Request: &http.Request{
			URL: &url.URL{
				Path: testpath,
			},
		},
	}
}

func Test_Inc(t *testing.T) {
	ca := NewCounterAspect()
	expect := 1
	tup := tuple{
		path: testpath,
		code: 404,
	}
	ca.increment(tup)
	ca.reset()
	if assert.Equal(t, expect, ca.RequestsSum, "Incrementation of counter does not work, expect %d but got %d %s",
		expect, ca.RequestsSum, ballotX) {
		t.Logf("Incrementation of counter works, expect %d and git %d %s",
			expect, ca.RequestsSum, checkMark)
	}
}

func Test_GetStats(t *testing.T) {
	ca := NewCounterAspect()
	if assert.NotNil(t, ca.GetStats(), "Return of Getstats() should not be nil") {
		t.Logf("Should be an interface %s", checkMark)
	}

	newCa := ca.GetStats().(CounterAspect)
	expect := 0
	if assert.Equal(t, expect, newCa.RequestsSum, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	tup := tuple{
		path: testpath,
		code: 404,
	}
	ca.increment(tup)

	if assert.Equal(t, expect, newCa.RequestsSum, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	if assert.Equal(t, expect, newCa.Requests[testpath], "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.Requests[testpath], ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.Requests[testpath], checkMark)
	}

	ca.reset()
	newCa = ca.GetStats().(CounterAspect)
	expect = 1
	if assert.Equal(t, expect, newCa.RequestsSum, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	if assert.Equal(t, expect, newCa.Requests[testpath], "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.Requests[testpath], ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.Requests[testpath], checkMark)
	}
}

func Test_Name(t *testing.T) {
	ca := NewCounterAspect()
	expect := "Counter"
	if assert.Equal(t, expect, ca.Name(), "Return of counter name does not work, expect %s but got %s %s",
		expect, ca.Name(), ballotX) {
		t.Logf("Return of counter name works, expect %s and got %s %s",
			expect, ca.Name(), checkMark)
	}
}

func Test_InRoot(t *testing.T) {
	ca := NewCounterAspect()
	expect := false
	if assert.Equal(t, expect, ca.InRoot(), "Expect %v but got %v %s",
		expect, ca.InRoot(), ballotX) {
		t.Logf("Expect %v and got %v %s",
			expect, ca.InRoot(), checkMark)
	}
}

// func Test_CounterHandler(t *testing.T) {
// 	gin.SetMode(TestMode)
// 	router := gin.New()
// 	ca := NewCounterAspect()
// 	expect := 1
// 	tup := tuple{
// 		path: testpath,
// 		code: 404,
// 	}
// 	ca.increment(tup)
// 	ca.reset()

// 	router.Use(CounterHandler(ca))
// 	tryRequest(router, "GET", "/")
// 	if assert.Equal(t, expect, ca.RequestsSum, "Incrementation of counter does not work, expect %d but got %d %s", expect, ca.RequestsSum, ballotX) {
// 		t.Logf("CounterHandler works, expect %d and got %d %s", expect, ca.RequestsSum, checkMark)
// 	}
// }

func tryRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}
<commit_msg>remove dead code<commit_after>package ginmon

import (
	"net/http"
	"net/url"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

const TestMode string = "test"

const checkMark = "\u2713"
const ballotX = "\u2717"

const testpath = "/foo/bar"

func internalGinCtx() *gin.Context {
	return &gin.Context{
		Request: &http.Request{
			URL: &url.URL{
				Path: testpath,
			},
		},
	}
}

func Test_Inc(t *testing.T) {
	ca := NewCounterAspect()
	expect := 1
	tup := tuple{
		path: testpath,
		code: 404,
	}
	ca.increment(tup)
	ca.reset()
	if assert.Equal(t, expect, ca.RequestsSum, "Incrementation of counter does not work, expect %d but got %d %s",
		expect, ca.RequestsSum, ballotX) {
		t.Logf("Incrementation of counter works, expect %d and git %d %s",
			expect, ca.RequestsSum, checkMark)
	}
}

func Test_GetStats(t *testing.T) {
	ca := NewCounterAspect()
	if assert.NotNil(t, ca.GetStats(), "Return of Getstats() should not be nil") {
		t.Logf("Should be an interface %s", checkMark)
	}

	newCa := ca.GetStats().(CounterAspect)
	expect := 0
	if assert.Equal(t, expect, newCa.RequestsSum, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	tup := tuple{
		path: testpath,
		code: 404,
	}
	ca.increment(tup)

	if assert.Equal(t, expect, newCa.RequestsSum, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	if assert.Equal(t, expect, newCa.Requests[testpath], "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.Requests[testpath], ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.Requests[testpath], checkMark)
	}

	ca.reset()
	newCa = ca.GetStats().(CounterAspect)
	expect = 1
	if assert.Equal(t, expect, newCa.RequestsSum, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	if assert.Equal(t, expect, newCa.Requests[testpath], "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.Requests[testpath], ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.Requests[testpath], checkMark)
	}
}

func Test_Name(t *testing.T) {
	ca := NewCounterAspect()
	expect := "Counter"
	if assert.Equal(t, expect, ca.Name(), "Return of counter name does not work, expect %s but got %s %s",
		expect, ca.Name(), ballotX) {
		t.Logf("Return of counter name works, expect %s and got %s %s",
			expect, ca.Name(), checkMark)
	}
}

func Test_InRoot(t *testing.T) {
	ca := NewCounterAspect()
	expect := false
	if assert.Equal(t, expect, ca.InRoot(), "Expect %v but got %v %s",
		expect, ca.InRoot(), ballotX) {
		t.Logf("Expect %v and got %v %s",
			expect, ca.InRoot(), checkMark)
	}
}
<|endoftext|>


<commit_before>package db

import (
	"context"
	"fmt"
	"time"

	"code.cloudfoundry.org/lager/lagerctx"
	sq "github.com/Masterminds/squirrel"
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/creds"
	"github.com/concourse/concourse/atc/db/lock"
)

//go:generate counterfeiter . Checkable

type Checkable interface {
	PipelineRef

	Name() string
	TeamID() int
	ResourceConfigScopeID() int
	TeamName() string
	Type() string
	Source() atc.Source
	Tags() atc.Tags
	CheckEvery() *atc.CheckEvery
	CheckTimeout() string
	LastCheckEndTime() time.Time
	CurrentPinnedVersion() atc.Version

	HasWebhook() bool

	CheckPlan(atc.Version, time.Duration, ResourceTypes, atc.Source) atc.CheckPlan
	CreateBuild(context.Context, bool, atc.Plan) (Build, bool, error)
}

//go:generate counterfeiter . CheckFactory

type CheckFactory interface {
	TryCreateCheck(context.Context, Checkable, ResourceTypes, atc.Version, bool) (Build, bool, error)
	Resources() ([]Resource, error)
	ResourceTypes() ([]ResourceType, error)
}

type checkFactory struct {
	conn        Conn
	lockFactory lock.LockFactory

	secrets       creds.Secrets
	varSourcePool creds.VarSourcePool

	planFactory atc.PlanFactory

	defaultCheckTimeout             time.Duration
	defaultCheckInterval            time.Duration
	defaultWithWebhookCheckInterval time.Duration
}

type CheckDurations struct {
	Timeout             time.Duration
	Interval            time.Duration
	IntervalWithWebhook time.Duration
}

func NewCheckFactory(
	conn Conn,
	lockFactory lock.LockFactory,
	secrets creds.Secrets,
	varSourcePool creds.VarSourcePool,
	durations CheckDurations,
) CheckFactory {
	return &checkFactory{
		conn:        conn,
		lockFactory: lockFactory,

		secrets:       secrets,
		varSourcePool: varSourcePool,

		planFactory: atc.NewPlanFactory(time.Now().Unix()),

		defaultCheckTimeout:             durations.Timeout,
		defaultCheckInterval:            durations.Interval,
		defaultWithWebhookCheckInterval: durations.IntervalWithWebhook,
	}
}

func (c *checkFactory) TryCreateCheck(ctx context.Context, checkable Checkable, resourceTypes ResourceTypes, from atc.Version, manuallyTriggered bool) (Build, bool, error) {
	logger := lagerctx.FromContext(ctx)

	var err error

	sourceDefaults := atc.Source{}
	parentType, found := resourceTypes.Parent(checkable)
	if found {
		if parentType.Version() == nil {
			return nil, false, fmt.Errorf("resource type '%s' has no version", parentType.Name())
		}
		sourceDefaults = parentType.Defaults()
	} else {
		defaults, found := atc.FindBaseResourceTypeDefaults(checkable.Type())
		if found {
			sourceDefaults = defaults
		}
	}

	interval := c.defaultCheckInterval
	if checkable.HasWebhook() {
		interval = c.defaultWithWebhookCheckInterval
	}
	if checkable.CheckEvery() != nil && !checkable.CheckEvery().Never {
		interval = checkable.CheckEvery().Interval
	}

	if !manuallyTriggered && time.Now().Before(checkable.LastCheckEndTime().Add(interval)) {
		// skip creating the check if its interval hasn't elapsed yet
		return nil, false, nil
	}

	checkPlan := checkable.CheckPlan(from, interval, resourceTypes.Filter(checkable), sourceDefaults)

	plan := c.planFactory.NewPlan(checkPlan)

	build, created, err := checkable.CreateBuild(ctx, manuallyTriggered, plan)
	if err != nil {
		return nil, false, fmt.Errorf("create build: %w", err)
	}

	if !created {
		return nil, false, nil
	}

	logger.Info("created-build", build.LagerData())

	return build, true, nil
}

func (c *checkFactory) Resources() ([]Resource, error) {
	var resources []Resource

	rows, err := resourcesQuery.
		LeftJoin("(select DISTINCT(resource_id) FROM job_inputs) ji ON ji.resource_id = r.id").
		LeftJoin("(select DISTINCT(resource_id) FROM job_outputs) jo ON jo.resource_id = r.id").
		Where(sq.Or{
			sq.And{
				sq.Eq{"p.paused": false},
				sq.Eq{"r.active": true},
				sq.NotEq{"ji.resource_id": nil},
			},
			sq.And{
				sq.Eq{"p.paused": false},
				sq.Expr("b.status IN ('aborted','failed','errored')"),
				sq.Eq{"ji.resource_id": nil},
			},
		}).
		RunWith(c.conn).
		Query()

	if err != nil {
		return nil, err
	}

	defer Close(rows)

	for rows.Next() {
		r := newEmptyResource(c.conn, c.lockFactory)
		err = scanResource(r, rows)
		if err != nil {
			return nil, err
		}

		resources = append(resources, r)
	}

	return resources, nil
}

func (c *checkFactory) ResourceTypes() ([]ResourceType, error) {
	var resourceTypes []ResourceType

	rows, err := resourceTypesQuery.
		RunWith(c.conn).
		Query()

	if err != nil {
		return nil, err
	}

	defer Close(rows)

	for rows.Next() {
		r := newEmptyResourceType(c.conn, c.lockFactory)
		err = scanResourceType(r, rows)
		if err != nil {
			return nil, err
		}

		resourceTypes = append(resourceTypes, r)
	}

	return resourceTypes, nil
}
<commit_msg>Refactor common parts of where clause<commit_after>package db

import (
	"context"
	"fmt"
	"time"

	"code.cloudfoundry.org/lager/lagerctx"
	sq "github.com/Masterminds/squirrel"
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/creds"
	"github.com/concourse/concourse/atc/db/lock"
)

//go:generate counterfeiter . Checkable

type Checkable interface {
	PipelineRef

	Name() string
	TeamID() int
	ResourceConfigScopeID() int
	TeamName() string
	Type() string
	Source() atc.Source
	Tags() atc.Tags
	CheckEvery() *atc.CheckEvery
	CheckTimeout() string
	LastCheckEndTime() time.Time
	CurrentPinnedVersion() atc.Version

	HasWebhook() bool

	CheckPlan(atc.Version, time.Duration, ResourceTypes, atc.Source) atc.CheckPlan
	CreateBuild(context.Context, bool, atc.Plan) (Build, bool, error)
}

//go:generate counterfeiter . CheckFactory

type CheckFactory interface {
	TryCreateCheck(context.Context, Checkable, ResourceTypes, atc.Version, bool) (Build, bool, error)
	Resources() ([]Resource, error)
	ResourceTypes() ([]ResourceType, error)
}

type checkFactory struct {
	conn        Conn
	lockFactory lock.LockFactory

	secrets       creds.Secrets
	varSourcePool creds.VarSourcePool

	planFactory atc.PlanFactory

	defaultCheckTimeout             time.Duration
	defaultCheckInterval            time.Duration
	defaultWithWebhookCheckInterval time.Duration
}

type CheckDurations struct {
	Timeout             time.Duration
	Interval            time.Duration
	IntervalWithWebhook time.Duration
}

func NewCheckFactory(
	conn Conn,
	lockFactory lock.LockFactory,
	secrets creds.Secrets,
	varSourcePool creds.VarSourcePool,
	durations CheckDurations,
) CheckFactory {
	return &checkFactory{
		conn:        conn,
		lockFactory: lockFactory,

		secrets:       secrets,
		varSourcePool: varSourcePool,

		planFactory: atc.NewPlanFactory(time.Now().Unix()),

		defaultCheckTimeout:             durations.Timeout,
		defaultCheckInterval:            durations.Interval,
		defaultWithWebhookCheckInterval: durations.IntervalWithWebhook,
	}
}

func (c *checkFactory) TryCreateCheck(ctx context.Context, checkable Checkable, resourceTypes ResourceTypes, from atc.Version, manuallyTriggered bool) (Build, bool, error) {
	logger := lagerctx.FromContext(ctx)

	var err error

	sourceDefaults := atc.Source{}
	parentType, found := resourceTypes.Parent(checkable)
	if found {
		if parentType.Version() == nil {
			return nil, false, fmt.Errorf("resource type '%s' has no version", parentType.Name())
		}
		sourceDefaults = parentType.Defaults()
	} else {
		defaults, found := atc.FindBaseResourceTypeDefaults(checkable.Type())
		if found {
			sourceDefaults = defaults
		}
	}

	interval := c.defaultCheckInterval
	if checkable.HasWebhook() {
		interval = c.defaultWithWebhookCheckInterval
	}
	if checkable.CheckEvery() != nil && !checkable.CheckEvery().Never {
		interval = checkable.CheckEvery().Interval
	}

	if !manuallyTriggered && time.Now().Before(checkable.LastCheckEndTime().Add(interval)) {
		// skip creating the check if its interval hasn't elapsed yet
		return nil, false, nil
	}

	checkPlan := checkable.CheckPlan(from, interval, resourceTypes.Filter(checkable), sourceDefaults)

	plan := c.planFactory.NewPlan(checkPlan)

	build, created, err := checkable.CreateBuild(ctx, manuallyTriggered, plan)
	if err != nil {
		return nil, false, fmt.Errorf("create build: %w", err)
	}

	if !created {
		return nil, false, nil
	}

	logger.Info("created-build", build.LagerData())

	return build, true, nil
}

func (c *checkFactory) Resources() ([]Resource, error) {
	var resources []Resource

	rows, err := resourcesQuery.
		LeftJoin("(select DISTINCT(resource_id) FROM job_inputs) ji ON ji.resource_id = r.id").
		LeftJoin("(select DISTINCT(resource_id) FROM job_outputs) jo ON jo.resource_id = r.id").
		Where(sq.And{
			sq.Eq{"p.paused": false},
			sq.Eq{"r.active": true},
		}).
		Where(sq.Or{
			sq.And{
				// find all resources that are inputs to jobs
				sq.NotEq{"ji.resource_id": nil},
			},
			sq.And{
				// find put-only resources that have errored
				sq.Expr("b.status IN ('aborted','failed','errored')"),
				sq.Eq{"ji.resource_id": nil},
			},
		}).
		RunWith(c.conn).
		Query()

	if err != nil {
		return nil, err
	}

	defer Close(rows)

	for rows.Next() {
		r := newEmptyResource(c.conn, c.lockFactory)
		err = scanResource(r, rows)
		if err != nil {
			return nil, err
		}

		resources = append(resources, r)
	}

	return resources, nil
}

func (c *checkFactory) ResourceTypes() ([]ResourceType, error) {
	var resourceTypes []ResourceType

	rows, err := resourceTypesQuery.
		RunWith(c.conn).
		Query()

	if err != nil {
		return nil, err
	}

	defer Close(rows)

	for rows.Next() {
		r := newEmptyResourceType(c.conn, c.lockFactory)
		err = scanResourceType(r, rows)
		if err != nil {
			return nil, err
		}

		resourceTypes = append(resourceTypes, r)
	}

	return resourceTypes, nil
}
<|endoftext|>


<commit_before>package batcher

import (
	"encoding/gob"
	"fmt"
	"os"
	"reflect"
	"testing"

	"github.com/kylelemons/godebug/pretty"
	"github.com/pagarme/teleport/action"
	"github.com/pagarme/teleport/client"
	"github.com/pagarme/teleport/config"
	"github.com/pagarme/teleport/database"
)

var db *database.Database
var batcher *Batcher

func init() {
	gob.Register(&StubAction{})

	config := config.New()
	err := config.ReadFromFile("../config_test.yml")

	if err != nil {
		fmt.Printf("Error opening config file: %v\n", err)
		os.Exit(1)
	}

	db = database.New(config.Database)

	// Start db
	if err = db.Start(); err != nil {
		fmt.Printf("Erro starting database: %v", err)
		os.Exit(1)
	}

	targets := make(map[string]*client.Client)

	for key, target := range config.Targets {
		targets[key] = client.New(target)
	}

	batcher = New(db, targets, -1)
}

// StubAction implements Action
type StubAction struct {
	ShouldFilter  bool
	SeparateBatch bool
}

func (a *StubAction) Execute(c *action.Context) error {
	return nil
}

func (a *StubAction) Filter(targetExpression string) bool {
	return a.ShouldFilter
}

func (a *StubAction) NeedsSeparatedBatch() bool {
	return a.SeparateBatch
}

func TestMarkEventsBatched(t *testing.T) {
	db.Db.Exec(`
		TRUNCATE teleport.event;
		TRUNCATE teleport.batch;
	`)

	tx := db.NewTransaction()
	stubEvent := &database.Event{
		Kind:          "ddl",
		Status:        "waiting_batch",
		TriggerTag:    "TAG",
		TriggerEvent:  "EVENT",
		TransactionId: "123",
	}
	stubEvent.InsertQuery(tx)
	tx.Commit()

	tx = db.NewTransaction()

	err := batcher.markEventsBatched([]*database.Event{stubEvent}, tx)

	if err != nil {
		t.Errorf("mark events batched returned error: %#v\n", err)
	}

	tx.Commit()

	batchedEvents, _ := db.GetEvents("batched", -1)
	var updatedEvent *database.Event

	for _, event := range batchedEvents {
		if stubEvent.Id == event.Id {
			updatedEvent = event
			break
		}
	}

	if updatedEvent == nil {
		t.Errorf("ignored event => nil, want %v", stubEvent)
	}
}

func TestCreateBatchesWithActions(t *testing.T) {
	testAction := &StubAction{true, false}
	separateAction := &StubAction{true, true}

	actionsForTarget := map[string][]action.Action{
		"test_target": []action.Action{
			testAction,
			testAction,
			separateAction,
			testAction,
		},
	}

	tx := batcher.db.NewTransaction()

	batches, err := batcher.CreateBatchesWithActions(actionsForTarget, tx)

	if err != nil {
		t.Errorf("create batches returned error: %#v", err)
	}

	if len(batches) != 3 {
		t.Errorf("batches => %d, want %d", len(batches), 3)
	}

	expectedActions := [][]action.Action{
		[]action.Action{
			testAction,
			testAction,
		},
		[]action.Action{
			separateAction,
		},
		[]action.Action{
			testAction,
		},
	}

	for i, batch := range batches {
		actions, _ := batch.GetActions()

		if !reflect.DeepEqual(expectedActions[i], actions) {
			t.Errorf(
				"actions for batch %d => %#v, want %#v",
				i,
				actions,
				expectedActions[i],
			)
		}
	}
}

func TestActionsForTarget(t *testing.T) {
	batcher.db.Schemas = map[string]*database.Schema{
		"public": &database.Schema{
			Tables: []*database.Table{
				&database.Table{
					RelationKind: "r",
					RelationName: "test_table",
					Columns: []*database.Column{
						&database.Column{
							Name:         "id",
							Num:          1,
							TypeName:     "int4",
							TypeSchema:   "pg_catalog",
							TypeOid:      "123",
							IsPrimaryKey: true,
							Table:        nil,
						},
						&database.Column{
							Name:         "content",
							Num:          2,
							TypeName:     "text",
							TypeSchema:   "pg_catalog",
							TypeOid:      "124",
							IsPrimaryKey: false,
							Table:        nil,
						},
					},
				},
			},
		},
	}

	dataEvent1 := `{
		"pre":[{"oid":"2200","schema_name":"public","owner_id":"10","classes":
			[{"oid":"16443","namespace_oid":"2200","relation_kind":"r","relation_name":"test_table","columns":
				[
					{"class_oid":"16443","attr_name":"id","attr_num":1,"type_name":"int4","type_oid":"23","is_primary_key":true,"type_schema":"pg_catalog"}
				]
			}]
		}],
		"post":[{"oid":"2200","schema_name":"public","owner_id":"10","classes":
			[{"oid":"16443","namespace_oid":"2200","relation_kind":"r","relation_name":"test_table","columns":
				[
					{"class_oid":"16443","attr_name":"id","attr_num":1,"type_name":"int4","type_oid":"23","is_primary_key":true,"type_schema":"pg_catalog"},
					{"class_oid":"16443","attr_name":"content","attr_num":2,"type_name":"text","type_oid":"24","type_schema":"pg_catalog"}
				]
			}]
		}]
	}`

	dataEvent2 := `{
		"pre":null,
		"post":{
			"id": 5
		}
	}`

	events := database.Events{
		&database.Event{
			Kind:          "ddl",
			Status:        "waiting_batch",
			TriggerTag:    "TAG",
			TriggerEvent:  "EVENT",
			TransactionId: "123",
			Data:          &dataEvent1,
		},
		&database.Event{
			Kind:          "dml",
			Status:        "waiting_batch",
			TriggerTag:    "public.test_table",
			TriggerEvent:  "INSERT",
			TransactionId: "123",
			Data:          &dataEvent2,
		},
	}

	actionsForTargets, err := batcher.actionsForTargets(events)

	if err != nil {
		t.Errorf("actions for targets returned error: %#v", err)
	}

	expectedActions := map[string][]action.Action{
		"test_target": []action.Action{
			&action.CreateColumn{
				SchemaName: "live",
				TableName:  "test_table",
				Column: action.Column{
					Name:         "content",
					Type:         "text",
					IsNativeType: true,
				},
			},
			&action.InsertRow{
				SchemaName:     "live",
				TableName:      "test_table",
				PrimaryKeyName: "id",
				Rows: action.Rows{
					action.Row{
						Value: 5,
						Column: action.Column{
							Name:         "id",
							Type:         "int4",
							IsNativeType: true,
						},
					},
				},
			},
		},
	}

	if diff := pretty.Compare(expectedActions, actionsForTargets); diff != "" {
		t.Errorf(
			"actions for target => %s",
			diff,
		)
	}
}
<commit_msg>Fix typo<commit_after>package batcher

import (
	"encoding/gob"
	"fmt"
	"os"
	"reflect"
	"testing"

	"github.com/kylelemons/godebug/pretty"
	"github.com/pagarme/teleport/action"
	"github.com/pagarme/teleport/client"
	"github.com/pagarme/teleport/config"
	"github.com/pagarme/teleport/database"
)

var db *database.Database
var batcher *Batcher

func init() {
	gob.Register(&StubAction{})

	config := config.New()
	err := config.ReadFromFile("../config_test.yml")

	if err != nil {
		fmt.Printf("Error opening config file: %v\n", err)
		os.Exit(1)
	}

	db = database.New(config.Database)

	// Start db
	if err = db.Start(); err != nil {
		fmt.Printf("Error starting database: %v", err)
		os.Exit(1)
	}

	targets := make(map[string]*client.Client)

	for key, target := range config.Targets {
		targets[key] = client.New(target)
	}

	batcher = New(db, targets, -1)
}

// StubAction implements Action
type StubAction struct {
	ShouldFilter  bool
	SeparateBatch bool
}

func (a *StubAction) Execute(c *action.Context) error {
	return nil
}

func (a *StubAction) Filter(targetExpression string) bool {
	return a.ShouldFilter
}

func (a *StubAction) NeedsSeparatedBatch() bool {
	return a.SeparateBatch
}

func TestMarkEventsBatched(t *testing.T) {
	db.Db.Exec(`
		TRUNCATE teleport.event;
		TRUNCATE teleport.batch;
	`)

	tx := db.NewTransaction()
	stubEvent := &database.Event{
		Kind:          "ddl",
		Status:        "waiting_batch",
		TriggerTag:    "TAG",
		TriggerEvent:  "EVENT",
		TransactionId: "123",
	}
	stubEvent.InsertQuery(tx)
	tx.Commit()

	tx = db.NewTransaction()

	err := batcher.markEventsBatched([]*database.Event{stubEvent}, tx)

	if err != nil {
		t.Errorf("mark events batched returned error: %#v\n", err)
	}

	tx.Commit()

	batchedEvents, _ := db.GetEvents("batched", -1)
	var updatedEvent *database.Event

	for _, event := range batchedEvents {
		if stubEvent.Id == event.Id {
			updatedEvent = event
			break
		}
	}

	if updatedEvent == nil {
		t.Errorf("ignored event => nil, want %v", stubEvent)
	}
}

func TestCreateBatchesWithActions(t *testing.T) {
	testAction := &StubAction{true, false}
	separateAction := &StubAction{true, true}

	actionsForTarget := map[string][]action.Action{
		"test_target": []action.Action{
			testAction,
			testAction,
			separateAction,
			testAction,
		},
	}

	tx := batcher.db.NewTransaction()

	batches, err := batcher.CreateBatchesWithActions(actionsForTarget, tx)

	if err != nil {
		t.Errorf("create batches returned error: %#v", err)
	}

	if len(batches) != 3 {
		t.Errorf("batches => %d, want %d", len(batches), 3)
	}

	expectedActions := [][]action.Action{
		[]action.Action{
			testAction,
			testAction,
		},
		[]action.Action{
			separateAction,
		},
		[]action.Action{
			testAction,
		},
	}

	for i, batch := range batches {
		actions, _ := batch.GetActions()

		if !reflect.DeepEqual(expectedActions[i], actions) {
			t.Errorf(
				"actions for batch %d => %#v, want %#v",
				i,
				actions,
				expectedActions[i],
			)
		}
	}
}

func TestActionsForTarget(t *testing.T) {
	batcher.db.Schemas = map[string]*database.Schema{
		"public": &database.Schema{
			Tables: []*database.Table{
				&database.Table{
					RelationKind: "r",
					RelationName: "test_table",
					Columns: []*database.Column{
						&database.Column{
							Name:         "id",
							Num:          1,
							TypeName:     "int4",
							TypeSchema:   "pg_catalog",
							TypeOid:      "123",
							IsPrimaryKey: true,
							Table:        nil,
						},
						&database.Column{
							Name:         "content",
							Num:          2,
							TypeName:     "text",
							TypeSchema:   "pg_catalog",
							TypeOid:      "124",
							IsPrimaryKey: false,
							Table:        nil,
						},
					},
				},
			},
		},
	}

	dataEvent1 := `{
		"pre":[{"oid":"2200","schema_name":"public","owner_id":"10","classes":
			[{"oid":"16443","namespace_oid":"2200","relation_kind":"r","relation_name":"test_table","columns":
				[
					{"class_oid":"16443","attr_name":"id","attr_num":1,"type_name":"int4","type_oid":"23","is_primary_key":true,"type_schema":"pg_catalog"}
				]
			}]
		}],
		"post":[{"oid":"2200","schema_name":"public","owner_id":"10","classes":
			[{"oid":"16443","namespace_oid":"2200","relation_kind":"r","relation_name":"test_table","columns":
				[
					{"class_oid":"16443","attr_name":"id","attr_num":1,"type_name":"int4","type_oid":"23","is_primary_key":true,"type_schema":"pg_catalog"},
					{"class_oid":"16443","attr_name":"content","attr_num":2,"type_name":"text","type_oid":"24","type_schema":"pg_catalog"}
				]
			}]
		}]
	}`

	dataEvent2 := `{
		"pre":null,
		"post":{
			"id": 5
		}
	}`

	events := database.Events{
		&database.Event{
			Kind:          "ddl",
			Status:        "waiting_batch",
			TriggerTag:    "TAG",
			TriggerEvent:  "EVENT",
			TransactionId: "123",
			Data:          &dataEvent1,
		},
		&database.Event{
			Kind:          "dml",
			Status:        "waiting_batch",
			TriggerTag:    "public.test_table",
			TriggerEvent:  "INSERT",
			TransactionId: "123",
			Data:          &dataEvent2,
		},
	}

	actionsForTargets, err := batcher.actionsForTargets(events)

	if err != nil {
		t.Errorf("actions for targets returned error: %#v", err)
	}

	expectedActions := map[string][]action.Action{
		"test_target": []action.Action{
			&action.CreateColumn{
				SchemaName: "live",
				TableName:  "test_table",
				Column: action.Column{
					Name:         "content",
					Type:         "text",
					IsNativeType: true,
				},
			},
			&action.InsertRow{
				SchemaName:     "live",
				TableName:      "test_table",
				PrimaryKeyName: "id",
				Rows: action.Rows{
					action.Row{
						Value: 5,
						Column: action.Column{
							Name:         "id",
							Type:         "int4",
							IsNativeType: true,
						},
					},
				},
			},
		},
	}

	if diff := pretty.Compare(expectedActions, actionsForTargets); diff != "" {
		t.Errorf(
			"actions for target => %s",
			diff,
		)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/catatsuy/isucon6-final/bench/fails"
	"github.com/catatsuy/isucon6-final/bench/scenario"
	"github.com/catatsuy/isucon6-final/bench/score"
)

var BenchmarkTimeout int
var InitialCheckOnly bool
var MatsuriNum = 10
var LoadIndexPageNum = 10
var DrawOnRandomRoomNum = 2

func main() {

	var urls string

	flag.StringVar(&urls, "urls", "", "ベンチマーク対象のURL（scheme, host, portまで。カンマ区切りで複数可。例： https://xxx.xxx.xxx.xxx,https://xxx.xxx.xxx.xxx:1443）")
	flag.IntVar(&BenchmarkTimeout, "timeout", 60, "ソフトタイムアウト")
	flag.BoolVar(&InitialCheckOnly, "initialcheck", false, "初期チェックだけ行う")

	flag.Parse()

	origins, err := makeOrigins(urls)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
	}

	initialCheck(origins)

	// 初期チェックのみモードではない、かつ、この時点でcriticalが出ていなければ負荷をかけにいく
	if !InitialCheckOnly && !fails.GetIsCritical() {
		benchmark(origins)
	}

	output()
}

func makeOrigins(urls string) ([]string, error) {
	if urls == "" {
		return nil, errors.New("urlsが指定されていません")
	}
	origins := strings.Split(urls, ",")
	for _, origin := range origins {
		u, err := url.Parse(origin)
		if err != nil {
			return nil, err
		}
		if u.Scheme != "https" || u.Host == "" {
			return nil, errors.New("urlsの指定が間違っています")
		}
	}
	return origins, nil
}

func initialCheck(origins []string) {
	scenario.CSRFTokenRefreshed(origins)
	scenario.StrokeReflectedToTop(origins)
	scenario.RoomWithoutStrokeNotShownAtTop(origins)
	scenario.CantDrawFirstStrokeOnSomeoneElsesRoom(origins)
	scenario.TopPageContent(origins)
	scenario.APIAndHTMLMustBeConsistent(origins)
	scenario.CheckStaticFiles(origins)
}

func benchmark(origins []string) {
	var wg sync.WaitGroup
	for i := 0; i < MatsuriNum; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			scenario.Matsuri(origins, BenchmarkTimeout-5)
		}()
	}

	loadIndexPageCh := makeChan(LoadIndexPageNum)
	drawOnRandomRoomCh := makeChan(DrawOnRandomRoomNum)
	timeoutCh := time.After(time.Duration(BenchmarkTimeout) * time.Second)

L:
	for {
		select {
		case <-loadIndexPageCh:
			go func() {
				scenario.LoadIndexPage(origins)
				time.Sleep(100 * time.Millisecond)
				loadIndexPageCh <- struct{}{}
			}()
		case <-drawOnRandomRoomCh:
			go func() {
				scenario.DrawOnRandomRoom(origins)
				time.Sleep(500 * time.Millisecond)
				drawOnRandomRoomCh <- struct{}{}
			}()
		case <-timeoutCh:
			break L
		}
	}

	wg.Wait()
}

func output() {
	b, _ := json.Marshal(score.Output{
		Pass:     !fails.GetIsCritical(),
		Score:    score.Get(),
		Messages: fails.GetUnique(),
	})

	fmt.Println(string(b))
}

func makeChan(len int) chan struct{} {
	ch := make(chan struct{}, len)
	for i := 0; i < len; i++ {
		ch <- struct{}{}
	}
	return ch
}
<commit_msg>Call random seed<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"math/rand"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/catatsuy/isucon6-final/bench/fails"
	"github.com/catatsuy/isucon6-final/bench/scenario"
	"github.com/catatsuy/isucon6-final/bench/score"
)

var BenchmarkTimeout int
var InitialCheckOnly bool
var MatsuriNum = 10
var LoadIndexPageNum = 10
var DrawOnRandomRoomNum = 2

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {

	var urls string

	flag.StringVar(&urls, "urls", "", "ベンチマーク対象のURL（scheme, host, portまで。カンマ区切りで複数可。例： https://xxx.xxx.xxx.xxx,https://xxx.xxx.xxx.xxx:1443）")
	flag.IntVar(&BenchmarkTimeout, "timeout", 60, "ソフトタイムアウト")
	flag.BoolVar(&InitialCheckOnly, "initialcheck", false, "初期チェックだけ行う")

	flag.Parse()

	origins, err := makeOrigins(urls)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
	}

	initialCheck(origins)

	// 初期チェックのみモードではない、かつ、この時点でcriticalが出ていなければ負荷をかけにいく
	if !InitialCheckOnly && !fails.GetIsCritical() {
		benchmark(origins)
	}

	output()
}

func makeOrigins(urls string) ([]string, error) {
	if urls == "" {
		return nil, errors.New("urlsが指定されていません")
	}
	origins := strings.Split(urls, ",")
	for _, origin := range origins {
		u, err := url.Parse(origin)
		if err != nil {
			return nil, err
		}
		if u.Scheme != "https" || u.Host == "" {
			return nil, errors.New("urlsの指定が間違っています")
		}
	}
	return origins, nil
}

func initialCheck(origins []string) {
	scenario.CSRFTokenRefreshed(origins)
	scenario.StrokeReflectedToTop(origins)
	scenario.RoomWithoutStrokeNotShownAtTop(origins)
	scenario.CantDrawFirstStrokeOnSomeoneElsesRoom(origins)
	scenario.TopPageContent(origins)
	scenario.APIAndHTMLMustBeConsistent(origins)
	scenario.CheckStaticFiles(origins)
}

func benchmark(origins []string) {
	var wg sync.WaitGroup
	for i := 0; i < MatsuriNum; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			scenario.Matsuri(origins, BenchmarkTimeout-5)
		}()
	}

	loadIndexPageCh := makeChan(LoadIndexPageNum)
	drawOnRandomRoomCh := makeChan(DrawOnRandomRoomNum)
	timeoutCh := time.After(time.Duration(BenchmarkTimeout) * time.Second)

L:
	for {
		select {
		case <-loadIndexPageCh:
			go func() {
				scenario.LoadIndexPage(origins)
				time.Sleep(100 * time.Millisecond)
				loadIndexPageCh <- struct{}{}
			}()
		case <-drawOnRandomRoomCh:
			go func() {
				scenario.DrawOnRandomRoom(origins)
				time.Sleep(500 * time.Millisecond)
				drawOnRandomRoomCh <- struct{}{}
			}()
		case <-timeoutCh:
			break L
		}
	}

	wg.Wait()
}

func output() {
	b, _ := json.Marshal(score.Output{
		Pass:     !fails.GetIsCritical(),
		Score:    score.Get(),
		Messages: fails.GetUnique(),
	})

	fmt.Println(string(b))
}

func makeChan(len int) chan struct{} {
	ch := make(chan struct{}, len)
	for i := 0; i < len; i++ {
		ch <- struct{}{}
	}
	return ch
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"time"

	"net/url"
	"strings"

	"errors"
	"os"

	"github.com/catatsuy/isucon6-final/bench/fails"
	"github.com/catatsuy/isucon6-final/bench/scenario"
	"github.com/catatsuy/isucon6-final/bench/score"
)

var BenchmarkTimeout int
var InitialCheckOnly bool
var MatsuriNum int = 10

func main() {

	var urls string

	flag.StringVar(&urls, "urls", "", "ベンチマーク対象のURL（scheme, host, portまで。カンマ区切りで複数可。例： https://xxx.xxx.xxx.xxx,https://xxx.xxx.xxx.xxx:1443）")
	flag.IntVar(&BenchmarkTimeout, "timeout", 60, "ソフトタイムアウト")
	flag.BoolVar(&InitialCheckOnly, "initialcheck", false, "初期チェックだけ行う")

	flag.Parse()

	origins, err := makeOrigins(urls)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
	}

	// 初期チェックで失敗したらそこで終了
	initialCheck(origins)
	if fails.GetIsCritical() || !InitialCheckOnly {
		output()
		return
	}

	benchmark(origins)
	output()
}

func makeOrigins(urls string) ([]string, error) {
	if urls == "" {
		return nil, errors.New("urlsが指定されていません")
	}
	origins := strings.Split(urls, ",")
	for _, origin := range origins {
		u, err := url.Parse(origin)
		if err != nil {
			return nil, err
		}
		if u.Scheme != "https" || u.Host == "" {
			return nil, errors.New("urlsの指定が間違っています")
		}
	}
	return origins, nil
}

func initialCheck(origins []string) {
	scenario.CSRFTokenRefreshed(origins)
	scenario.StrokeReflectedToTop(origins)
	scenario.RoomWithoutStrokeNotShownAtTop(origins)
	scenario.StrokeReflectedToSVG(origins)
	scenario.CantDrawFirstStrokeOnSomeoneElsesRoom(origins)
}

func benchmark(origins []string) {
	loadIndexPageCh := makeChan(2)
	loadRoomPageCh := makeChan(2)
	checkCSRFTokenRefreshedCh := makeChan(1)
	matsuriCh := makeChan(MatsuriNum)
	matsuriEndCh := make(chan struct{}, MatsuriNum)

	timeoutCh := time.After(time.Duration(BenchmarkTimeout) * time.Second)

L:
	for {
		select {
		case <-loadIndexPageCh:
			go func() {
				scenario.LoadIndexPage(origins)
				loadIndexPageCh <- struct{}{}
			}()
		case <-loadRoomPageCh:
			go func() {
				scenario.LoadRoomPage(origins)
				loadRoomPageCh <- struct{}{}
			}()
		case <-checkCSRFTokenRefreshedCh:
			go func() {
				scenario.CSRFTokenRefreshed(origins)
				checkCSRFTokenRefreshedCh <- struct{}{}
			}()
		case <-matsuriCh:
			go func() {
				scenario.Matsuri(origins, BenchmarkTimeout-5)
				//matsuriRoomCh <- struct{}{} // Never again.
				matsuriEndCh <- struct{}{}
			}()
		case <-timeoutCh:
			break L
		}
	}
	for i := 0; i < MatsuriNum; i++ {
		<-matsuriEndCh
	}
}

func output() {
	s := score.Get()
	pass := true
	if fails.GetIsCritical() {
		s = 0
		pass = false
	}
	b, _ := json.Marshal(score.Output{
		Pass:     pass,
		Score:    s,
		Messages: fails.GetUnique(),
	})

	fmt.Println(string(b))
}

func makeChan(len int) chan struct{} {
	ch := make(chan struct{}, len)
	for i := 0; i < len; i++ {
		ch <- struct{}{}
	}
	return ch
}
<commit_msg>Remove checkCSRFTokenRefreshedCh<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"time"

	"net/url"
	"strings"

	"errors"
	"os"

	"github.com/catatsuy/isucon6-final/bench/fails"
	"github.com/catatsuy/isucon6-final/bench/scenario"
	"github.com/catatsuy/isucon6-final/bench/score"
)

var BenchmarkTimeout int
var InitialCheckOnly bool
var MatsuriNum int = 10

func main() {

	var urls string

	flag.StringVar(&urls, "urls", "", "ベンチマーク対象のURL（scheme, host, portまで。カンマ区切りで複数可。例： https://xxx.xxx.xxx.xxx,https://xxx.xxx.xxx.xxx:1443）")
	flag.IntVar(&BenchmarkTimeout, "timeout", 60, "ソフトタイムアウト")
	flag.BoolVar(&InitialCheckOnly, "initialcheck", false, "初期チェックだけ行う")

	flag.Parse()

	origins, err := makeOrigins(urls)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
	}

	// 初期チェックで失敗したらそこで終了
	initialCheck(origins)
	if fails.GetIsCritical() || !InitialCheckOnly {
		output()
		return
	}

	benchmark(origins)
	output()
}

func makeOrigins(urls string) ([]string, error) {
	if urls == "" {
		return nil, errors.New("urlsが指定されていません")
	}
	origins := strings.Split(urls, ",")
	for _, origin := range origins {
		u, err := url.Parse(origin)
		if err != nil {
			return nil, err
		}
		if u.Scheme != "https" || u.Host == "" {
			return nil, errors.New("urlsの指定が間違っています")
		}
	}
	return origins, nil
}

func initialCheck(origins []string) {
	scenario.CSRFTokenRefreshed(origins)
	scenario.StrokeReflectedToTop(origins)
	scenario.RoomWithoutStrokeNotShownAtTop(origins)
	scenario.StrokeReflectedToSVG(origins)
	scenario.CantDrawFirstStrokeOnSomeoneElsesRoom(origins)
}

func benchmark(origins []string) {
	loadIndexPageCh := makeChan(2)
	loadRoomPageCh := makeChan(2)
	matsuriCh := makeChan(MatsuriNum)
	matsuriEndCh := make(chan struct{}, MatsuriNum)

	timeoutCh := time.After(time.Duration(BenchmarkTimeout) * time.Second)

L:
	for {
		select {
		case <-loadIndexPageCh:
			go func() {
				scenario.LoadIndexPage(origins)
				loadIndexPageCh <- struct{}{}
			}()
		case <-loadRoomPageCh:
			go func() {
				scenario.LoadRoomPage(origins)
				loadRoomPageCh <- struct{}{}
			}()
		case <-matsuriCh:
			go func() {
				scenario.Matsuri(origins, BenchmarkTimeout-5)
				//matsuriRoomCh <- struct{}{} // Never again.
				matsuriEndCh <- struct{}{}
			}()
		case <-timeoutCh:
			break L
		}
	}
	for i := 0; i < MatsuriNum; i++ {
		<-matsuriEndCh
	}
}

func output() {
	s := score.Get()
	pass := true
	if fails.GetIsCritical() {
		s = 0
		pass = false
	}
	b, _ := json.Marshal(score.Output{
		Pass:     pass,
		Score:    s,
		Messages: fails.GetUnique(),
	})

	fmt.Println(string(b))
}

func makeChan(len int) chan struct{} {
	ch := make(chan struct{}, len)
	for i := 0; i < len; i++ {
		ch <- struct{}{}
	}
	return ch
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC 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 remote

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"

	"github.com/google/go-containerregistry/authn"
	"github.com/google/go-containerregistry/name"
	"github.com/google/go-containerregistry/v1"
	"github.com/google/go-containerregistry/v1/remote/transport"
)

// WriteOptions are used to expose optional information to guide or
// control the image write.
type WriteOptions struct {
	// The set of paths from which to attempt to mount blobs.
	MountPaths []name.Repository
	// TODO(mattmoor): Expose "threads" to limit parallelism?
}

// Write pushes the provided img to the specified image reference.
func Write(ref name.Reference, img v1.Image, auth authn.Authenticator, t http.RoundTripper,
	wo WriteOptions) error {

	scopes := []string{ref.Scope(transport.PushScope)}
	for _, mp := range wo.MountPaths {
		scopes = append(scopes, mp.Scope(transport.PullScope))
	}

	tr, err := transport.New(ref.Context().Registry, auth, t, scopes)
	if err != nil {
		return err
	}
	w := writer{
		ref:     ref,
		client:  &http.Client{Transport: tr},
		img:     img,
		options: wo,
	}

	bs, err := img.BlobSet()
	if err != nil {
		return err
	}

	// Spin up go routines to publish each of the members of BlobSet(),
	// and use an error channel to collect their results.
	errCh := make(chan error)
	defer close(errCh)
	for h := range bs {
		go func(h v1.Hash) {
			errCh <- w.uploadOne(h)
		}(h)
	}

	// Now wait for all of the blob uploads to complete.
	var errors []error
	for _ = range bs {
		if err := <-errCh; err != nil {
			errors = append(errors, err)
		}
	}
	if len(errors) > 0 {
		// Return the first error we encountered.
		return errors[0]
	}

	// With all of the constituent elements uploaded, upload the manifest
	// to commit the image.
	return w.commitImage()
}

// writer writes the elements of an image to a remote image reference.
type writer struct {
	ref     name.Reference
	client  *http.Client
	img     v1.Image
	options WriteOptions
}

// url returns a url.Url for the specified path in the context of this remote image reference.
func (w *writer) url(path string) url.URL {
	return url.URL{
		Scheme: transport.Scheme(w.ref.Context().Registry),
		Host:   w.ref.Context().RegistryStr(),
		Path:   path,
	}
}

// nextLocation extracts the fully-qualified URL to which we should send the next request in an upload sequence.
func (w *writer) nextLocation(resp *http.Response) (string, error) {
	loc := resp.Header.Get("Location")
	if len(loc) == 0 {
		return "", errors.New("missing Location header")
	}
	u, err := url.Parse(loc)
	if err != nil {
		return "", err
	}

	// If the location header returned is just a url path, then fully qualify it.
	// We cannot simply call w.url, since there might be an embedded query string.
	return resp.Request.URL.ResolveReference(u).String(), nil
}

// initiateUpload initiates the blob upload, which starts with a POST that can
// optionally include the hash of the layer and a list of repositories from
// which that layer might be read. On failure, an error is returned.
// On success, the layer was either mounted (nothing more to do) or a blob
// upload was initiated and the body of that blob should be sent to the returned
// location.
func (w *writer) initiateUpload(h v1.Hash) (location string, mounted bool, err error) {
	u := w.url(fmt.Sprintf("/v2/%s/blobs/uploads/", w.ref.Context().RepositoryStr()))
	uv := url.Values{
		"mount": []string{h.String()},
	}
	var from []string
	for _, m := range w.options.MountPaths {
		from = append(from, m.RepositoryStr())
	}
	// We currently avoid HEAD because it's semi-redundant with the mount that is part
	// of initiating the blob upload.  GCR will perform an existence check on the initiation
	// if "mount" is specified, even if no "from" sources are specified.  If this turns out
	// to not be broadly applicable then we should replace mounts without "from"s with a HEAD.
	if len(from) > 0 {
		uv["from"] = from
	}
	u.RawQuery = uv.Encode()

	// Make the request to initiate the blob upload.
	resp, err := w.client.Post(u.String(), "application/json", nil)
	if err != nil {
		return "", false, err
	}
	defer resp.Body.Close()

	if err := checkError(resp, http.StatusCreated, http.StatusAccepted); err != nil {
		return "", false, err
	}

	// Check the response code to determine the result.
	switch resp.StatusCode {
	case http.StatusCreated:
		// We're done, we were able to fast-path.
		return "", true, nil
	case http.StatusAccepted:
		// Proceed to PATCH, upload has begun.
		loc, err := w.nextLocation(resp)
		return loc, false, err
	default:
		panic("Unreachable: initiateUpload")
	}
}

// streamBlob streams the contents of the blob to the specified location.
// On failure, this will return an error.  On success, this will return the location
// header indicating how to commit the streamed blob.
func (w *writer) streamBlob(h v1.Hash, streamLocation string) (commitLocation string, err error) {
	l, err := w.img.LayerByDigest(h)
	if err != nil {
		return "", err
	}
	blob, err := l.Compressed()
	if err != nil {
		return "", err
	}
	defer blob.Close()

	req, err := http.NewRequest(http.MethodPatch, streamLocation, blob)
	if err != nil {
		return "", err
	}

	resp, err := w.client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if err := checkError(resp, http.StatusNoContent, http.StatusAccepted, http.StatusCreated); err != nil {
		return "", err
	}

	// The blob has been uploaded, return the location header indicating
	// how to commit this layer.
	return w.nextLocation(resp)
}

// commitBlob commits this blob by sending a PUT to the location returned from streaming the blob.
func (w *writer) commitBlob(h v1.Hash, location string) (err error) {
	u, err := url.Parse(location)
	if err != nil {
		return err
	}
	v := u.Query()
	v.Set("digest", h.String())
	u.RawQuery = v.Encode()

	req, err := http.NewRequest(http.MethodPut, u.String(), nil)
	if err != nil {
		return err
	}

	resp, err := w.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	return checkError(resp, http.StatusCreated)
}

// uploadOne performs a complete upload of a single layer.
func (w *writer) uploadOne(h v1.Hash) error {
	location, mounted, err := w.initiateUpload(h)
	if err != nil {
		return err
	} else if mounted {
		log.Printf("mounted blob: %v", h)
		return nil
	}

	location, err = w.streamBlob(h, location)
	if err != nil {
		return err
	}

	if err := w.commitBlob(h, location); err != nil {
		return err
	}
	log.Printf("pushed blob %v", h)
	return nil
}

// commitImage does a PUT of the image's manifest.
func (w *writer) commitImage() error {
	raw, err := w.img.RawManifest()
	if err != nil {
		return err
	}
	mt, err := w.img.MediaType()
	if err != nil {
		return err
	}

	u := w.url(fmt.Sprintf("/v2/%s/manifests/%s", w.ref.Context().RepositoryStr(), w.ref.Identifier()))

	// Make the request to PUT the serialized manifest
	req, err := http.NewRequest(http.MethodPut, u.String(), bytes.NewBuffer(raw))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", string(mt))

	resp, err := w.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if err := checkError(resp, http.StatusOK, http.StatusCreated, http.StatusAccepted); err != nil {
		return err
	}

	digest, err := w.img.Digest()
	if err != nil {
		return err
	}

	// The image was successfully pushed!
	fmt.Printf("%v: digest: %v size: %d\n", w.ref, digest, len(raw))
	return nil
}

// TODO(mattmoor): WriteIndex
<commit_msg>Log pushed digest to stderr (#176)<commit_after>// Copyright 2018 Google LLC 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 remote

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"

	"github.com/google/go-containerregistry/authn"
	"github.com/google/go-containerregistry/name"
	"github.com/google/go-containerregistry/v1"
	"github.com/google/go-containerregistry/v1/remote/transport"
)

// WriteOptions are used to expose optional information to guide or
// control the image write.
type WriteOptions struct {
	// The set of paths from which to attempt to mount blobs.
	MountPaths []name.Repository
	// TODO(mattmoor): Expose "threads" to limit parallelism?
}

// Write pushes the provided img to the specified image reference.
func Write(ref name.Reference, img v1.Image, auth authn.Authenticator, t http.RoundTripper,
	wo WriteOptions) error {

	scopes := []string{ref.Scope(transport.PushScope)}
	for _, mp := range wo.MountPaths {
		scopes = append(scopes, mp.Scope(transport.PullScope))
	}

	tr, err := transport.New(ref.Context().Registry, auth, t, scopes)
	if err != nil {
		return err
	}
	w := writer{
		ref:     ref,
		client:  &http.Client{Transport: tr},
		img:     img,
		options: wo,
	}

	bs, err := img.BlobSet()
	if err != nil {
		return err
	}

	// Spin up go routines to publish each of the members of BlobSet(),
	// and use an error channel to collect their results.
	errCh := make(chan error)
	defer close(errCh)
	for h := range bs {
		go func(h v1.Hash) {
			errCh <- w.uploadOne(h)
		}(h)
	}

	// Now wait for all of the blob uploads to complete.
	var errors []error
	for _ = range bs {
		if err := <-errCh; err != nil {
			errors = append(errors, err)
		}
	}
	if len(errors) > 0 {
		// Return the first error we encountered.
		return errors[0]
	}

	// With all of the constituent elements uploaded, upload the manifest
	// to commit the image.
	return w.commitImage()
}

// writer writes the elements of an image to a remote image reference.
type writer struct {
	ref     name.Reference
	client  *http.Client
	img     v1.Image
	options WriteOptions
}

// url returns a url.Url for the specified path in the context of this remote image reference.
func (w *writer) url(path string) url.URL {
	return url.URL{
		Scheme: transport.Scheme(w.ref.Context().Registry),
		Host:   w.ref.Context().RegistryStr(),
		Path:   path,
	}
}

// nextLocation extracts the fully-qualified URL to which we should send the next request in an upload sequence.
func (w *writer) nextLocation(resp *http.Response) (string, error) {
	loc := resp.Header.Get("Location")
	if len(loc) == 0 {
		return "", errors.New("missing Location header")
	}
	u, err := url.Parse(loc)
	if err != nil {
		return "", err
	}

	// If the location header returned is just a url path, then fully qualify it.
	// We cannot simply call w.url, since there might be an embedded query string.
	return resp.Request.URL.ResolveReference(u).String(), nil
}

// initiateUpload initiates the blob upload, which starts with a POST that can
// optionally include the hash of the layer and a list of repositories from
// which that layer might be read. On failure, an error is returned.
// On success, the layer was either mounted (nothing more to do) or a blob
// upload was initiated and the body of that blob should be sent to the returned
// location.
func (w *writer) initiateUpload(h v1.Hash) (location string, mounted bool, err error) {
	u := w.url(fmt.Sprintf("/v2/%s/blobs/uploads/", w.ref.Context().RepositoryStr()))
	uv := url.Values{
		"mount": []string{h.String()},
	}
	var from []string
	for _, m := range w.options.MountPaths {
		from = append(from, m.RepositoryStr())
	}
	// We currently avoid HEAD because it's semi-redundant with the mount that is part
	// of initiating the blob upload.  GCR will perform an existence check on the initiation
	// if "mount" is specified, even if no "from" sources are specified.  If this turns out
	// to not be broadly applicable then we should replace mounts without "from"s with a HEAD.
	if len(from) > 0 {
		uv["from"] = from
	}
	u.RawQuery = uv.Encode()

	// Make the request to initiate the blob upload.
	resp, err := w.client.Post(u.String(), "application/json", nil)
	if err != nil {
		return "", false, err
	}
	defer resp.Body.Close()

	if err := checkError(resp, http.StatusCreated, http.StatusAccepted); err != nil {
		return "", false, err
	}

	// Check the response code to determine the result.
	switch resp.StatusCode {
	case http.StatusCreated:
		// We're done, we were able to fast-path.
		return "", true, nil
	case http.StatusAccepted:
		// Proceed to PATCH, upload has begun.
		loc, err := w.nextLocation(resp)
		return loc, false, err
	default:
		panic("Unreachable: initiateUpload")
	}
}

// streamBlob streams the contents of the blob to the specified location.
// On failure, this will return an error.  On success, this will return the location
// header indicating how to commit the streamed blob.
func (w *writer) streamBlob(h v1.Hash, streamLocation string) (commitLocation string, err error) {
	l, err := w.img.LayerByDigest(h)
	if err != nil {
		return "", err
	}
	blob, err := l.Compressed()
	if err != nil {
		return "", err
	}
	defer blob.Close()

	req, err := http.NewRequest(http.MethodPatch, streamLocation, blob)
	if err != nil {
		return "", err
	}

	resp, err := w.client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if err := checkError(resp, http.StatusNoContent, http.StatusAccepted, http.StatusCreated); err != nil {
		return "", err
	}

	// The blob has been uploaded, return the location header indicating
	// how to commit this layer.
	return w.nextLocation(resp)
}

// commitBlob commits this blob by sending a PUT to the location returned from streaming the blob.
func (w *writer) commitBlob(h v1.Hash, location string) (err error) {
	u, err := url.Parse(location)
	if err != nil {
		return err
	}
	v := u.Query()
	v.Set("digest", h.String())
	u.RawQuery = v.Encode()

	req, err := http.NewRequest(http.MethodPut, u.String(), nil)
	if err != nil {
		return err
	}

	resp, err := w.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	return checkError(resp, http.StatusCreated)
}

// uploadOne performs a complete upload of a single layer.
func (w *writer) uploadOne(h v1.Hash) error {
	location, mounted, err := w.initiateUpload(h)
	if err != nil {
		return err
	} else if mounted {
		log.Printf("mounted blob: %v", h)
		return nil
	}

	location, err = w.streamBlob(h, location)
	if err != nil {
		return err
	}

	if err := w.commitBlob(h, location); err != nil {
		return err
	}
	log.Printf("pushed blob %v", h)
	return nil
}

// commitImage does a PUT of the image's manifest.
func (w *writer) commitImage() error {
	raw, err := w.img.RawManifest()
	if err != nil {
		return err
	}
	mt, err := w.img.MediaType()
	if err != nil {
		return err
	}

	u := w.url(fmt.Sprintf("/v2/%s/manifests/%s", w.ref.Context().RepositoryStr(), w.ref.Identifier()))

	// Make the request to PUT the serialized manifest
	req, err := http.NewRequest(http.MethodPut, u.String(), bytes.NewBuffer(raw))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", string(mt))

	resp, err := w.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if err := checkError(resp, http.StatusOK, http.StatusCreated, http.StatusAccepted); err != nil {
		return err
	}

	digest, err := w.img.Digest()
	if err != nil {
		return err
	}

	// The image was successfully pushed!
	log.Printf("%v: digest: %v size: %d\n", w.ref, digest, len(raw))
	return nil
}

// TODO(mattmoor): WriteIndex
<|endoftext|>


<commit_before>package PlatformProvider

import (
	"BabelProxy/DataShare"
	"BabelProxy/Protocol"
	"BabelProxy/Utils"
	ejson "encoding/json"
	"encoding/xml"
	"errors"
	_ "fmt"
	json "github.com/bitly/go-simplejson"
	_ "github.com/robfig/cron"
	"github.com/spf13/viper"
	_ "io/ioutil"
	"launchpad.net/xmlpath"
	"net/http"
	"strconv"
	"time"
)

type TextResponse struct {
	XMLName      xml.Name `xml:"xml"`
	ToUserName   string   `xml:"ToUserName"`
	FromUserName string   `xml:"FromUserName"`
	CreateTime   int64    `xml:"CreateTime"`
	MsgType      string   `xml:"MsgType"`
	Content      string   `xml:"Content"`
}

type WechatPlatformProvider struct {
	name string
	meta map[string]string
}

func (wPP *WechatPlatformProvider) GetName() string {
	return wPP.name
}

func (wPP *WechatPlatformProvider) GetMeta() map[string]string {
	return wPP.meta
}

func (wPP *WechatPlatformProvider) ReConfigure(f string) (bool, error) {
	return true, nil
}

func (wPP *WechatPlatformProvider) SendMsg(msg *Protocol.Message) bool {
	switch msg.GetMsgType() {
	case "text":
		tR := &TextResponse{ToUserName: msg.GetSender(), FromUserName: wPP.GetMeta()["account"], CreateTime: time.Now().Unix(), MsgType: "text", Content: msg.GetMsgBody()}
		Utils.Logger.Println(tR)
		xmlstring, _ := xml.Marshal(tR)
		jsonString, _ := ejson.Marshal(tR)
		Utils.Logger.Println(string(xmlstring[:]))
		Utils.Logger.Println(string(jsonString[:]))
		return true
	case "image":
		tR := &TextResponse{ToUserName: msg.GetSender(), FromUserName: wPP.GetMeta()["account"], CreateTime: time.Now().Unix(), MsgType: "text", Content: msg.GetMsgBody()}
		xmlstring, _ := xml.Marshal(tR)
		Utils.Logger.Println(xmlstring)
		return true
	case "voice":
		tR := &TextResponse{ToUserName: msg.GetSender(), FromUserName: wPP.GetMeta()["account"], CreateTime: time.Now().Unix(), MsgType: "text", Content: msg.GetMsgBody()}
		xmlstring, _ := xml.Marshal(tR)
		Utils.Logger.Println(xmlstring)
		return true
	case "video":
		tR := &TextResponse{ToUserName: msg.GetSender(), FromUserName: wPP.GetMeta()["account"], CreateTime: time.Now().Unix(), MsgType: "text", Content: msg.GetMsgBody()}
		xmlstring, _ := xml.Marshal(tR)
		Utils.Logger.Println(xmlstring)
		return true
	case "news":
		tR := &TextResponse{ToUserName: msg.GetSender(), FromUserName: wPP.GetMeta()["account"], CreateTime: time.Now().Unix(), MsgType: "text", Content: msg.GetMsgBody()}
		xmlstring, _ := xml.Marshal(tR)
		Utils.Logger.Println(xmlstring)
		return true
	case "resp":
		tR := &TextResponse{ToUserName: msg.GetSender(), FromUserName: wPP.GetMeta()["account"], CreateTime: time.Now().Unix(), MsgType: "text", Content: msg.GetMsgBody()}
		xmlstring, _ := xml.Marshal(tR)
		Utils.Logger.Println(xmlstring)
		return true
	default:
		tR := &TextResponse{ToUserName: msg.GetSender(), FromUserName: wPP.GetMeta()["account"], CreateTime: time.Now().Unix(), MsgType: "text", Content: msg.GetMsgBody()}
		xmlstring, _ := xml.Marshal(tR)
		Utils.Logger.Println(xmlstring)
		return true

	}
}

func (wPP *WechatPlatformProvider) UpdateToken() {
	for {
		Utils.Logger.Println("Update Access Token For WeChat")
		url := "https://api.weixin.qq.com/cgi-bin/token?grant_type=client_credential"
		updateUrl := url + "&appid=" + wPP.GetMeta()["appId"] + "&secret=" + wPP.GetMeta()["appsecret"]
		Utils.Logger.Println(updateUrl)
		resp, err := http.Get(updateUrl)
		if err == nil {
			js, _ := json.NewFromReader(resp.Body)
			value, flag := js.CheckGet("access_token")
			if flag {
				wPP.meta["access_token"] = value.MustString()
				Utils.Logger.Println("Update Access Token to ", wPP.meta["access_token"])
				time.Sleep(3600 * time.Second)
			}
		} else {
			Utils.Logger.Println("Get Token Failed", err)
			time.Sleep(3600 * time.Second)

		}
	}

}

func (wPP *WechatPlatformProvider) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	content, err := xmlpath.Parse(r.Body)
	if err != nil {
		Utils.Logger.Println("Cann't Parse Message from ", r.URL)
	} else {
		senderPath := xmlpath.MustCompile("/xml/FromUserName")
		createTimePath := xmlpath.MustCompile("/xml/CreateTime")
		sender, _ := senderPath.String(content)
		createTimeStampStr, _ := createTimePath.String(content)
		createTimeStamp, _ := strconv.Atoi(createTimeStampStr)
		createTime := time.Unix(int64(createTimeStamp), 0)
		msgTypePath := xmlpath.MustCompile("/xml/MsgType")
		switch msgType, _ := msgTypePath.String(content); msgType {
		case "text":
			Utils.Logger.Println("Get a Text From ", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/Content")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "text", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			DataShare.MsgQ <- newMsg
			return
		case "image":
			Utils.Logger.Println("Get a Image From ", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/MediaId")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			picUrlPath := xmlpath.MustCompile("/xml/PicUrl")
			picUrl, _ := picUrlPath.String(content)
			msgMeta["PicUrl"] = picUrl
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "image", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		case "voice":
			Utils.Logger.Println("Get a Voice Msg From", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/MediaId")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			formatPath := xmlpath.MustCompile("/xml/Format")
			format, _ := formatPath.String(content)
			msgMeta["Format"] = format
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "voice", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		case "video":
			Utils.Logger.Println("Get a Video Msg From", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/MediaId")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			thumbMediaIdPath := xmlpath.MustCompile("/xml/ThumbMediaId")
			thumbMediaId, _ := thumbMediaIdPath.String(content)
			msgMeta["ThumbMediaId"] = thumbMediaId
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "voice", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		case "shortvideo":
			Utils.Logger.Println("Get a Short Video Msg From", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/MediaId")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			thumbMediaIdPath := xmlpath.MustCompile("/xml/ThumbMediaId")
			thumbMediaId, _ := thumbMediaIdPath.String(content)
			msgMeta["ThumbMediaId"] = thumbMediaId
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "voice", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		case "location":
			Utils.Logger.Println("Get a Location Msg From", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/Label")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			Location_XPath := xmlpath.MustCompile("/xml/Location_X")
			Location_X, _ := Location_XPath.String(content)
			msgMeta["Location_X"] = Location_X
			Location_YPath := xmlpath.MustCompile("/xml/Location_Y")
			Location_Y, _ := Location_YPath.String(content)
			msgMeta["Location_Y"] = Location_Y
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "Location", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		default:
			Utils.Logger.Println("default choice")
		}
	}
}

func CreateWechatPlatformProvider(f string) (*WechatPlatformProvider, error) {
	Utils.Logger.Println("Start Creating WeChatPlatformProvider")
	viper.SetConfigFile(f)
	err := viper.ReadInConfig()
	if err != nil {
		Utils.Logger.Println("Cannot Load Wechat Configure File at " + f)
		return &WechatPlatformProvider{}, errors.New("Load Configuration Failed Error")
	}
	var wPP = &WechatPlatformProvider{}
	wPP.name = viper.GetString("name")
	wPP.meta = make(map[string]string)
	wPP.meta["account"] = viper.GetString("account")
	wPP.meta["appId"] = viper.GetString("appId")
	wPP.meta["appsecret"] = viper.GetString("appsecret")
	wPP.meta["url"] = viper.GetString("url")
	//	Utils.Logger.Println("Finish Creating WeChatPlatformProvider")
	//	c := cron.New()
	//	c.AddFunc("@every 1h30m", func() { wPP.UpdateToken() })
	//	c.Start()
	go wPP.UpdateToken()
	return wPP, nil
}
<commit_msg>add sending support<commit_after>package PlatformProvider

import (
	"BabelProxy/DataShare"
	"BabelProxy/Protocol"
	"BabelProxy/Utils"
	ejson "encoding/json"
	_"encoding/xml"
	"errors"
	_ "fmt"
	json "github.com/bitly/go-simplejson"
	_ "github.com/robfig/cron"
	"github.com/spf13/viper"
	_ "io/ioutil"
	"launchpad.net/xmlpath"
	"net/http"
	"strconv"
	"time"
	"bytes"
)

//type TextResponse struct {
//	XMLName      xml.Name `xml:"xml"`
//	ToUserName   string   `xml:"ToUserName"`
//	FromUserName string   `xml:"FromUserName"`
//	CreateTime   int64    `xml:"CreateTime"`
//	MsgType      string   `xml:"MsgType"`
//	Content      string   `xml:"Content"`
//}

type TextFrame struct{
	Content string `json:"content"`
}

type VoiceFrame struct{
	MediaId string   `json:"media_id"`
}

type VideoFrame struct{
	MediaId string `json:"media_id"`
	ThumbMediaI string `json:"thumb_media_id"`
	Title string `json:"title"`
	Description string `json:"description"`
}

type TextResponse struct{
		ToUser string `json:"touser"`
		MsgType string `json:"msgtype"`
		Video TextFrame `json:"text"`

}

type VoiceResponse struct{
		ToUser string `json:"touser"`
		MsgType string `json:"msgtype"`
		Voice VoiceFrame `json:"voice"`

}

type VideoResponse struct{
		ToUser string `json:"touser"`
		MsgType string `json:"msgtype"`
		Video VideoFrame `json:"video"`

}

type WechatPlatformProvider struct {
	name string
	meta map[string]string
}

func (wPP *WechatPlatformProvider) GetName() string {
	return wPP.name
}

func (wPP *WechatPlatformProvider) GetMeta() map[string]string {
	return wPP.meta
}

func (wPP *WechatPlatformProvider) ReConfigure(f string) (bool, error) {
	return true, nil
}

func (wPP *WechatPlatformProvider) SendMsg(msg *Protocol.Message) bool {
	url := "https://api.weixin.qq.com/cgi-bin/message/custom/send?access_token="+wPP.GetMeta()["access_token"]
	switch msg.GetMsgType() {
	case "text":
		tR := &TextResponse{msg.GetSender(),"text",TextFrame{msg.GetMsgBody()}}
		Utils.Logger.Println(tR)
		jsonString, _ := ejson.Marshal(tR)
		Utils.Logger.Println("Json to Sent ",string(jsonString[:]))
		resp,err := http.Post(url,"application/json",bytes.NewReader(jsonString))
		if err != nil{
			Utils.Logger.Println("Failed to Sent Package")

		}
		Utils.Logger.Println(resp.Body)

		return true
	case "image":
		tR := &TextResponse{msg.GetSender(),"text",TextFrame{msg.GetMsgBody()}}
		Utils.Logger.Println(tR)
		return true
	case "voice":
		tR := &TextResponse{msg.GetSender(),"text",TextFrame{msg.GetMsgBody()}}
		Utils.Logger.Println(tR)
		return true
	case "video":
		tR := &TextResponse{msg.GetSender(),"text",TextFrame{msg.GetMsgBody()}}
		Utils.Logger.Println(tR)
		return true
	default:
//		tR := &TextResponse{ToUserName: msg.GetSender(), FromUserName: wPP.GetMeta()["account"], CreateTime: time.Now().Unix(), MsgType: "text", Content: msg.GetMsgBody()}
		return true

	}
}

func (wPP *WechatPlatformProvider) UpdateToken() {
	for {
		Utils.Logger.Println("Update Access Token For WeChat")
		url := "https://api.weixin.qq.com/cgi-bin/token?grant_type=client_credential"
		updateUrl := url + "&appid=" + wPP.GetMeta()["appId"] + "&secret=" + wPP.GetMeta()["appsecret"]
		Utils.Logger.Println(updateUrl)
		resp, err := http.Get(updateUrl)
		if err == nil {
			js, _ := json.NewFromReader(resp.Body)
			value, flag := js.CheckGet("access_token")
			if flag {
				wPP.meta["access_token"] = value.MustString()
				Utils.Logger.Println("Update Access Token to ", wPP.meta["access_token"])
				time.Sleep(3600 * time.Second)
			}
		} else {
			Utils.Logger.Println("Get Token Failed", err)
			time.Sleep(3600 * time.Second)

		}
	}

}

func (wPP *WechatPlatformProvider) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	content, err := xmlpath.Parse(r.Body)
	if err != nil {
		Utils.Logger.Println("Cann't Parse Message from ", r.URL)
	} else {
		senderPath := xmlpath.MustCompile("/xml/FromUserName")
		createTimePath := xmlpath.MustCompile("/xml/CreateTime")
		sender, _ := senderPath.String(content)
		createTimeStampStr, _ := createTimePath.String(content)
		createTimeStamp, _ := strconv.Atoi(createTimeStampStr)
		createTime := time.Unix(int64(createTimeStamp), 0)
		msgTypePath := xmlpath.MustCompile("/xml/MsgType")
		switch msgType, _ := msgTypePath.String(content); msgType {
		case "text":
			Utils.Logger.Println("Get a Text From ", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/Content")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "text", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			DataShare.MsgQ <- newMsg
			return
		case "image":
			Utils.Logger.Println("Get a Image From ", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/MediaId")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			picUrlPath := xmlpath.MustCompile("/xml/PicUrl")
			picUrl, _ := picUrlPath.String(content)
			msgMeta["PicUrl"] = picUrl
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "image", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		case "voice":
			Utils.Logger.Println("Get a Voice Msg From", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/MediaId")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			formatPath := xmlpath.MustCompile("/xml/Format")
			format, _ := formatPath.String(content)
			msgMeta["Format"] = format
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "voice", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		case "video":
			Utils.Logger.Println("Get a Video Msg From", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/MediaId")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			thumbMediaIdPath := xmlpath.MustCompile("/xml/ThumbMediaId")
			thumbMediaId, _ := thumbMediaIdPath.String(content)
			msgMeta["ThumbMediaId"] = thumbMediaId
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "voice", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		case "shortvideo":
			Utils.Logger.Println("Get a Short Video Msg From", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/MediaId")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			thumbMediaIdPath := xmlpath.MustCompile("/xml/ThumbMediaId")
			thumbMediaId, _ := thumbMediaIdPath.String(content)
			msgMeta["ThumbMediaId"] = thumbMediaId
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "voice", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		case "location":
			Utils.Logger.Println("Get a Location Msg From", r.URL)
			msgBodyPath := xmlpath.MustCompile("/xml/Label")
			msgBody, _ := msgBodyPath.String(content)
			msgMeta := make(map[string]string)
			Location_XPath := xmlpath.MustCompile("/xml/Location_X")
			Location_X, _ := Location_XPath.String(content)
			msgMeta["Location_X"] = Location_X
			Location_YPath := xmlpath.MustCompile("/xml/Location_Y")
			Location_Y, _ := Location_YPath.String(content)
			msgMeta["Location_Y"] = Location_Y
			newMsg := Protocol.CreateMsg(msgBody, sender, wPP.GetName(), "Location", createTime, msgMeta)
			Utils.Logger.Println(newMsg.GetMsgBody())
			return
		default:
			Utils.Logger.Println("default choice")
		}
	}
}

func CreateWechatPlatformProvider(f string) (*WechatPlatformProvider, error) {
	Utils.Logger.Println("Start Creating WeChatPlatformProvider")
	viper.SetConfigFile(f)
	err := viper.ReadInConfig()
	if err != nil {
		Utils.Logger.Println("Cannot Load Wechat Configure File at " + f)
		return &WechatPlatformProvider{}, errors.New("Load Configuration Failed Error")
	}
	var wPP = &WechatPlatformProvider{}
	wPP.name = viper.GetString("name")
	wPP.meta = make(map[string]string)
	wPP.meta["account"] = viper.GetString("account")
	wPP.meta["appId"] = viper.GetString("appId")
	wPP.meta["appsecret"] = viper.GetString("appsecret")
	wPP.meta["url"] = viper.GetString("url")
	//	Utils.Logger.Println("Finish Creating WeChatPlatformProvider")
	//	c := cron.New()
	//	c.AddFunc("@every 1h30m", func() { wPP.UpdateToken() })
	//	c.Start()
	go wPP.UpdateToken()
	return wPP, nil
}
<|endoftext|>