<commit_before>package metrics

import (
	"fmt"
	"time"

	"github.com/fweikert/continuous-integration/metrics/clients"
	"github.com/fweikert/continuous-integration/metrics/data"
	timestamp "github.com/golang/protobuf/ptypes/timestamp"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	monitoredres "google.golang.org/genproto/googleapis/api/monitoredres"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

const baseMetricType = "custom.googleapis.com/bazel/ci"

type PlatformLoad struct {
	client  *clients.BuildkiteClient
	columns []Column
	builds  int
}

func (pl *PlatformLoad) Name() string {
	return "platform_load"
}

func (pl *PlatformLoad) Columns() []Column {
	return pl.columns
}

func (pl *PlatformLoad) Collect() (data.DataSet, error) {
	builds, err := pl.client.GetMostRecentBuilds("all", pl.builds)
	if err != nil {
		return nil, fmt.Errorf("Cannot get builds to determine platform load: %v", err)
	}

	result := &loadDataSet{headers: GetColumnNames(pl.columns), ts: time.Now()}
	allPlatforms := make(map[string]bool)
	waiting := make(map[string]int)
	running := make(map[string]int)
	for _, build := range builds {
		for _, job := range build.Jobs {
			// Do not use getPlatform() since it may return "rbe", but here we're only interested in the actual worker OS (which would be "linux" in the rbe case).
			platform := getPlatformFromAgentQueryRules(job.AgentQueryRules)
			if platform == "" || job.CreatedAt == nil || job.FinishedAt != nil {
				continue
			}
			allPlatforms[platform] = true
			switch *job.State {
			case "running":
				running[platform] += 1
			case "scheduled":
				/*
					State "scheduled" = waiting for a worker to become available
					State "waiting" / "waiting_failed" = waiting for another task to finish

					We're only interested in "scheduled" jobs since they may indicate a shortage of workers.
				*/
				waiting[platform] += 1
			}
		}
	}

	result.rows = make([]loadDataRow, 0)
	for platform := range allPlatforms {
		row := loadDataRow{platform: platform, waitingJobs: waiting[platform], runningJobs: running[platform]}
		result.rows = append(result.rows, row)
	}
	return result, nil
}

// CREATE TABLE platform_load (timestamp DATETIME, platform VARCHAR(255), waiting_jobs INT, running_jobs INT, PRIMARY KEY(timestamp, platform));
func CreatePlatformLoad(client *clients.BuildkiteClient, builds int) *PlatformLoad {
	columns := []Column{Column{"timestamp", true}, Column{"platform", true}, Column{"waiting_jobs", false}, Column{"running_jobs", false}}
	return &PlatformLoad{client: client, columns: columns, builds: builds}
}

type loadDataRow struct {
	platform    string
	waitingJobs int
	runningJobs int
}

type loadDataSet struct {
	headers []string
	ts      time.Time
	rows    []loadDataRow
}

func (lds *loadDataSet) GetData() *data.LegacyDataSet {
	rawSet := data.CreateDataSet(lds.headers)
	for _, row := range lds.rows {
		rawRow := []interface{}{lds.ts, row.platform, row.waitingJobs, row.runningJobs}
		rawSet.Data = append(rawSet.Data, rawRow)
	}
	return rawSet
}

func (lds *loadDataSet) CreateTimeSeriesRequest(projectID string) *monitoringpb.CreateTimeSeriesRequest {
	ts := &timestamp.Timestamp{
		Seconds: lds.ts.Unix(),
	}
	series := make([]*monitoringpb.TimeSeries, len(lds.rows)*2)
	for i, row := range lds.rows {
		series[i] = createTimeSeries(ts, row.platform, "waiting_jobs", row.waitingJobs)
		series[i] = createTimeSeries(ts, row.platform, "running_jobs", row.runningJobs)
	}

	return &monitoringpb.CreateTimeSeriesRequest{
		Name:       "projects/" + projectID,
		TimeSeries: series,
	}
}

func createTimeSeries(ts *timestamp.Timestamp, platform, metricSuffix string, value int) *monitoringpb.TimeSeries {
	return &monitoringpb.TimeSeries{
		Metric: &metricpb.Metric{
			Type: fmt.Sprintf("%s/%s/%s", baseMetricType, platform, metricSuffix),
		},
		Resource: &monitoredres.MonitoredResource{
			Type: "global",
		},
		Points: []*monitoringpb.Point{{
			Interval: &monitoringpb.TimeInterval{
				StartTime: ts,
				EndTime:   ts,
			},
			Value: &monitoringpb.TypedValue{
				Value: &monitoringpb.TypedValue_Int64Value{
					Int64Value: int64(value),
				},
			},
		}},
	}
}
<commit_msg>Fix list indexes of time series.<commit_after>package metrics

import (
	"fmt"
	"time"

	"github.com/fweikert/continuous-integration/metrics/clients"
	"github.com/fweikert/continuous-integration/metrics/data"
	timestamp "github.com/golang/protobuf/ptypes/timestamp"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	monitoredres "google.golang.org/genproto/googleapis/api/monitoredres"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

const baseMetricType = "custom.googleapis.com/bazel/ci"

type PlatformLoad struct {
	client  *clients.BuildkiteClient
	columns []Column
	builds  int
}

func (pl *PlatformLoad) Name() string {
	return "platform_load"
}

func (pl *PlatformLoad) Columns() []Column {
	return pl.columns
}

func (pl *PlatformLoad) Collect() (data.DataSet, error) {
	builds, err := pl.client.GetMostRecentBuilds("all", pl.builds)
	if err != nil {
		return nil, fmt.Errorf("Cannot get builds to determine platform load: %v", err)
	}

	result := &loadDataSet{headers: GetColumnNames(pl.columns), ts: time.Now()}
	allPlatforms := make(map[string]bool)
	waiting := make(map[string]int)
	running := make(map[string]int)
	for _, build := range builds {
		for _, job := range build.Jobs {
			// Do not use getPlatform() since it may return "rbe", but here we're only interested in the actual worker OS (which would be "linux" in the rbe case).
			platform := getPlatformFromAgentQueryRules(job.AgentQueryRules)
			if platform == "" || job.CreatedAt == nil || job.FinishedAt != nil {
				continue
			}
			allPlatforms[platform] = true
			switch *job.State {
			case "running":
				running[platform] += 1
			case "scheduled":
				/*
					State "scheduled" = waiting for a worker to become available
					State "waiting" / "waiting_failed" = waiting for another task to finish

					We're only interested in "scheduled" jobs since they may indicate a shortage of workers.
				*/
				waiting[platform] += 1
			}
		}
	}

	result.rows = make([]loadDataRow, 0)
	for platform := range allPlatforms {
		row := loadDataRow{platform: platform, waitingJobs: waiting[platform], runningJobs: running[platform]}
		result.rows = append(result.rows, row)
	}
	return result, nil
}

// CREATE TABLE platform_load (timestamp DATETIME, platform VARCHAR(255), waiting_jobs INT, running_jobs INT, PRIMARY KEY(timestamp, platform));
func CreatePlatformLoad(client *clients.BuildkiteClient, builds int) *PlatformLoad {
	columns := []Column{Column{"timestamp", true}, Column{"platform", true}, Column{"waiting_jobs", false}, Column{"running_jobs", false}}
	return &PlatformLoad{client: client, columns: columns, builds: builds}
}

type loadDataRow struct {
	platform    string
	waitingJobs int
	runningJobs int
}

type loadDataSet struct {
	headers []string
	ts      time.Time
	rows    []loadDataRow
}

func (lds *loadDataSet) GetData() *data.LegacyDataSet {
	rawSet := data.CreateDataSet(lds.headers)
	for _, row := range lds.rows {
		rawRow := []interface{}{lds.ts, row.platform, row.waitingJobs, row.runningJobs}
		rawSet.Data = append(rawSet.Data, rawRow)
	}
	return rawSet
}

func (lds *loadDataSet) CreateTimeSeriesRequest(projectID string) *monitoringpb.CreateTimeSeriesRequest {
	ts := &timestamp.Timestamp{
		Seconds: lds.ts.Unix(),
	}
	series := make([]*monitoringpb.TimeSeries, len(lds.rows)*2)
	for i, row := range lds.rows {
		series[2*i] = createTimeSeries(ts, row.platform, "waiting_jobs", row.waitingJobs)
		series[2*i+1] = createTimeSeries(ts, row.platform, "running_jobs", row.runningJobs)
	}

	return &monitoringpb.CreateTimeSeriesRequest{
		Name:       "projects/" + projectID,
		TimeSeries: series,
	}
}

func createTimeSeries(ts *timestamp.Timestamp, platform, metricSuffix string, value int) *monitoringpb.TimeSeries {
	return &monitoringpb.TimeSeries{
		Metric: &metricpb.Metric{
			Type: fmt.Sprintf("%s/%s/%s", baseMetricType, platform, metricSuffix),
		},
		Resource: &monitoredres.MonitoredResource{
			Type: "global",
		},
		Points: []*monitoringpb.Point{{
			Interval: &monitoringpb.TimeInterval{
				StartTime: ts,
				EndTime:   ts,
			},
			Value: &monitoringpb.TypedValue{
				Value: &monitoringpb.TypedValue_Int64Value{
					Int64Value: int64(value),
				},
			},
		}},
	}
}
<|endoftext|>


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

// HTTP reverse proxy handler

package proxy

import (
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Director must be a function which modifies
	// the request into a new request to be sent
	// using Transport. Its response is then copied
	// back to the original client unmodified.
	Director func(*http.Request)

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

// NewSingleHostReverseProxy returns a new ReverseProxy that rewrites
// URLs to the scheme, host, and base path provided in target. If the
// target's path is "/base" and the incoming request was for "/dir",
// the target request will be for /base/dir.
func NewSingleHostReverseProxy(target *url.URL) *ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}
	return &ReverseProxy{Director: director}
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request, extraHeaders http.Header) error {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	p.Director(outreq)
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	if extraHeaders != nil {
		for k, v := range extraHeaders {
			outreq.Header[k] = v
		}
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode == http.StatusSwitchingProtocols && outreq.Header.Get("Upgrade") == "websocket" {
		hj, ok := rw.(http.Hijacker)
		if !ok {
			return nil
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			return err
		}

		backendConn, err := net.Dial("tcp", outreq.Host)
		if err != nil {
			conn.Close()
			return err
		}

		outreq.Write(backendConn)

		go func() {
			io.Copy(backendConn, conn) // write tcp stream to backend.
			backendConn.Close()
		}()

		io.Copy(conn, backendConn) // read tcp stream from backend.
		conn.Close()
	} else {
		for _, h := range hopHeaders {
			res.Header.Del(h)
		}

		copyHeader(rw.Header(), res.Header)

		rw.WriteHeader(res.StatusCode)
		p.copyResponse(rw, res.Body)
	}

	return nil
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	io.Copy(dst, src)
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<commit_msg>check server response instead of client<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// HTTP reverse proxy handler

package proxy

import (
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Director must be a function which modifies
	// the request into a new request to be sent
	// using Transport. Its response is then copied
	// back to the original client unmodified.
	Director func(*http.Request)

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

// NewSingleHostReverseProxy returns a new ReverseProxy that rewrites
// URLs to the scheme, host, and base path provided in target. If the
// target's path is "/base" and the incoming request was for "/dir",
// the target request will be for /base/dir.
func NewSingleHostReverseProxy(target *url.URL) *ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}
	return &ReverseProxy{Director: director}
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request, extraHeaders http.Header) error {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	p.Director(outreq)
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	if extraHeaders != nil {
		for k, v := range extraHeaders {
			outreq.Header[k] = v
		}
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode == http.StatusSwitchingProtocols && res.Header.Get("Upgrade") == "websocket" {
		hj, ok := rw.(http.Hijacker)
		if !ok {
			return nil
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			return err
		}

		backendConn, err := net.Dial("tcp", outreq.Host)
		if err != nil {
			conn.Close()
			return err
		}

		outreq.Write(backendConn)

		go func() {
			io.Copy(backendConn, conn) // write tcp stream to backend.
			backendConn.Close()
		}()

		io.Copy(conn, backendConn) // read tcp stream from backend.
		conn.Close()
	} else {
		for _, h := range hopHeaders {
			res.Header.Del(h)
		}

		copyHeader(rw.Header(), res.Header)

		rw.WriteHeader(res.StatusCode)
		p.copyResponse(rw, res.Body)
	}

	return nil
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	io.Copy(dst, src)
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<|endoftext|>


<commit_before>package main

import (
	"eaciit/gdrj/model"
	"eaciit/gdrj/modules"
	"os"

	"flag"
	"github.com/eaciit/dbox"
	"github.com/eaciit/orm/v1"
	"github.com/eaciit/toolkit"
	"strings"
	"time"
)

var conn dbox.IConnection
var count int

// var mwg sync.WaitGroup

var (
	t0                                time.Time
	custgroup, prodgroup, plcode, ref string
	value, fiscalyear                 int
	globalval                         float64
	mapkeysvalue                      map[string]float64
	masters                           toolkit.M
)

func setinitialconnection() {
	var err error
	conn, err = modules.GetDboxIConnection("db_godrej")

	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}

	err = gdrj.SetDb(conn)
	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}
}

func buildmap(holder interface{},
	fnModel func() orm.IModel,
	filter *dbox.Filter,
	fnIter func(holder interface{}, obj interface{})) interface{} {
	crx, _ := gdrj.Find(fnModel(), filter, nil)
	defer crx.Close()
	for {
		s := fnModel()
		e := crx.Fetch(s, 1, false)
		if e != nil {
			break
		}
		fnIter(holder, s)
	}
	return holder
}

func prepmaster() {
	masters = toolkit.M{}
	masters.Set("plmodel", buildmap(map[string]*gdrj.PLModel{},
		func() orm.IModel {
			return new(gdrj.PLModel)
		},
		nil,
		func(holder, obj interface{}) {
			h := holder.(map[string]*gdrj.PLModel)
			o := obj.(*gdrj.PLModel)
			h[o.ID] = o
		}).(map[string]*gdrj.PLModel))
}

func main() {
	t0 = time.Now()
	mapkeysvalue = make(map[string]float64)

	flag.IntVar(&value, "value", 0, "representation of value need to be allocation. Default is 0")
	//branch,channel,group
	flag.StringVar(&custgroup, "custgroup", "global", "group of customer. Default is global")
	//brand,group,skuid
	flag.StringVar(&prodgroup, "prodgroup", "global", "group of product. Default is global")
	flag.IntVar(&fiscalyear, "year", 2015, "YYYY representation of godrej fiscal year. Default is 2015")
	flag.StringVar(&plcode, "plcode", "", "PL Code to take the value. Default is blank")
	flag.StringVar(&ref, "ref", "", "Reference from document or other. Default is blank")
	flag.Parse()

	eperiode := time.Date(fiscalyear, 4, 1, 0, 0, 0, 0, time.UTC)
	speriode := eperiode.AddDate(-1, 0, 0)

	filter := dbox.And(dbox.Gte("date.date", speriode), dbox.Lt("date.date", eperiode))

	if plcode == "" || value == 0 {
		toolkit.Println("PLCode and Value are mandatory to fill")
		os.Exit(1)
	}

	setinitialconnection()
	defer gdrj.CloseDb()

	toolkit.Println("Get Data Master...")
	prepmaster()

	toolkit.Println("Start Data Process...")
	c, _ := gdrj.Find(new(gdrj.SalesPL), filter, nil)
	defer c.Close()

	splcount := c.Count()
	step := splcount / 100
	i := 0

	for {
		i++

		spl := new(gdrj.SalesPL)
		e := c.Fetch(spl, 1, false)
		if e != nil {
			break
		}
		key := toolkit.Sprintf("%d_%d", spl.Date.Month, spl.Date.Year)
		k := ""

		if spl.Customer != nil {
			switch custgroup {
			case "branch":
				key = toolkit.Sprintf("%v_%v", key, spl.Customer.BranchID)
			case "channel":
				k = toolkit.Sprintf("%v|%v", spl.Customer.BranchID, spl.Customer.ChannelID)
				key = toolkit.Sprintf("%v_%v", key, k)
			case "group":
				k = toolkit.Sprintf("%v|%v", spl.Customer.BranchID, spl.Customer.CustomerGroup)
				key = toolkit.Sprintf("%v_%v", key, k)
			default:
				key = toolkit.Sprintf("%v_", key)
			}
		}

		if spl.Product != nil {
			switch prodgroup {
			case "brand":
				key = toolkit.Sprintf("%v_%v", key, spl.Product.Brand)
			case "group":
				k = toolkit.Sprintf("%v|%v", spl.Product.Brand, spl.Product.BrandCategoryID)
				key = toolkit.Sprintf("%v_%v", key, k)
			case "skuid":
				k = toolkit.Sprintf("%v|%v|%v", spl.Product.Brand, spl.Product.BrandCategoryID, spl.Product.ID)
				key = toolkit.Sprintf("%v_%v", key, k)
			default:
				key = toolkit.Sprintf("%v_", key)
			}

		}

		mapkeysvalue[key] += spl.GrossAmount
		globalval += spl.GrossAmount

		if i > step {
			step += splcount / 100
			toolkit.Printfn("Processing %d of %d in %s", i, splcount,
				time.Since(t0).String())
		}

	}

	toolkit.Printfn("Preparing the data")
	mcount := len(mapkeysvalue)
	i = 0
	for k, v := range mapkeysvalue {
		i++

		akey := strings.Split(k, "_")
		tcustomer := new(gdrj.Customer)
		tproduct := new(gdrj.Product)

		if len(akey) > 2 && akey[2] != "" {
			switch custgroup {
			case "branch":
				tcustomer.BranchID = akey[2]
			case "channel":
				skey := strings.Split(akey[2], "|")
				tcustomer.BranchID = skey[0]
				if len(skey) > 1 {
					tcustomer.ChannelID = skey[1]
				}
			case "group":
				skey := strings.Split(akey[2], "|")
				tcustomer.BranchID = skey[0]
				if len(skey) > 1 {
					tcustomer.CustomerGroup = skey[1]
				}
			}
		}

		if len(akey) > 3 && akey[3] != "" {
			switch prodgroup {
			case "brand":
				tproduct.Brand = akey[3]
			case "group":
				skey := strings.Split(akey[3], "|")
				tproduct.Brand = skey[0]
				if len(skey) > 1 {
					tproduct.BrandCategoryID = skey[1]
				}
			case "skuid":
				skey := strings.Split(akey[3], "|")
				tproduct.Brand = skey[0]
				if len(skey) > 1 {
					tproduct.BrandCategoryID = skey[1]
				}
				if len(skey) > 2 {
					tproduct.ID = skey[2]
				}
			}
		}

		spl := new(gdrj.SalesPL)
		spl.ID = toolkit.Sprintf("%v%v%v%v", akey[0], akey[1], "", "")
		spl.Date = gdrj.SetDate(time.Date(toolkit.ToInt(akey[1], toolkit.RoundingAuto),
			time.Month(toolkit.ToInt(akey[0], toolkit.RoundingAuto)),
			1, 0, 0, 0, 0, time.UTC))
		spl.Customer = tcustomer
		spl.Product = tproduct
		spl.Source = "ADJUST"
		spl.Ref = ref

		spl.PC = new(gdrj.ProfitCenter)
		spl.CC = new(gdrj.CostCenter)

		amount := toolkit.ToFloat64(value, 6, toolkit.RoundingAuto) * (v / globalval)
		plmodels := masters.Get("plmodel").(map[string]*gdrj.PLModel)
		spl.AddData(plcode, amount, plmodels)
		spl.CalcSum(masters)

		gdrj.Save(spl)

		toolkit.Printfn("Saving %d of %d in %s", i, mcount,
			time.Since(t0).String())
	}

	toolkit.Printfn("Processing done in %s",
		time.Since(t0).String())
}
<commit_msg>update<commit_after>package main

import (
	"eaciit/gdrj/model"
	"eaciit/gdrj/modules"
	"os"

	"flag"
	"github.com/eaciit/dbox"
	"github.com/eaciit/orm/v1"
	"github.com/eaciit/toolkit"
	"strings"
	"time"
)

var conn dbox.IConnection
var count int

// var mwg sync.WaitGroup

var (
	t0                                time.Time
	custgroup, prodgroup, plcode, ref string
	value, fiscalyear                 int
	globalval                         float64
	mapkeysvalue                      map[string]float64
	masters                           toolkit.M
)

func setinitialconnection() {
	var err error
	conn, err = modules.GetDboxIConnection("db_godrej")

	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}

	err = gdrj.SetDb(conn)
	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}
}

func getCursor(obj orm.IModel) dbox.ICursor {
	c, e := gdrj.Find(obj,
		nil, nil)
	//toolkit.M{}.Set("take", 10))
	if e != nil {
		return nil
	}
	return c
}

func buildmap(holder interface{},
	fnModel func() orm.IModel,
	filter *dbox.Filter,
	fnIter func(holder interface{}, obj interface{})) interface{} {
	crx, _ := gdrj.Find(fnModel(), filter, nil)
	defer crx.Close()
	for {
		s := fnModel()
		e := crx.Fetch(s, 1, false)
		if e != nil {
			break
		}
		fnIter(holder, s)
	}
	return holder
}

func prepmaster() {
	masters = toolkit.M{}
	masters.Set("plmodel", buildmap(map[string]*gdrj.PLModel{},
		func() orm.IModel {
			return new(gdrj.PLModel)
		},
		nil,
		func(holder, obj interface{}) {
			h := holder.(map[string]*gdrj.PLModel)
			o := obj.(*gdrj.PLModel)
			h[o.ID] = o
		}).(map[string]*gdrj.PLModel))
}

func prepmasterclean() {
	subchannels := toolkit.M{}
	toolkit.Println("--> Sub Channel")
	csr, _ := conn.NewQuery().From("subchannels").Cursor(nil)
	defer csr.Close()
	for {
		m := toolkit.M{}
		e := csr.Fetch(&m, 1, false)
		if e != nil {
			break
		}
		subchannels.Set(m.GetString("_id"), m.GetString("title"))
	}
	masters.Set("subchannels", subchannels)

	customers := toolkit.M{}
	toolkit.Println("--> Customer")
	ccb := getCursor(new(gdrj.Customer))
	defer ccb.Close()
	for {
		cust := new(gdrj.Customer)
		e := ccb.Fetch(cust, 1, false)
		if e != nil {
			break
		}

		customers.Set(cust.ID, cust)
	}
	masters.Set("customers", customers)

	branchs := toolkit.M{}
	cmb := getCursor(new(gdrj.MasterBranch))
	defer cmb.Close()
	for {
		stx := toolkit.M{}
		e := cmb.Fetch(&stx, 1, false)
		if e != nil {
			break
		}

		branchs.Set(stx.Get("_id", "").(string), stx)
	}
	masters.Set("branchs", branchs)

	rdlocations := toolkit.M{}
	crdloc, _ := conn.NewQuery().From("outletgeo").Cursor(nil)
	defer crdloc.Close()
	for {
		stx := toolkit.M{}
		e := crdloc.Fetch(&stx, 1, false)
		if e != nil {
			break
		}

		branchs.Set(stx.GetString("_id"), stx)
	}
	masters.Set("rdlocations", rdlocations)
}

func main() {
	t0 = time.Now()
	mapkeysvalue = make(map[string]float64)

	flag.IntVar(&value, "value", 0, "representation of value need to be allocation. Default is 0")
	//branch,channel,group
	flag.StringVar(&custgroup, "custgroup", "global", "group of customer. Default is global")
	//brand,group,skuid
	flag.StringVar(&prodgroup, "prodgroup", "global", "group of product. Default is global")
	flag.IntVar(&fiscalyear, "year", 2015, "YYYY representation of godrej fiscal year. Default is 2015")
	flag.StringVar(&plcode, "plcode", "", "PL Code to take the value. Default is blank")
	flag.StringVar(&ref, "ref", "", "Reference from document or other. Default is blank")
	flag.Parse()

	eperiode := time.Date(fiscalyear, 4, 1, 0, 0, 0, 0, time.UTC)
	speriode := eperiode.AddDate(-1, 0, 0)

	filter := dbox.And(dbox.Gte("date.date", speriode), dbox.Lt("date.date", eperiode))

	if plcode == "" || value == 0 {
		toolkit.Println("PLCode and Value are mandatory to fill")
		os.Exit(1)
	}

	setinitialconnection()
	defer gdrj.CloseDb()

	toolkit.Println("Get Data Master...")
	prepmaster()
	prepmasterclean()

	toolkit.Println("Start Data Process...")
	c, _ := gdrj.Find(new(gdrj.SalesPL), filter, nil)
	defer c.Close()

	splcount := c.Count()
	step := splcount / 100
	i := 0

	for {
		i++

		spl := new(gdrj.SalesPL)
		e := c.Fetch(spl, 1, false)
		if e != nil {
			break
		}
		key := toolkit.Sprintf("%d_%d", spl.Date.Month, spl.Date.Year)
		k := ""

		if spl.Customer != nil {
			switch custgroup {
			case "branch":
				key = toolkit.Sprintf("%v_%v", key, spl.Customer.BranchID)
			case "channel":
				k = toolkit.Sprintf("%v|%v", spl.Customer.BranchID, spl.Customer.ChannelID)
				key = toolkit.Sprintf("%v_%v", key, k)
			case "group":
				k = toolkit.Sprintf("%v|%v", spl.Customer.BranchID, spl.Customer.CustomerGroup)
				key = toolkit.Sprintf("%v_%v", key, k)
			default:
				key = toolkit.Sprintf("%v_", key)
			}
		}

		if prodgroup == "skuid" && (spl.Product.ID == "" || len(spl.Product.ID) < 3) {
			continue
		}

		if spl.Product != nil {
			switch prodgroup {
			case "brand":
				key = toolkit.Sprintf("%v_%v", key, spl.Product.Brand)
			case "group":
				k = toolkit.Sprintf("%v|%v", spl.Product.Brand, spl.Product.BrandCategoryID)
				key = toolkit.Sprintf("%v_%v", key, k)
			case "skuid":
				k = toolkit.Sprintf("%v|%v|%v", spl.Product.Brand, spl.Product.BrandCategoryID, spl.Product.ID)
				key = toolkit.Sprintf("%v_%v", key, k)
			default:
				key = toolkit.Sprintf("%v_", key)
			}

		}

		mapkeysvalue[key] += spl.GrossAmount
		globalval += spl.GrossAmount

		if i > step {
			step += splcount / 100
			toolkit.Printfn("Processing %d of %d in %s", i, splcount,
				time.Since(t0).String())
		}

	}

	toolkit.Printfn("Preparing the data")
	mcount := len(mapkeysvalue)
	i = 0
	for k, v := range mapkeysvalue {
		i++

		akey := strings.Split(k, "_")
		tcustomer := new(gdrj.Customer)
		tproduct := new(gdrj.Product)

		if len(akey) > 2 && akey[2] != "" {
			switch custgroup {
			case "branch":
				tcustomer.BranchID = akey[2]
			case "channel":
				skey := strings.Split(akey[2], "|")
				tcustomer.BranchID = skey[0]
				if len(skey) > 1 {
					tcustomer.ChannelID = skey[1]
				}
			case "group":
				skey := strings.Split(akey[2], "|")
				tcustomer.BranchID = skey[0]
				if len(skey) > 1 {
					tcustomer.CustomerGroup = skey[1]
				}
			}
		}

		if len(akey) > 3 && akey[3] != "" {
			switch prodgroup {
			case "brand":
				tproduct.Brand = akey[3]
			case "group":
				skey := strings.Split(akey[3], "|")
				tproduct.Brand = skey[0]
				if len(skey) > 1 {
					tproduct.BrandCategoryID = skey[1]
				}
			case "skuid":
				skey := strings.Split(akey[3], "|")
				tproduct.Brand = skey[0]
				if len(skey) > 1 {
					tproduct.BrandCategoryID = skey[1]
				}
				if len(skey) > 2 {
					tproduct.ID = skey[2]
				}
			}
		}

		spl := new(gdrj.SalesPL)
		spl.Date = gdrj.SetDate(time.Date(toolkit.ToInt(akey[1], toolkit.RoundingAuto),
			time.Month(toolkit.ToInt(akey[0], toolkit.RoundingAuto)),
			1, 0, 0, 0, 0, time.UTC))
		spl.Customer = tcustomer
		spl.Product = tproduct
		spl.Source = "ADJUST"
		spl.Ref = ref

		spl.PC = new(gdrj.ProfitCenter)
		spl.CC = new(gdrj.CostCenter)

		amount := toolkit.ToFloat64(value, 6, toolkit.RoundingAuto) * (v / globalval)
		plmodels := masters.Get("plmodel").(map[string]*gdrj.PLModel)

		spl.ID = toolkit.Sprintf("%v_%v_%v_%v", "ALLOCATION-SCRIPT", spl.Date.Fiscal, akey[0], akey[1])

		spl.CleanAndClasify(masters)
		spl.AddData(plcode, amount, plmodels)
		spl.CalcSum(masters)

		gdrj.Save(spl)

		toolkit.Printfn("Saving %d of %d in %s", i, mcount,
			time.Since(t0).String())
	}

	toolkit.Printfn("Processing done in %s",
		time.Since(t0).String())
}
<|endoftext|>


<commit_before>package router

import (
	"fmt"
	"strings"
	"sync"

	"github.com/micro/go-micro/registry"
	"github.com/olekukonko/tablewriter"
)

// router provides default router implementation
type router struct {
	opts Options
	exit chan struct{}
	wg   *sync.WaitGroup
}

// newRouter creates new router and returns it
func newRouter(opts ...Option) Router {
	// TODO: we need to add default GW entry here
	// Should default GW be part of router options?

	// get default options
	options := DefaultOptions()

	// apply requested options
	for _, o := range opts {
		o(&options)
	}

	return &router{
		opts: options,
		exit: make(chan struct{}),
		wg:   &sync.WaitGroup{},
	}
}

// Init initializes router with given options
func (r *router) Init(opts ...Option) error {
	for _, o := range opts {
		o(&r.opts)
	}
	return nil
}

// Options returns router options
func (r *router) Options() Options {
	return r.opts
}

// ID returns router ID
func (r *router) ID() string {
	return r.opts.ID
}

// Table returns routing table
func (r *router) Table() Table {
	return r.opts.Table
}

// Address returns router's bind address
func (r *router) Address() string {
	return r.opts.Address
}

// Network returns the address router advertises to the network
func (r *router) Network() string {
	return r.opts.Network
}

// Advertise advertises the routes to the network. It is a blocking function.
// It returns error if any of the launched goroutines fail with error.
func (r *router) Advertise() error {
	// add local service routes into the routing table
	if err := r.addServiceRoutes(r.opts.Registry, "local", DefaultLocalMetric); err != nil {
		return fmt.Errorf("failed adding routes: %v", err)
	}

	localWatcher, err := r.opts.Registry.Watch()
	if err != nil {
		return fmt.Errorf("failed to create registry watcher: %v", err)
	}

	// error channel collecting goroutine errors
	errChan := make(chan error, 1)

	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		// watch local registry and register routes in routine table
		errChan <- r.manageServiceRoutes(localWatcher, DefaultLocalMetric)
	}()

	return <-errChan
}

// addServiceRoutes adds all services in given registry to the routing table.
// NOTE: this is a one-off operation done when bootstrapping the routing table
// It returns error if either the services failed to be listed or
// if the routes could not be added to the routing table.
func (r *router) addServiceRoutes(reg registry.Registry, network string, metric int) error {
	services, err := reg.ListServices()
	if err != nil {
		return fmt.Errorf("failed to list services: %v", err)
	}

	// add each service node as a separate route;
	for _, service := range services {
		for _, node := range service.Nodes {
			gw := node.Address
			if node.Port > 0 {
				gw = fmt.Sprintf("%s:%d", node.Address, node.Port)
			}
			route := Route{
				Destination: service.Name,
				Gateway:     gw,
				Router:      r.opts.Address,
				Network:     r.opts.Network,
				Metric:      metric,
			}
			if err := r.opts.Table.Add(route); err != nil && err != ErrDuplicateRoute {
				return fmt.Errorf("error adding route for service %s: %s", service.Name, err)
			}
		}
	}

	return nil
}

// manageServiceRoutes watches services in given registry and updates the routing table accordingly.
// It returns error if the service registry watcher has stopped or if the routing table failed to be updated.
func (r *router) manageServiceRoutes(w registry.Watcher, metric int) error {
	// wait in the background for the router to stop
	// when the router stops, stop the watcher and exit
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		<-r.exit
		w.Stop()
	}()

	var watchErr error

	for {
		res, err := w.Next()
		if err == registry.ErrWatcherStopped {
			break
		}

		if err != nil {
			watchErr = err
			break
		}

		route := Route{
			Destination: res.Service.Name,
			Router:      r.opts.Address,
			Network:     r.opts.Network,
			Metric:      metric,
		}

		switch res.Action {
		case "create":
			// only return error if the route is not duplicate, but something else has failed
			if err := r.opts.Table.Add(route); err != nil && err != ErrDuplicateRoute {
				return fmt.Errorf("failed to add route for service %v: %s", res.Service.Name, err)
			}
		case "delete":
			// only return error if the route is not in the table, but something else has failed
			if err := r.opts.Table.Delete(route); err != nil && err != ErrRouteNotFound {
				return fmt.Errorf("failed to delete route for service %v: %s", res.Service.Name, err)
			}
		}
	}

	return watchErr
}

// Stop stops the router
func (r *router) Stop() error {
	// notify all goroutines to finish
	close(r.exit)

	// wait for all goroutines to finish
	r.wg.Wait()

	return nil
}

// String prints debugging information about router
func (r *router) String() string {
	sb := &strings.Builder{}

	table := tablewriter.NewWriter(sb)
	table.SetHeader([]string{"ID", "Address", "Network", "Table"})

	data := []string{
		r.opts.ID,
		r.opts.Address,
		r.opts.Network,
		fmt.Sprintf("%d", r.opts.Table.Size()),
	}
	table.Append(data)

	// render table into sb
	table.Render()

	return sb.String()
}
<commit_msg>Lookup every service. FML<commit_after>package router

import (
	"fmt"
	"strings"
	"sync"

	"github.com/micro/go-micro/registry"
	"github.com/olekukonko/tablewriter"
)

// router provides default router implementation
type router struct {
	opts Options
	exit chan struct{}
	wg   *sync.WaitGroup
}

// newRouter creates new router and returns it
func newRouter(opts ...Option) Router {
	// TODO: we need to add default GW entry here
	// Should default GW be part of router options?

	// get default options
	options := DefaultOptions()

	// apply requested options
	for _, o := range opts {
		o(&options)
	}

	return &router{
		opts: options,
		exit: make(chan struct{}),
		wg:   &sync.WaitGroup{},
	}
}

// Init initializes router with given options
func (r *router) Init(opts ...Option) error {
	for _, o := range opts {
		o(&r.opts)
	}
	return nil
}

// Options returns router options
func (r *router) Options() Options {
	return r.opts
}

// ID returns router ID
func (r *router) ID() string {
	return r.opts.ID
}

// Table returns routing table
func (r *router) Table() Table {
	return r.opts.Table
}

// Address returns router's bind address
func (r *router) Address() string {
	return r.opts.Address
}

// Network returns the address router advertises to the network
func (r *router) Network() string {
	return r.opts.Network
}

// Advertise advertises the routes to the network. It is a blocking function.
// It returns error if any of the launched goroutines fail with error.
func (r *router) Advertise() error {
	// add local service routes into the routing table
	if err := r.addServiceRoutes(r.opts.Registry, "local", DefaultLocalMetric); err != nil {
		return fmt.Errorf("failed adding routes: %v", err)
	}

	localWatcher, err := r.opts.Registry.Watch()
	if err != nil {
		return fmt.Errorf("failed to create registry watcher: %v", err)
	}

	// error channel collecting goroutine errors
	errChan := make(chan error, 1)

	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		// watch local registry and register routes in routine table
		errChan <- r.manageServiceRoutes(localWatcher, DefaultLocalMetric)
	}()

	return <-errChan
}

// addServiceRoutes adds all services in given registry to the routing table.
// NOTE: this is a one-off operation done when bootstrapping the routing table
// It returns error if either the services failed to be listed or
// if the routes could not be added to the routing table.
func (r *router) addServiceRoutes(reg registry.Registry, network string, metric int) error {
	services, err := reg.ListServices()
	if err != nil {
		return fmt.Errorf("failed to list services: %v", err)
	}

	// add each service node as a separate route;
	for _, service := range services {
		// get the service to retrieve all its info
		srvs, err := reg.GetService(service.Name)
		if err != nil {
			continue
		}

		// create a flat slide of nodes
		var nodes []*registry.Node
		for _, s := range srvs {
			nodes = append(nodes, s.Nodes...)
		}

		// range over the flat slice of nodes
		for _, node := range nodes {
			gw := node.Address
			if node.Port > 0 {
				gw = fmt.Sprintf("%s:%d", node.Address, node.Port)
			}
			route := Route{
				Destination: service.Name,
				Gateway:     gw,
				Router:      r.opts.Address,
				Network:     r.opts.Network,
				Metric:      metric,
			}
			if err := r.opts.Table.Add(route); err != nil && err != ErrDuplicateRoute {
				return fmt.Errorf("error adding route for service %s: %s", service.Name, err)
			}
		}
	}

	return nil
}

// manageServiceRoutes watches services in given registry and updates the routing table accordingly.
// It returns error if the service registry watcher has stopped or if the routing table failed to be updated.
func (r *router) manageServiceRoutes(w registry.Watcher, metric int) error {
	// wait in the background for the router to stop
	// when the router stops, stop the watcher and exit
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		<-r.exit
		w.Stop()
	}()

	var watchErr error

	for {
		res, err := w.Next()
		if err == registry.ErrWatcherStopped {
			break
		}

		if err != nil {
			watchErr = err
			break
		}

		route := Route{
			Destination: res.Service.Name,
			Router:      r.opts.Address,
			Network:     r.opts.Network,
			Metric:      metric,
		}

		switch res.Action {
		case "create":
			// only return error if the route is not duplicate, but something else has failed
			if err := r.opts.Table.Add(route); err != nil && err != ErrDuplicateRoute {
				return fmt.Errorf("failed to add route for service %v: %s", res.Service.Name, err)
			}
		case "delete":
			// only return error if the route is not in the table, but something else has failed
			if err := r.opts.Table.Delete(route); err != nil && err != ErrRouteNotFound {
				return fmt.Errorf("failed to delete route for service %v: %s", res.Service.Name, err)
			}
		}
	}

	return watchErr
}

// Stop stops the router
func (r *router) Stop() error {
	// notify all goroutines to finish
	close(r.exit)

	// wait for all goroutines to finish
	r.wg.Wait()

	return nil
}

// String prints debugging information about router
func (r *router) String() string {
	sb := &strings.Builder{}

	table := tablewriter.NewWriter(sb)
	table.SetHeader([]string{"ID", "Address", "Network", "Table"})

	data := []string{
		r.opts.ID,
		r.opts.Address,
		r.opts.Network,
		fmt.Sprintf("%d", r.opts.Table.Size()),
	}
	table.Append(data)

	// render table into sb
	table.Render()

	return sb.String()
}
<|endoftext|>


<commit_before><commit_msg>Add operation<commit_after><|endoftext|>


<commit_before>//go:build !windows && !plan9 && !js
// +build !windows,!plan9,!js

package daemon

import (
	"io"
	"os"
	"os/user"
	"testing"

	"src.elv.sh/pkg/daemon/daemondefs"
	"src.elv.sh/pkg/testutil"
)

func TestActivate_InterruptsOutdatedServerAndSpawnsNewServer(t *testing.T) {
	activated := 0
	setupForActivate(t, func(name string, argv []string, attr *os.ProcAttr) error {
		startServer(t, argv)
		activated++
		return nil
	})
	version := api.Version - 1
	oldServer := startServerOpts(t, cli("sock", "db"), ServeOpts{Version: &version})

	_, err := Activate(io.Discard,
		&daemondefs.SpawnConfig{DbPath: "db", SockPath: "sock", RunDir: "."})
	if err != nil {
		t.Errorf("got error %v, want nil", err)
	}
	if activated != 1 {
		t.Errorf("got activated %v times, want 1", activated)
	}
	oldServer.WaitQuit()
}

func TestActivate_FailsIfUnableToRemoveHangingSocket(t *testing.T) {
	if u, err := user.Current(); err != nil || u.Uid == "0" {
		t.Skip("current user is root or unknown")
	}
	activated := 0
	setupForActivate(t, func(name string, argv []string, attr *os.ProcAttr) error {
		activated++
		return nil
	})
	testutil.MustMkdirAll("d")
	makeHangingUNIXSocket(t, "d/sock")
	// Remove write permission so that removing d/sock will fail
	os.Chmod("d", 0600)
	defer os.Chmod("d", 0700)

	_, err := Activate(io.Discard,
		&daemondefs.SpawnConfig{DbPath: "db", SockPath: "d/sock", RunDir: "."})
	if err == nil {
		t.Errorf("got error nil, want non-nil")
	}
	if activated != 0 {
		t.Errorf("got activated %v times, want 0", activated)
	}
}
<commit_msg>pkg/daemon: Fix activate_unix_test.go.<commit_after>//go:build !windows && !plan9 && !js
// +build !windows,!plan9,!js

package daemon

import (
	"io"
	"os"
	"os/user"
	"testing"

	"src.elv.sh/pkg/daemon/daemondefs"
	"src.elv.sh/pkg/daemon/internal/api"
	"src.elv.sh/pkg/testutil"
)

func TestActivate_InterruptsOutdatedServerAndSpawnsNewServer(t *testing.T) {
	activated := 0
	setupForActivate(t, func(name string, argv []string, attr *os.ProcAttr) error {
		startServer(t, argv)
		activated++
		return nil
	})
	version := api.Version - 1
	oldServer := startServerOpts(t, cli("sock", "db"), ServeOpts{Version: &version})

	_, err := Activate(io.Discard,
		&daemondefs.SpawnConfig{DbPath: "db", SockPath: "sock", RunDir: "."})
	if err != nil {
		t.Errorf("got error %v, want nil", err)
	}
	if activated != 1 {
		t.Errorf("got activated %v times, want 1", activated)
	}
	oldServer.WaitQuit()
}

func TestActivate_FailsIfUnableToRemoveHangingSocket(t *testing.T) {
	if u, err := user.Current(); err != nil || u.Uid == "0" {
		t.Skip("current user is root or unknown")
	}
	activated := 0
	setupForActivate(t, func(name string, argv []string, attr *os.ProcAttr) error {
		activated++
		return nil
	})
	testutil.MustMkdirAll("d")
	makeHangingUNIXSocket(t, "d/sock")
	// Remove write permission so that removing d/sock will fail
	os.Chmod("d", 0600)
	defer os.Chmod("d", 0700)

	_, err := Activate(io.Discard,
		&daemondefs.SpawnConfig{DbPath: "db", SockPath: "d/sock", RunDir: "."})
	if err == nil {
		t.Errorf("got error nil, want non-nil")
	}
	if activated != 0 {
		t.Errorf("got activated %v times, want 0", activated)
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package adapter

import (
	"context"
	"net/url"
	"sort"
	"sync"
	"time"

	"github.com/pingcap-incubator/tidb-dashboard/pkg/apiserver"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"go.uber.org/zap"

	"github.com/pingcap/log"
	"github.com/pingcap/pd/v4/pkg/logutil"
	"github.com/pingcap/pd/v4/server"
	"github.com/pingcap/pd/v4/server/cluster"
)

var (
	// CheckInterval represents the time interval of running check.
	CheckInterval = time.Second
)

// Manager is used to control dashboard.
type Manager struct {
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup

	srv        *server.Server
	service    *apiserver.Service
	redirector *Redirector

	isLeader bool
	members  []*pdpb.Member
}

// NewManager creates a new Manager.
func NewManager(srv *server.Server, s *apiserver.Service, redirector *Redirector) *Manager {
	ctx, cancel := context.WithCancel(srv.Context())
	return &Manager{
		ctx:        ctx,
		cancel:     cancel,
		srv:        srv,
		service:    s,
		redirector: redirector,
	}
}

// Start monitoring the dynamic config and control the dashboard.
func (m *Manager) Start() {
	m.wg.Add(1)
	go m.serviceLoop()
}

// Stop monitoring the dynamic config and control the dashboard.
func (m *Manager) Stop() {
	m.cancel()
	m.wg.Wait()
	log.Info("exit dashboard loop")
}

func (m *Manager) serviceLoop() {
	defer logutil.LogPanic()
	defer m.wg.Done()

	ticker := time.NewTicker(CheckInterval)
	defer ticker.Stop()

	for {
		select {
		case <-m.ctx.Done():
			m.stopService()
			return
		case <-ticker.C:
			m.updateInfo()
			m.checkAddress()
		}
	}
}

// updateInfo updates information from the server.
func (m *Manager) updateInfo() {
	if !m.srv.GetMember().IsLeader() {
		m.isLeader = false
		m.members = nil
		m.srv.GetPersistOptions().Reload(m.srv.GetStorage())
		return
	}

	m.isLeader = true

	var err error
	if m.members, err = cluster.GetMembers(m.srv.GetClient()); err != nil {
		log.Warn("failed to get members", zap.Error(err))
		m.members = nil
		return
	}

	allHasClientUrls := true
	for _, member := range m.members {
		if len(member.GetClientUrls()) == 0 {
			allHasClientUrls = false
		}
	}
	if !allHasClientUrls {
		log.Warn("failed to get member client urls")
		m.members = nil
	}
}

// checkDashboardAddress checks if the dashboard service needs to change due to dashboard address is changed.
func (m *Manager) checkAddress() {
	dashboardAddress := m.srv.GetPersistOptions().GetDashboardAddress()
	switch dashboardAddress {
	case "auto":
		if m.isLeader && len(m.members) > 0 {
			m.setNewAddress()
		}
		return
	case "none":
		m.redirector.SetAddress("")
		m.stopService()
		return
	default:
		if _, err := url.Parse(dashboardAddress); err != nil {
			log.Error("illegal dashboard address", zap.String("address", dashboardAddress))
			return
		}

		if m.isLeader && m.needResetAddress(dashboardAddress) {
			m.setNewAddress()
			return
		}
	}

	m.redirector.SetAddress(dashboardAddress)

	clientUrls := m.srv.GetMemberInfo().GetClientUrls()
	if len(clientUrls) > 0 && clientUrls[0] == dashboardAddress {
		m.startService()
	} else {
		m.stopService()
	}
}

func (m *Manager) needResetAddress(addr string) bool {
	if len(m.members) == 0 {
		return false
	}

	for _, member := range m.members {
		if member.GetClientUrls()[0] == addr {
			return false
		}
	}

	return true
}

func (m *Manager) setNewAddress() {
	// get new dashboard address
	members := m.members
	var addr string
	switch len(members) {
	case 1:
		addr = members[0].GetClientUrls()[0]
	default:
		addr = members[0].GetClientUrls()[0]
		leaderID := m.srv.GetMemberInfo().MemberId
		sort.Slice(members, func(i, j int) bool { return members[i].GetMemberId() < members[j].GetMemberId() })
		for _, member := range members {
			if member.MemberId != leaderID {
				addr = member.GetClientUrls()[0]
				break
			}
		}
	}
	// set new dashboard address
	cfg := m.srv.GetPersistOptions().GetPDServerConfig().Clone()
	cfg.DashboardAddress = addr
	m.srv.SetPDServerConfig(*cfg)
}

func (m *Manager) startService() {
	if m.service.IsRunning() {
		return
	}
	if err := m.service.Start(m.ctx); err != nil {
		log.Error("Can not start dashboard server", zap.Error(err))
	} else {
		log.Info("Dashboard server is started")
	}
}

func (m *Manager) stopService() {
	if !m.service.IsRunning() {
		return
	}
	if err := m.service.Stop(context.Background()); err != nil {
		log.Error("Stop dashboard server error", zap.Error(err))
	} else {
		log.Info("Dashboard server is stopped")
	}
}
<commit_msg>dashboard: generate Context every time the manager starts (#2466)<commit_after>// Copyright 2020 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package adapter

import (
	"context"
	"net/url"
	"sort"
	"sync"
	"time"

	"github.com/pingcap-incubator/tidb-dashboard/pkg/apiserver"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"go.uber.org/zap"

	"github.com/pingcap/log"
	"github.com/pingcap/pd/v4/pkg/logutil"
	"github.com/pingcap/pd/v4/server"
	"github.com/pingcap/pd/v4/server/cluster"
)

var (
	// CheckInterval represents the time interval of running check.
	CheckInterval = time.Second
)

// Manager is used to control dashboard.
type Manager struct {
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup

	srv        *server.Server
	service    *apiserver.Service
	redirector *Redirector

	isLeader bool
	members  []*pdpb.Member
}

// NewManager creates a new Manager.
func NewManager(srv *server.Server, s *apiserver.Service, redirector *Redirector) *Manager {
	return &Manager{
		srv:        srv,
		service:    s,
		redirector: redirector,
	}
}

// Start monitoring the dynamic config and control the dashboard.
func (m *Manager) Start() {
	m.ctx, m.cancel = context.WithCancel(m.srv.Context())
	m.wg.Add(1)
	m.isLeader = false
	m.members = nil
	go m.serviceLoop()
}

// Stop monitoring the dynamic config and control the dashboard.
func (m *Manager) Stop() {
	m.cancel()
	m.wg.Wait()
	log.Info("exit dashboard loop")
}

func (m *Manager) serviceLoop() {
	defer logutil.LogPanic()
	defer m.wg.Done()

	ticker := time.NewTicker(CheckInterval)
	defer ticker.Stop()

	for {
		select {
		case <-m.ctx.Done():
			m.stopService()
			return
		case <-ticker.C:
			m.updateInfo()
			m.checkAddress()
		}
	}
}

// updateInfo updates information from the server.
func (m *Manager) updateInfo() {
	if !m.srv.GetMember().IsLeader() {
		m.isLeader = false
		m.members = nil
		m.srv.GetPersistOptions().Reload(m.srv.GetStorage())
		return
	}

	m.isLeader = true

	var err error
	if m.members, err = cluster.GetMembers(m.srv.GetClient()); err != nil {
		log.Warn("failed to get members", zap.Error(err))
		m.members = nil
		return
	}

	allHasClientUrls := true
	for _, member := range m.members {
		if len(member.GetClientUrls()) == 0 {
			allHasClientUrls = false
		}
	}
	if !allHasClientUrls {
		log.Warn("failed to get member client urls")
		m.members = nil
	}
}

// checkDashboardAddress checks if the dashboard service needs to change due to dashboard address is changed.
func (m *Manager) checkAddress() {
	dashboardAddress := m.srv.GetPersistOptions().GetDashboardAddress()
	switch dashboardAddress {
	case "auto":
		if m.isLeader && len(m.members) > 0 {
			m.setNewAddress()
		}
		return
	case "none":
		m.redirector.SetAddress("")
		m.stopService()
		return
	default:
		if _, err := url.Parse(dashboardAddress); err != nil {
			log.Error("illegal dashboard address", zap.String("address", dashboardAddress))
			return
		}

		if m.isLeader && m.needResetAddress(dashboardAddress) {
			m.setNewAddress()
			return
		}
	}

	m.redirector.SetAddress(dashboardAddress)

	clientUrls := m.srv.GetMemberInfo().GetClientUrls()
	if len(clientUrls) > 0 && clientUrls[0] == dashboardAddress {
		m.startService()
	} else {
		m.stopService()
	}
}

func (m *Manager) needResetAddress(addr string) bool {
	if len(m.members) == 0 {
		return false
	}

	for _, member := range m.members {
		if member.GetClientUrls()[0] == addr {
			return false
		}
	}

	return true
}

func (m *Manager) setNewAddress() {
	// get new dashboard address
	members := m.members
	var addr string
	switch len(members) {
	case 1:
		addr = members[0].GetClientUrls()[0]
	default:
		addr = members[0].GetClientUrls()[0]
		leaderID := m.srv.GetMemberInfo().MemberId
		sort.Slice(members, func(i, j int) bool { return members[i].GetMemberId() < members[j].GetMemberId() })
		for _, member := range members {
			if member.MemberId != leaderID {
				addr = member.GetClientUrls()[0]
				break
			}
		}
	}
	// set new dashboard address
	cfg := m.srv.GetPersistOptions().GetPDServerConfig().Clone()
	cfg.DashboardAddress = addr
	m.srv.SetPDServerConfig(*cfg)
}

func (m *Manager) startService() {
	if m.service.IsRunning() {
		return
	}
	if err := m.service.Start(m.ctx); err != nil {
		log.Error("Can not start dashboard server", zap.Error(err))
	} else {
		log.Info("Dashboard server is started")
	}
}

func (m *Manager) stopService() {
	if !m.service.IsRunning() {
		return
	}
	if err := m.service.Stop(context.Background()); err != nil {
		log.Error("Stop dashboard server error", zap.Error(err))
	} else {
		log.Info("Dashboard server is stopped")
	}
}
<|endoftext|>


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

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

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

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

package download

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"cloud.google.com/go/storage"
	"google.golang.org/api/option"
	"k8s.io/minikube/pkg/minikube/detect"

	"github.com/pkg/errors"
	"github.com/spf13/viper"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/minikube/driver"
	"k8s.io/minikube/pkg/minikube/localpath"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/style"
)

const (
	// PreloadVersion is the current version of the preloaded tarball
	//
	// NOTE: You may need to bump this version up when upgrading auxiliary docker images
	PreloadVersion = "v14"
	// PreloadBucket is the name of the GCS bucket where preloaded volume tarballs exist
	PreloadBucket = "minikube-preloaded-volume-tarballs"
)

var (
	preloadStates = make(map[string]map[string]bool)
)

// TarballName returns name of the tarball
func TarballName(k8sVersion, containerRuntime string) string {
	if containerRuntime == "crio" {
		containerRuntime = "cri-o"
	}
	var storageDriver string
	if containerRuntime == "cri-o" {
		storageDriver = "overlay"
	} else {
		storageDriver = "overlay2"
	}
	arch := detect.EffectiveArch()
	return fmt.Sprintf("preloaded-images-k8s-%s-%s-%s-%s-%s.tar.lz4", PreloadVersion, k8sVersion, containerRuntime, storageDriver, arch)
}

// returns the name of the checksum file
func checksumName(k8sVersion, containerRuntime string) string {
	return fmt.Sprintf("%s.checksum", TarballName(k8sVersion, containerRuntime))
}

// returns target dir for all cached items related to preloading
func targetDir() string {
	return localpath.MakeMiniPath("cache", "preloaded-tarball")
}

// PreloadChecksumPath returns the local path to the cached checksum file
func PreloadChecksumPath(k8sVersion, containerRuntime string) string {
	return filepath.Join(targetDir(), checksumName(k8sVersion, containerRuntime))
}

// TarballPath returns the local path to the cached preload tarball
func TarballPath(k8sVersion, containerRuntime string) string {
	return filepath.Join(targetDir(), TarballName(k8sVersion, containerRuntime))
}

// remoteTarballURL returns the URL for the remote tarball in GCS
func remoteTarballURL(k8sVersion, containerRuntime string) string {
	return fmt.Sprintf("https://%s/%s/%s", downloadHost, PreloadBucket, TarballName(k8sVersion, containerRuntime))
}

func setPreloadState(k8sVersion, containerRuntime string, value bool) {
	cRuntimes, ok := preloadStates[k8sVersion]
	if !ok {
		cRuntimes = make(map[string]bool)
		preloadStates[k8sVersion] = cRuntimes
	}
	cRuntimes[containerRuntime] = value
}

var checkRemotePreloadExists = func(k8sVersion, containerRuntime string) bool {
	url := remoteTarballURL(k8sVersion, containerRuntime)
	resp, err := http.Head(url)
	if err != nil {
		klog.Warningf("%s fetch error: %v", url, err)
		return false
	}

	// note: err won't be set if it's a 404
	if resp.StatusCode != http.StatusOK {
		klog.Warningf("%s status code: %d", url, resp.StatusCode)
		return false
	}

	klog.Infof("Found remote preload: %s", url)
	return true
}

// PreloadExists returns true if there is a preloaded tarball that can be used
func PreloadExists(k8sVersion, containerRuntime, driverName string, forcePreload ...bool) bool {
	// TODO (#8166): Get rid of the need for this and viper at all
	force := false
	if len(forcePreload) > 0 {
		force = forcePreload[0]
	}

	// TODO: debug why this func is being called two times
	klog.Infof("Checking if preload exists for k8s version %s and runtime %s", k8sVersion, containerRuntime)
	// If `driverName` is BareMetal, there is no preload. Note: some uses of
	// `PreloadExists` assume that the driver is irrelevant unless BareMetal.
	if !driver.AllowsPreload(driverName) || !viper.GetBool("preload") && !force {
		return false
	}

	// If the preload existence is cached, just return that value.
	preloadState, ok := preloadStates[k8sVersion][containerRuntime]
	if ok {
		return preloadState
	}

	// Omit remote check if tarball exists locally
	targetPath := TarballPath(k8sVersion, containerRuntime)
	if _, err := checkCache(targetPath); err == nil {
		klog.Infof("Found local preload: %s", targetPath)
		setPreloadState(k8sVersion, containerRuntime, true)
		return true
	}

	existence := checkRemotePreloadExists(k8sVersion, containerRuntime)
	setPreloadState(k8sVersion, containerRuntime, existence)
	return existence
}

var checkPreloadExists = PreloadExists

// Preload caches the preloaded images tarball on the host machine
func Preload(k8sVersion, containerRuntime, driverName string) error {
	targetPath := TarballPath(k8sVersion, containerRuntime)
	targetLock := targetPath + ".lock"

	releaser, err := lockDownload(targetLock)
	if releaser != nil {
		defer releaser.Release()
	}
	if err != nil {
		return err
	}

	if _, err := checkCache(targetPath); err == nil {
		klog.Infof("Found %s in cache, skipping download", targetPath)
		return nil
	}

	// Make sure we support this k8s version
	if !checkPreloadExists(k8sVersion, containerRuntime, driverName) {
		klog.Infof("Preloaded tarball for k8s version %s does not exist", k8sVersion)
		return nil
	}

	out.Step(style.FileDownload, "Downloading Kubernetes {{.version}} preload ...", out.V{"version": k8sVersion})
	url := remoteTarballURL(k8sVersion, containerRuntime)

	checksum, err := getChecksum(k8sVersion, containerRuntime)
	var realPath string
	if err != nil {
		klog.Warningf("No checksum for preloaded tarball for k8s version %s: %v", k8sVersion, err)
		realPath = targetPath
		tmp, err := os.CreateTemp(targetDir(), TarballName(k8sVersion, containerRuntime)+".*")
		if err != nil {
			return errors.Wrap(err, "tempfile")
		}
		targetPath = tmp.Name()
	} else if checksum != nil {
		// add URL parameter for go-getter to automatically verify the checksum
		url += fmt.Sprintf("?checksum=md5:%s", hex.EncodeToString(checksum))
	}

	if err := download(url, targetPath); err != nil {
		return errors.Wrapf(err, "download failed: %s", url)
	}

	if err := ensureChecksumValid(k8sVersion, containerRuntime, targetPath, checksum); err != nil {
		return err
	}

	if realPath != "" {
		klog.Infof("renaming tempfile to %s ...", TarballName(k8sVersion, containerRuntime))
		err := os.Rename(targetPath, realPath)
		if err != nil {
			return errors.Wrap(err, "rename")
		}
	}

	// If the download was successful, mark off that the preload exists in the cache.
	setPreloadState(k8sVersion, containerRuntime, true)
	return nil
}

func getStorageAttrs(name string) (*storage.ObjectAttrs, error) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithoutAuthentication())
	if err != nil {
		return nil, errors.Wrap(err, "getting storage client")
	}
	attrs, err := client.Bucket(PreloadBucket).Object(name).Attrs(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting storage object")
	}
	return attrs, nil
}

// getChecksum returns the MD5 checksum of the preload tarball
var getChecksum = func(k8sVersion, containerRuntime string) ([]byte, error) {
	klog.Infof("getting checksum for %s ...", TarballName(k8sVersion, containerRuntime))
	attrs, err := getStorageAttrs(TarballName(k8sVersion, containerRuntime))
	if err != nil {
		return nil, err
	}
	return attrs.MD5, nil
}

// saveChecksumFile saves the checksum to a local file for later verification
func saveChecksumFile(k8sVersion, containerRuntime string, checksum []byte) error {
	klog.Infof("saving checksum for %s ...", TarballName(k8sVersion, containerRuntime))
	return os.WriteFile(PreloadChecksumPath(k8sVersion, containerRuntime), checksum, 0o644)
}

// verifyChecksum returns true if the checksum of the local binary matches
// the checksum of the remote binary
func verifyChecksum(k8sVersion, containerRuntime, path string) error {
	klog.Infof("verifying checksumm of %s ...", path)
	// get md5 checksum of tarball path
	contents, err := os.ReadFile(path)
	if err != nil {
		return errors.Wrap(err, "reading tarball")
	}
	checksum := md5.Sum(contents)

	remoteChecksum, err := os.ReadFile(PreloadChecksumPath(k8sVersion, containerRuntime))
	if err != nil {
		return errors.Wrap(err, "reading checksum file")
	}

	// create a slice of checksum, which is [16]byte
	if string(remoteChecksum) != string(checksum[:]) {
		return fmt.Errorf("checksum of %s does not match remote checksum (%s != %s)", path, string(remoteChecksum), string(checksum[:]))
	}
	return nil
}

// ensureChecksumValid saves and verifies local binary checksum matches remote binary checksum
var ensureChecksumValid = func(k8sVersion, containerRuntime, targetPath string, checksum []byte) error {
	if err := saveChecksumFile(k8sVersion, containerRuntime, checksum); err != nil {
		return errors.Wrap(err, "saving checksum file")
	}

	if err := verifyChecksum(k8sVersion, containerRuntime, targetPath); err != nil {
		return errors.Wrap(err, "verify")
	}

	return nil
}

// CleanUpOlderPreloads deletes preload files beloning to older minikube versions
// checks the current preload version and then if the saved tar file is belongs to older minikube it will delete it
// in case of failure only logs to the user
func CleanUpOlderPreloads() {
	files, err := os.ReadDir(targetDir())
	if err != nil {
		klog.Warningf("Failed to list preload files: %v", err)
	}

	for _, file := range files {
		splited := strings.Split(file.Name(), "-")
		if len(splited) < 4 {
			continue
		}
		ver := splited[3]
		if ver != PreloadVersion {
			fn := path.Join(targetDir(), file.Name())
			klog.Infof("deleting older generation preload %s", fn)
			err := os.Remove(fn)
			if err != nil {
				klog.Warningf("Failed to clean up older preload files, consider running `minikube delete --all --purge`")
			}
		}
	}
}
<commit_msg>Bump PreloadVersion const to v15<commit_after>/*
Copyright 2020 The Kubernetes Authors All rights reserved.

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

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

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

package download

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"cloud.google.com/go/storage"
	"google.golang.org/api/option"
	"k8s.io/minikube/pkg/minikube/detect"

	"github.com/pkg/errors"
	"github.com/spf13/viper"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/minikube/driver"
	"k8s.io/minikube/pkg/minikube/localpath"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/style"
)

const (
	// PreloadVersion is the current version of the preloaded tarball
	//
	// NOTE: You may need to bump this version up when upgrading auxiliary docker images
	PreloadVersion = "v15"
	// PreloadBucket is the name of the GCS bucket where preloaded volume tarballs exist
	PreloadBucket = "minikube-preloaded-volume-tarballs"
)

var (
	preloadStates = make(map[string]map[string]bool)
)

// TarballName returns name of the tarball
func TarballName(k8sVersion, containerRuntime string) string {
	if containerRuntime == "crio" {
		containerRuntime = "cri-o"
	}
	var storageDriver string
	if containerRuntime == "cri-o" {
		storageDriver = "overlay"
	} else {
		storageDriver = "overlay2"
	}
	arch := detect.EffectiveArch()
	return fmt.Sprintf("preloaded-images-k8s-%s-%s-%s-%s-%s.tar.lz4", PreloadVersion, k8sVersion, containerRuntime, storageDriver, arch)
}

// returns the name of the checksum file
func checksumName(k8sVersion, containerRuntime string) string {
	return fmt.Sprintf("%s.checksum", TarballName(k8sVersion, containerRuntime))
}

// returns target dir for all cached items related to preloading
func targetDir() string {
	return localpath.MakeMiniPath("cache", "preloaded-tarball")
}

// PreloadChecksumPath returns the local path to the cached checksum file
func PreloadChecksumPath(k8sVersion, containerRuntime string) string {
	return filepath.Join(targetDir(), checksumName(k8sVersion, containerRuntime))
}

// TarballPath returns the local path to the cached preload tarball
func TarballPath(k8sVersion, containerRuntime string) string {
	return filepath.Join(targetDir(), TarballName(k8sVersion, containerRuntime))
}

// remoteTarballURL returns the URL for the remote tarball in GCS
func remoteTarballURL(k8sVersion, containerRuntime string) string {
	return fmt.Sprintf("https://%s/%s/%s", downloadHost, PreloadBucket, TarballName(k8sVersion, containerRuntime))
}

func setPreloadState(k8sVersion, containerRuntime string, value bool) {
	cRuntimes, ok := preloadStates[k8sVersion]
	if !ok {
		cRuntimes = make(map[string]bool)
		preloadStates[k8sVersion] = cRuntimes
	}
	cRuntimes[containerRuntime] = value
}

var checkRemotePreloadExists = func(k8sVersion, containerRuntime string) bool {
	url := remoteTarballURL(k8sVersion, containerRuntime)
	resp, err := http.Head(url)
	if err != nil {
		klog.Warningf("%s fetch error: %v", url, err)
		return false
	}

	// note: err won't be set if it's a 404
	if resp.StatusCode != http.StatusOK {
		klog.Warningf("%s status code: %d", url, resp.StatusCode)
		return false
	}

	klog.Infof("Found remote preload: %s", url)
	return true
}

// PreloadExists returns true if there is a preloaded tarball that can be used
func PreloadExists(k8sVersion, containerRuntime, driverName string, forcePreload ...bool) bool {
	// TODO (#8166): Get rid of the need for this and viper at all
	force := false
	if len(forcePreload) > 0 {
		force = forcePreload[0]
	}

	// TODO: debug why this func is being called two times
	klog.Infof("Checking if preload exists for k8s version %s and runtime %s", k8sVersion, containerRuntime)
	// If `driverName` is BareMetal, there is no preload. Note: some uses of
	// `PreloadExists` assume that the driver is irrelevant unless BareMetal.
	if !driver.AllowsPreload(driverName) || !viper.GetBool("preload") && !force {
		return false
	}

	// If the preload existence is cached, just return that value.
	preloadState, ok := preloadStates[k8sVersion][containerRuntime]
	if ok {
		return preloadState
	}

	// Omit remote check if tarball exists locally
	targetPath := TarballPath(k8sVersion, containerRuntime)
	if _, err := checkCache(targetPath); err == nil {
		klog.Infof("Found local preload: %s", targetPath)
		setPreloadState(k8sVersion, containerRuntime, true)
		return true
	}

	existence := checkRemotePreloadExists(k8sVersion, containerRuntime)
	setPreloadState(k8sVersion, containerRuntime, existence)
	return existence
}

var checkPreloadExists = PreloadExists

// Preload caches the preloaded images tarball on the host machine
func Preload(k8sVersion, containerRuntime, driverName string) error {
	targetPath := TarballPath(k8sVersion, containerRuntime)
	targetLock := targetPath + ".lock"

	releaser, err := lockDownload(targetLock)
	if releaser != nil {
		defer releaser.Release()
	}
	if err != nil {
		return err
	}

	if _, err := checkCache(targetPath); err == nil {
		klog.Infof("Found %s in cache, skipping download", targetPath)
		return nil
	}

	// Make sure we support this k8s version
	if !checkPreloadExists(k8sVersion, containerRuntime, driverName) {
		klog.Infof("Preloaded tarball for k8s version %s does not exist", k8sVersion)
		return nil
	}

	out.Step(style.FileDownload, "Downloading Kubernetes {{.version}} preload ...", out.V{"version": k8sVersion})
	url := remoteTarballURL(k8sVersion, containerRuntime)

	checksum, err := getChecksum(k8sVersion, containerRuntime)
	var realPath string
	if err != nil {
		klog.Warningf("No checksum for preloaded tarball for k8s version %s: %v", k8sVersion, err)
		realPath = targetPath
		tmp, err := os.CreateTemp(targetDir(), TarballName(k8sVersion, containerRuntime)+".*")
		if err != nil {
			return errors.Wrap(err, "tempfile")
		}
		targetPath = tmp.Name()
	} else if checksum != nil {
		// add URL parameter for go-getter to automatically verify the checksum
		url += fmt.Sprintf("?checksum=md5:%s", hex.EncodeToString(checksum))
	}

	if err := download(url, targetPath); err != nil {
		return errors.Wrapf(err, "download failed: %s", url)
	}

	if err := ensureChecksumValid(k8sVersion, containerRuntime, targetPath, checksum); err != nil {
		return err
	}

	if realPath != "" {
		klog.Infof("renaming tempfile to %s ...", TarballName(k8sVersion, containerRuntime))
		err := os.Rename(targetPath, realPath)
		if err != nil {
			return errors.Wrap(err, "rename")
		}
	}

	// If the download was successful, mark off that the preload exists in the cache.
	setPreloadState(k8sVersion, containerRuntime, true)
	return nil
}

func getStorageAttrs(name string) (*storage.ObjectAttrs, error) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithoutAuthentication())
	if err != nil {
		return nil, errors.Wrap(err, "getting storage client")
	}
	attrs, err := client.Bucket(PreloadBucket).Object(name).Attrs(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting storage object")
	}
	return attrs, nil
}

// getChecksum returns the MD5 checksum of the preload tarball
var getChecksum = func(k8sVersion, containerRuntime string) ([]byte, error) {
	klog.Infof("getting checksum for %s ...", TarballName(k8sVersion, containerRuntime))
	attrs, err := getStorageAttrs(TarballName(k8sVersion, containerRuntime))
	if err != nil {
		return nil, err
	}
	return attrs.MD5, nil
}

// saveChecksumFile saves the checksum to a local file for later verification
func saveChecksumFile(k8sVersion, containerRuntime string, checksum []byte) error {
	klog.Infof("saving checksum for %s ...", TarballName(k8sVersion, containerRuntime))
	return os.WriteFile(PreloadChecksumPath(k8sVersion, containerRuntime), checksum, 0o644)
}

// verifyChecksum returns true if the checksum of the local binary matches
// the checksum of the remote binary
func verifyChecksum(k8sVersion, containerRuntime, path string) error {
	klog.Infof("verifying checksumm of %s ...", path)
	// get md5 checksum of tarball path
	contents, err := os.ReadFile(path)
	if err != nil {
		return errors.Wrap(err, "reading tarball")
	}
	checksum := md5.Sum(contents)

	remoteChecksum, err := os.ReadFile(PreloadChecksumPath(k8sVersion, containerRuntime))
	if err != nil {
		return errors.Wrap(err, "reading checksum file")
	}

	// create a slice of checksum, which is [16]byte
	if string(remoteChecksum) != string(checksum[:]) {
		return fmt.Errorf("checksum of %s does not match remote checksum (%s != %s)", path, string(remoteChecksum), string(checksum[:]))
	}
	return nil
}

// ensureChecksumValid saves and verifies local binary checksum matches remote binary checksum
var ensureChecksumValid = func(k8sVersion, containerRuntime, targetPath string, checksum []byte) error {
	if err := saveChecksumFile(k8sVersion, containerRuntime, checksum); err != nil {
		return errors.Wrap(err, "saving checksum file")
	}

	if err := verifyChecksum(k8sVersion, containerRuntime, targetPath); err != nil {
		return errors.Wrap(err, "verify")
	}

	return nil
}

// CleanUpOlderPreloads deletes preload files beloning to older minikube versions
// checks the current preload version and then if the saved tar file is belongs to older minikube it will delete it
// in case of failure only logs to the user
func CleanUpOlderPreloads() {
	files, err := os.ReadDir(targetDir())
	if err != nil {
		klog.Warningf("Failed to list preload files: %v", err)
	}

	for _, file := range files {
		splited := strings.Split(file.Name(), "-")
		if len(splited) < 4 {
			continue
		}
		ver := splited[3]
		if ver != PreloadVersion {
			fn := path.Join(targetDir(), file.Name())
			klog.Infof("deleting older generation preload %s", fn)
			err := os.Remove(fn)
			if err != nil {
				klog.Warningf("Failed to clean up older preload files, consider running `minikube delete --all --purge`")
			}
		}
	}
}
<|endoftext|>


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

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

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

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

package shell

import (
	"bytes"
	"os"
	"strings"
	"testing"
)

func TestGenerateUsageHint(t *testing.T) {
	var testCases = []struct {
		ec       EnvConfig
		expected string
	}{
		{EnvConfig{""}, `# foo
# eval $(bar)`},
		{EnvConfig{"powershell"}, `# foo
# & bar | Invoke-Expression`},
		{EnvConfig{"bash"}, `# foo
# eval $(bar)`},
		{EnvConfig{"powershell"}, `# foo
# & bar | Invoke-Expression`},
		{EnvConfig{"emacs"}, `;; foo
;; (with-temp-buffer (shell-command "bar" (current-buffer)) (eval-buffer))`},
		{EnvConfig{"fish"}, `# foo
# bar | source`},
		{EnvConfig{"none"}, `# foo
# eval $(bar)`},
	}
	for _, tc := range testCases {
		tc := tc
		t.Run(tc.ec.Shell, func(t *testing.T) {
			got := strings.TrimSpace(generateUsageHint(tc.ec, "foo", "bar"))
			expected := strings.TrimSpace(tc.expected)
			if got != expected {
				t.Errorf("Expected '%v' but got '%v'", expected, got)
			}
		})
	}
}

func TestCfgSet(t *testing.T) {
	var testCases = []struct {
		plz, cmd string
		ec       EnvConfig
		expected string
	}{
		{"", "eval", EnvConfig{""}, `"`},
		{"", "eval", EnvConfig{"bash"}, `"`},
		{"", "eval", EnvConfig{"powershell"}, `"`},
		{"", "eval", EnvConfig{"cmd"}, ``},
		{"", "eval", EnvConfig{"emacs"}, `")`},
		{"", "eval", EnvConfig{"none"}, ``},
		{"", "eval", EnvConfig{"fish"}, `";`},
	}
	for _, tc := range testCases {
		tc := tc
		t.Run(tc.ec.Shell, func(t *testing.T) {
			conf := CfgSet(tc.ec, tc.plz, tc.cmd)
			expected := strings.TrimSpace(tc.expected)
			got := strings.TrimSpace(conf.Suffix)
			if expected != got {
				t.Errorf("Expected suffix '%v' but got '%v'", expected, got)
			}
		})
	}
}

func TestUnsetScript(t *testing.T) {
	var testCases = []struct {
		vars     []string
		ec       EnvConfig
		expected string
	}{
		{[]string{"baz", "bar"}, EnvConfig{""}, `unset baz bar`},
		{[]string{"baz", "bar"}, EnvConfig{"bash"}, `unset baz bar`},
		{[]string{"baz", "bar"}, EnvConfig{"powershell"}, `Remove-Item Env:\\baz Env:\\bar`},
		{[]string{"baz", "bar"}, EnvConfig{"cmd"}, `SET baz=
SET bar=`},
		{[]string{"baz", "bar"}, EnvConfig{"fish"}, `set -e baz;
set -e bar;`},
		{[]string{"baz", "bar"}, EnvConfig{"emacs"}, `(setenv "baz" nil)
(setenv "bar" nil)`},
		{[]string{"baz", "bar"}, EnvConfig{"none"}, `baz bar`},
	}
	for _, tc := range testCases {
		tc := tc
		t.Run(tc.ec.Shell, func(t *testing.T) {
			var b bytes.Buffer

			if err := UnsetScript(tc.ec, &b, tc.vars); err != nil {
				t.Fatalf("Unexpected error when unseting script happen: %v", err)
			} else {
				writtenMessage := strings.TrimSpace(b.String())
				expected := strings.TrimSpace(tc.expected)
				if writtenMessage != expected {
					t.Fatalf("Expected '%v' but got '%v' ", tc.expected, writtenMessage)
				}
			}
		})
	}
}

func TestDetectSet(t *testing.T) {
	orgShellEnv := os.Getenv("SHELL")
	defer os.Setenv("SHELL", orgShellEnv)

	os.Setenv("SHELL", "/bin/bash")
	if s, err := Detect(); err != nil {
		t.Fatalf("unexpected error: '%v' during shell detection. Returned shell: %s", err, s)
	} else if s == "" {
		t.Fatalf("Detected shell expected to be non empty string")
	}
}

func TestDetectUnset(t *testing.T) {
	orgShellEnv := os.Getenv("SHELL")
	defer os.Setenv("SHELL", orgShellEnv)

	os.Unsetenv("SHELL")
	if s, err := Detect(); err != nil {
		t.Fatalf("unexpected error: '%v' during shell detection. Returned shell: %s", err, s)
	} else if s == "" {
		t.Fatalf("Detected shell expected to be non empty string")
	}
}

func TestSetScript(t *testing.T) {
	ec := EnvConfig{"bash"}
	var w bytes.Buffer
	if err := SetScript(ec, &w, "foo", nil); err != nil {
		t.Fatalf("Unexpected error: '%v' during Setting script", err)
	}
	if w.String() != "foo" {
		t.Fatalf("Expected foo writed by SetScript, but got '%v'", w.String())
	}
	if ec.Shell == "" {
		t.Fatalf("Expected no empty shell")
	}
}
<commit_msg>Fix shell test to reflect the fact that shell=none no longer has a usage msg.<commit_after>/*
Copyright 2020 The Kubernetes Authors All rights reserved.

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

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

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

package shell

import (
	"bytes"
	"os"
	"strings"
	"testing"
)

func TestGenerateUsageHint(t *testing.T) {
	var testCases = []struct {
		ec       EnvConfig
		expected string
	}{
		{EnvConfig{""}, `# foo
# eval $(bar)`},
		{EnvConfig{"powershell"}, `# foo
# & bar | Invoke-Expression`},
		{EnvConfig{"bash"}, `# foo
# eval $(bar)`},
		{EnvConfig{"powershell"}, `# foo
# & bar | Invoke-Expression`},
		{EnvConfig{"emacs"}, `;; foo
;; (with-temp-buffer (shell-command "bar" (current-buffer)) (eval-buffer))`},
		{EnvConfig{"fish"}, `# foo
# bar | source`},
		{EnvConfig{"none"}, ``},
	}
	for _, tc := range testCases {
		tc := tc
		t.Run(tc.ec.Shell, func(t *testing.T) {
			got := strings.TrimSpace(generateUsageHint(tc.ec, "foo", "bar"))
			expected := strings.TrimSpace(tc.expected)
			if got != expected {
				t.Errorf("Expected '%v' but got '%v'", expected, got)
			}
		})
	}
}

func TestCfgSet(t *testing.T) {
	var testCases = []struct {
		plz, cmd string
		ec       EnvConfig
		expected string
	}{
		{"", "eval", EnvConfig{""}, `"`},
		{"", "eval", EnvConfig{"bash"}, `"`},
		{"", "eval", EnvConfig{"powershell"}, `"`},
		{"", "eval", EnvConfig{"cmd"}, ``},
		{"", "eval", EnvConfig{"emacs"}, `")`},
		{"", "eval", EnvConfig{"none"}, ``},
		{"", "eval", EnvConfig{"fish"}, `";`},
	}
	for _, tc := range testCases {
		tc := tc
		t.Run(tc.ec.Shell, func(t *testing.T) {
			conf := CfgSet(tc.ec, tc.plz, tc.cmd)
			expected := strings.TrimSpace(tc.expected)
			got := strings.TrimSpace(conf.Suffix)
			if expected != got {
				t.Errorf("Expected suffix '%v' but got '%v'", expected, got)
			}
		})
	}
}

func TestUnsetScript(t *testing.T) {
	var testCases = []struct {
		vars     []string
		ec       EnvConfig
		expected string
	}{
		{[]string{"baz", "bar"}, EnvConfig{""}, `unset baz bar`},
		{[]string{"baz", "bar"}, EnvConfig{"bash"}, `unset baz bar`},
		{[]string{"baz", "bar"}, EnvConfig{"powershell"}, `Remove-Item Env:\\baz Env:\\bar`},
		{[]string{"baz", "bar"}, EnvConfig{"cmd"}, `SET baz=
SET bar=`},
		{[]string{"baz", "bar"}, EnvConfig{"fish"}, `set -e baz;
set -e bar;`},
		{[]string{"baz", "bar"}, EnvConfig{"emacs"}, `(setenv "baz" nil)
(setenv "bar" nil)`},
		{[]string{"baz", "bar"}, EnvConfig{"none"}, `baz bar`},
	}
	for _, tc := range testCases {
		tc := tc
		t.Run(tc.ec.Shell, func(t *testing.T) {
			var b bytes.Buffer

			if err := UnsetScript(tc.ec, &b, tc.vars); err != nil {
				t.Fatalf("Unexpected error when unseting script happen: %v", err)
			} else {
				writtenMessage := strings.TrimSpace(b.String())
				expected := strings.TrimSpace(tc.expected)
				if writtenMessage != expected {
					t.Fatalf("Expected '%v' but got '%v' ", tc.expected, writtenMessage)
				}
			}
		})
	}
}

func TestDetectSet(t *testing.T) {
	orgShellEnv := os.Getenv("SHELL")
	defer os.Setenv("SHELL", orgShellEnv)

	os.Setenv("SHELL", "/bin/bash")
	if s, err := Detect(); err != nil {
		t.Fatalf("unexpected error: '%v' during shell detection. Returned shell: %s", err, s)
	} else if s == "" {
		t.Fatalf("Detected shell expected to be non empty string")
	}
}

func TestDetectUnset(t *testing.T) {
	orgShellEnv := os.Getenv("SHELL")
	defer os.Setenv("SHELL", orgShellEnv)

	os.Unsetenv("SHELL")
	if s, err := Detect(); err != nil {
		t.Fatalf("unexpected error: '%v' during shell detection. Returned shell: %s", err, s)
	} else if s == "" {
		t.Fatalf("Detected shell expected to be non empty string")
	}
}

func TestSetScript(t *testing.T) {
	ec := EnvConfig{"bash"}
	var w bytes.Buffer
	if err := SetScript(ec, &w, "foo", nil); err != nil {
		t.Fatalf("Unexpected error: '%v' during Setting script", err)
	}
	if w.String() != "foo" {
		t.Fatalf("Expected foo writed by SetScript, but got '%v'", w.String())
	}
	if ec.Shell == "" {
		t.Fatalf("Expected no empty shell")
	}
}
<|endoftext|>


<commit_before>package gotest

import (
	"flag"
	"regexp"
	"strings"
	"testing"
	"time"

	"github.com/jstemmer/go-junit-report/v2/pkg/gtr"

	"github.com/google/go-cmp/cmp"
)

var matchTest = flag.String("match", "", "only test testdata matching this pattern")

var tests = []struct {
	name   string
	input  string
	events []gtr.Event
}{
	{
		"run",
		inp("=== RUN TestOne",
			"=== RUN   TestTwo/Subtest",
		),
		[]gtr.Event{
			{Type: "run_test", Name: "TestOne"},
			{Type: "run_test", Name: "TestTwo/Subtest"},
		},
	},
	{
		"pause",
		"=== PAUSE TestOne",
		[]gtr.Event{
			{Type: "pause_test", Name: "TestOne"},
		},
	},
	{
		"cont",
		"=== CONT  TestOne",
		[]gtr.Event{
			{Type: "cont_test", Name: "TestOne"},
		},
	},
	{
		"end",
		inp("--- PASS: TestOne (12.34 seconds)",
			"    --- SKIP: TestOne/Subtest (0.00s)",
			"        --- FAIL: TestOne/Subtest/#01 (0.35s)",
			"some text--- PASS: TestTwo (0.06 seconds)",
		),
		[]gtr.Event{
			{Type: "end_test", Name: "TestOne", Result: "PASS", Duration: 12_340 * time.Millisecond},
			{Type: "end_test", Name: "TestOne/Subtest", Result: "SKIP", Indent: 1},
			{Type: "end_test", Name: "TestOne/Subtest/#01", Result: "FAIL", Duration: 350 * time.Millisecond, Indent: 2},
			{Type: "output", Data: "some text"},
			{Type: "end_test", Name: "TestTwo", Result: "PASS", Duration: 60 * time.Millisecond},
		},
	},
	{
		"status",
		inp("PASS",
			"FAIL",
			"SKIP",
		),
		[]gtr.Event{
			{Type: "status", Result: "PASS"},
			{Type: "status", Result: "FAIL"},
			{Type: "status", Result: "SKIP"},
		},
	},
	{
		"summary",
		inp("ok      package/name/ok 0.100s",
			"FAIL    package/name/failing [build failed]",
			"FAIL    package/other/failing [setup failed]",
			"ok package/other     (cached)",
			"ok  	package/name 0.400s  coverage: 10.0% of statements",
			"ok  	package/name 4.200s  coverage: 99.8% of statements in fmt, encoding/xml",
			"?   	package/name	[no test files]",
			"ok  	package/name	0.001s [no tests to run]",
			"ok  	package/name	(cached) [no tests to run]",
		),
		[]gtr.Event{
			{Type: "summary", Name: "package/name/ok", Result: "ok", Duration: 100 * time.Millisecond},
			{Type: "summary", Name: "package/name/failing", Result: "FAIL", Data: "[build failed]"},
			{Type: "summary", Name: "package/other/failing", Result: "FAIL", Data: "[setup failed]"},
			{Type: "summary", Name: "package/other", Result: "ok", Data: "(cached)"},
			{Type: "summary", Name: "package/name", Result: "ok", Duration: 400 * time.Millisecond, CovPct: 10},
			{Type: "summary", Name: "package/name", Result: "ok", Duration: 4200 * time.Millisecond, CovPct: 99.8, CovPackages: []string{"fmt", "encoding/xml"}},
			{Type: "summary", Name: "package/name", Result: "?", Data: "[no test files]"},
			{Type: "summary", Name: "package/name", Result: "ok", Duration: 1 * time.Millisecond, Data: "[no tests to run]"},
			{Type: "summary", Name: "package/name", Result: "ok", Data: "(cached) [no tests to run]"},
		},
	},
	{
		"coverage",
		inp("coverage: 10% of statements",
			"coverage: 10% of statements in fmt, encoding/xml",
			"coverage: 13.37% of statements",
			"coverage: 99.8% of statements in fmt, encoding/xml",
		),
		[]gtr.Event{
			{Type: "coverage", CovPct: 10},
			{Type: "coverage", CovPct: 10, CovPackages: []string{"fmt", "encoding/xml"}},
			{Type: "coverage", CovPct: 13.37},
			{Type: "coverage", CovPct: 99.8, CovPackages: []string{"fmt", "encoding/xml"}},
		},
	},
	{
		"benchmark",
		inp("BenchmarkOne-8                     2000000	       604 ns/op",
			"BenchmarkTwo-16 30000	52568 ns/op	24879 B/op	494 allocs/op",
			"BenchmarkThree      2000000000	         0.26 ns/op",
			"BenchmarkFour-8         	   10000	    104427 ns/op	  95.76 MB/s	   40629 B/op	       5 allocs/op",
		),
		[]gtr.Event{
			{Type: "benchmark", Name: "BenchmarkOne", Iterations: 2_000_000, NsPerOp: 604},
			{Type: "benchmark", Name: "BenchmarkTwo", Iterations: 30_000, NsPerOp: 52_568, BytesPerOp: 24_879, AllocsPerOp: 494},
			{Type: "benchmark", Name: "BenchmarkThree", Iterations: 2_000_000_000, NsPerOp: 0.26},
			{Type: "benchmark", Name: "BenchmarkFour", Iterations: 10_000, NsPerOp: 104_427, MBPerSec: 95.76, BytesPerOp: 40_629, AllocsPerOp: 5},
		},
	},
	{
		"build output",
		inp("# package/name/failing1",
			"# package/name/failing2 [package/name/failing2.test]",
		),
		[]gtr.Event{
			{Type: "build_output", Name: "package/name/failing1"},
			{Type: "build_output", Name: "package/name/failing2"},
		},
	},
	{
		"output",
		inp("single line stdout",
			"# some more output",
			"\tfile_test.go:11: Error message",
			"\tfile_test.go:12: Longer",
			"\t\terror",
			"\t\tmessage.",
		),
		[]gtr.Event{
			{Type: "output", Data: "single line stdout"},
			{Type: "output", Data: "# some more output"},
			{Type: "output", Data: "\tfile_test.go:11: Error message"},
			{Type: "output", Data: "\tfile_test.go:12: Longer"},
			{Type: "output", Data: "\t\terror"},
			{Type: "output", Data: "\t\tmessage."},
		},
	},
}

func TestParse(t *testing.T) {
	matchRegex := compileMatch(t)
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			if !matchRegex.MatchString(test.name) || test.input == "" {
				t.SkipNow()
			}
			testParse(t, test.name, test.input, test.events)
		})
	}
}

func testParse(t *testing.T, name, input string, want []gtr.Event) {
	got, err := Parse(strings.NewReader(input))
	if err != nil {
		t.Errorf("Parse(%s) error: %v", name, err)
		return
	}

	if diff := cmp.Diff(got, want); diff != "" {
		t.Errorf("Parse returned unexpected events, input:\n%s\ndiff (-got, +want):\n%v", input, diff)
	}
}

func compileMatch(t *testing.T) *regexp.Regexp {
	rx, err := regexp.Compile(*matchTest)
	if err != nil {
		t.Fatalf("Error compiling -match flag %q: %v", *matchTest, err)
	}
	return rx
}

func inp(lines ...string) string {
	return strings.Join(lines, "\n")
}
<commit_msg>parser/gotest: refactor tests to run one test per input line<commit_after>package gotest

import (
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/jstemmer/go-junit-report/v2/pkg/gtr"

	"github.com/google/go-cmp/cmp"
)

type parseLineTest struct {
	input  string
	events interface{}
}

var parseLineTests = []parseLineTest{
	{
		"=== RUN TestOne",
		gtr.Event{Type: "run_test", Name: "TestOne"},
	},
	{
		"=== RUN   TestTwo/Subtest",
		gtr.Event{Type: "run_test", Name: "TestTwo/Subtest"},
	},
	{
		"=== PAUSE TestOne",
		gtr.Event{Type: "pause_test", Name: "TestOne"},
	},
	{
		"=== CONT  TestOne",
		gtr.Event{Type: "cont_test", Name: "TestOne"},
	},
	{
		"--- PASS: TestOne (12.34 seconds)",
		gtr.Event{Type: "end_test", Name: "TestOne", Result: "PASS", Duration: 12_340 * time.Millisecond},
	},
	{
		"    --- SKIP: TestOne/Subtest (0.00s)",
		gtr.Event{Type: "end_test", Name: "TestOne/Subtest", Result: "SKIP", Indent: 1},
	},
	{
		"        --- FAIL: TestOne/Subtest/#01 (0.35s)",
		gtr.Event{Type: "end_test", Name: "TestOne/Subtest/#01", Result: "FAIL", Duration: 350 * time.Millisecond, Indent: 2},
	},
	{
		"some text--- PASS: TestTwo (0.06 seconds)",
		[]gtr.Event{
			{Type: "output", Data: "some text"},
			{Type: "end_test", Name: "TestTwo", Result: "PASS", Duration: 60 * time.Millisecond},
		},
	},
	{
		"PASS",
		gtr.Event{Type: "status", Result: "PASS"},
	},
	{
		"FAIL",
		gtr.Event{Type: "status", Result: "FAIL"},
	},
	{
		"SKIP",
		gtr.Event{Type: "status", Result: "SKIP"},
	},
	{
		"ok      package/name/ok 0.100s",
		gtr.Event{Type: "summary", Name: "package/name/ok", Result: "ok", Duration: 100 * time.Millisecond},
	},
	{
		"FAIL    package/name/failing [build failed]",
		gtr.Event{Type: "summary", Name: "package/name/failing", Result: "FAIL", Data: "[build failed]"},
	},
	{
		"FAIL    package/other/failing [setup failed]",
		gtr.Event{Type: "summary", Name: "package/other/failing", Result: "FAIL", Data: "[setup failed]"},
	},
	{
		"ok package/other     (cached)",
		gtr.Event{Type: "summary", Name: "package/other", Result: "ok", Data: "(cached)"},
	},
	{
		"ok  	package/name 0.400s  coverage: 10.0% of statements",
		gtr.Event{Type: "summary", Name: "package/name", Result: "ok", Duration: 400 * time.Millisecond, CovPct: 10},
	},
	{
		"ok  	package/name 4.200s  coverage: 99.8% of statements in fmt, encoding/xml",
		gtr.Event{Type: "summary", Name: "package/name", Result: "ok", Duration: 4200 * time.Millisecond, CovPct: 99.8, CovPackages: []string{"fmt", "encoding/xml"}},
	},
	{
		"?   	package/name	[no test files]",
		gtr.Event{Type: "summary", Name: "package/name", Result: "?", Data: "[no test files]"},
	},
	{
		"ok  	package/name	0.001s [no tests to run]",
		gtr.Event{Type: "summary", Name: "package/name", Result: "ok", Duration: 1 * time.Millisecond, Data: "[no tests to run]"},
	},
	{
		"ok  	package/name	(cached) [no tests to run]",
		gtr.Event{Type: "summary", Name: "package/name", Result: "ok", Data: "(cached) [no tests to run]"},
	},
	{
		"coverage: 10% of statements",
		gtr.Event{Type: "coverage", CovPct: 10},
	},
	{
		"coverage: 10% of statements in fmt, encoding/xml",
		gtr.Event{Type: "coverage", CovPct: 10, CovPackages: []string{"fmt", "encoding/xml"}},
	},
	{
		"coverage: 13.37% of statements",
		gtr.Event{Type: "coverage", CovPct: 13.37},
	},
	{
		"coverage: 99.8% of statements in fmt, encoding/xml",
		gtr.Event{Type: "coverage", CovPct: 99.8, CovPackages: []string{"fmt", "encoding/xml"}},
	},
	{
		"BenchmarkOne-8                     2000000	       604 ns/op",
		gtr.Event{Type: "benchmark", Name: "BenchmarkOne", Iterations: 2_000_000, NsPerOp: 604},
	},
	{
		"BenchmarkTwo-16 30000	52568 ns/op	24879 B/op	494 allocs/op",
		gtr.Event{Type: "benchmark", Name: "BenchmarkTwo", Iterations: 30_000, NsPerOp: 52_568, BytesPerOp: 24_879, AllocsPerOp: 494},
	},
	{
		"BenchmarkThree      2000000000	         0.26 ns/op",
		gtr.Event{Type: "benchmark", Name: "BenchmarkThree", Iterations: 2_000_000_000, NsPerOp: 0.26},
	},
	{
		"BenchmarkFour-8         	   10000	    104427 ns/op	  95.76 MB/s	   40629 B/op	       5 allocs/op",
		gtr.Event{Type: "benchmark", Name: "BenchmarkFour", Iterations: 10_000, NsPerOp: 104_427, MBPerSec: 95.76, BytesPerOp: 40_629, AllocsPerOp: 5},
	},
	{
		"# package/name/failing1",
		gtr.Event{Type: "build_output", Name: "package/name/failing1"},
	},
	{
		"# package/name/failing2 [package/name/failing2.test]",
		gtr.Event{Type: "build_output", Name: "package/name/failing2"},
	},
	{
		"single line stdout",
		gtr.Event{Type: "output", Data: "single line stdout"},
	},
	{
		"# some more output",
		gtr.Event{Type: "output", Data: "# some more output"},
	},
	{
		"\tfile_test.go:11: Error message",
		gtr.Event{Type: "output", Data: "\tfile_test.go:11: Error message"},
	},
	{
		"\tfile_test.go:12: Longer",
		gtr.Event{Type: "output", Data: "\tfile_test.go:12: Longer"},
	},
	{
		"\t\terror",
		gtr.Event{Type: "output", Data: "\t\terror"},
	},
	{
		"\t\tmessage.",
		gtr.Event{Type: "output", Data: "\t\tmessage."},
	},
}

func TestParseLine(t *testing.T) {
	for i, test := range parseLineTests {
		var want []gtr.Event
		switch e := test.events.(type) {
		case gtr.Event:
			want = []gtr.Event{e}
		case []gtr.Event:
			want = e
		default:
			panic("invalid events type")
		}

		var types []string
		for _, e := range want {
			types = append(types, e.Type)
		}

		name := fmt.Sprintf("%d %s", i+1, strings.Join(types, ","))
		t.Run(name, func(t *testing.T) {
			testParseLine(t, &parser{}, test.input, want)
		})
	}
}

func testParseLine(t *testing.T, parser *parser, input string, want []gtr.Event) {
	parser.parseLine(input)
	got := parser.events
	if diff := cmp.Diff(got, want); diff != "" {
		t.Errorf("parseLine(%q) returned unexpected events, diff (-got, +want):\n%v", input, diff)
	}
}
<|endoftext|>