<commit_before>// Simple console progress bars
package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"
)

// Current version
const Version = "1.0.23"

const (
	// Default refresh rate - 200ms
	DEFAULT_REFRESH_RATE = time.Millisecond * 200
	FORMAT               = "[=>-]"
)

// DEPRECATED
// variables for backward compatibility, from now do not work
// use pb.Format and pb.SetRefreshRate
var (
	DefaultRefreshRate                         = DEFAULT_REFRESH_RATE
	BarStart, BarEnd, Empty, Current, CurrentN string
)

// Create new progress bar object
func New(total int) *ProgressBar {
	return New64(int64(total))
}

// Create new progress bar object using int64 as total
func New64(total int64) *ProgressBar {
	pb := &ProgressBar{
		Total:           total,
		RefreshRate:     DEFAULT_REFRESH_RATE,
		ShowPercent:     true,
		ShowCounters:    true,
		ShowBar:         true,
		ShowTimeLeft:    true,
		ShowElapsedTime: false,
		ShowFinalTime:   true,
		Units:           U_NO,
		ManualUpdate:    false,
		finish:          make(chan struct{}),
	}
	return pb.Format(FORMAT)
}

// Create new object and start
func StartNew(total int) *ProgressBar {
	return New(total).Start()
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current  int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)
	previous int64

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	ShowFinalTime, ShowElapsedTime   bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            Units
	Width                            int
	ForceWidth                       bool
	ManualUpdate                     bool
	AutoStat                         bool

	// Default width for the time box.
	UnitsWidth   int
	TimeBoxWidth int

	finishOnce sync.Once //Guards isFinish
	finish     chan struct{}
	isFinish   bool

	startTime  time.Time
	startValue int64

	changeTime time.Time

	prefix, postfix string

	mu        sync.Mutex
	lastPrint string

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string

	AlwaysUpdate bool
}

// Start print
func (pb *ProgressBar) Start() *ProgressBar {
	pb.startTime = time.Now()
	pb.startValue = atomic.LoadInt64(&pb.current)
	if atomic.LoadInt64(&pb.Total) == 0 {
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
		pb.AutoStat = false
	}
	if !pb.ManualUpdate {
		pb.Update() // Initial printing of the bar before running the bar refresher.
		go pb.refresher()
	}
	return pb
}

// Increment current value
func (pb *ProgressBar) Increment() int {
	return pb.Add(1)
}

// Get current value
func (pb *ProgressBar) Get() int64 {
	c := atomic.LoadInt64(&pb.current)
	return c
}

// Set current value
func (pb *ProgressBar) Set(current int) *ProgressBar {
	return pb.Set64(int64(current))
}

// Set64 sets the current value as int64
func (pb *ProgressBar) Set64(current int64) *ProgressBar {
	atomic.StoreInt64(&pb.current, current)
	return pb
}

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(pb.Add64(int64(add)))
}

func (pb *ProgressBar) Add64(add int64) int64 {
	return atomic.AddInt64(&pb.current, add)
}

// Set prefix string
func (pb *ProgressBar) Prefix(prefix string) *ProgressBar {
	pb.prefix = prefix
	return pb
}

// Set postfix string
func (pb *ProgressBar) Postfix(postfix string) *ProgressBar {
	pb.postfix = postfix
	return pb
}

// Set custom format for bar
// Example: bar.Format("[=>_]")
// Example: bar.Format("[\x00=\x00>\x00-\x00]") // \x00 is the delimiter
func (pb *ProgressBar) Format(format string) *ProgressBar {
	var formatEntries []string
	if utf8.RuneCountInString(format) == 5 {
		formatEntries = strings.Split(format, "")
	} else {
		formatEntries = strings.Split(format, "\x00")
	}
	if len(formatEntries) == 5 {
		pb.BarStart = formatEntries[0]
		pb.BarEnd = formatEntries[4]
		pb.Empty = formatEntries[3]
		pb.Current = formatEntries[1]
		pb.CurrentN = formatEntries[2]
	}
	return pb
}

// Set bar refresh rate
func (pb *ProgressBar) SetRefreshRate(rate time.Duration) *ProgressBar {
	pb.RefreshRate = rate
	return pb
}

// Set units
// bar.SetUnits(U_NO) - by default
// bar.SetUnits(U_BYTES) - for Mb, Kb, etc
func (pb *ProgressBar) SetUnits(units Units) *ProgressBar {
	pb.Units = units
	return pb
}

// Set max width, if width is bigger than terminal width, will be ignored
func (pb *ProgressBar) SetMaxWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = false
	return pb
}

// Set bar width
func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = true
	return pb
}

// End print
func (pb *ProgressBar) Finish() {
	//Protect multiple calls
	pb.finishOnce.Do(func() {
		close(pb.finish)
		pb.write(atomic.LoadInt64(&pb.Total), atomic.LoadInt64(&pb.current))
		pb.mu.Lock()
		defer pb.mu.Unlock()
		switch {
		case pb.Output != nil:
			fmt.Fprintln(pb.Output)
		case !pb.NotPrint:
			fmt.Println()
		}
		pb.isFinish = true
	})
}

// IsFinished return boolean
func (pb *ProgressBar) IsFinished() bool {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.isFinish
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	if pb.Output != nil {
		fmt.Fprintln(pb.Output, str)
	} else {
		fmt.Println(str)
	}
}

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// Create new proxy reader over bar
// Takes io.Reader or io.ReadCloser
func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
	return &Reader{r, pb}
}

func (pb *ProgressBar) write(total, current int64) {
	width := pb.GetWidth()

	var percentBox, countersBox, timeLeftBox, timeSpentBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		var percent float64
		if total > 0 {
			percent = float64(current) / (float64(total) / float64(100))
		} else {
			percent = float64(current) / float64(100)
		}
		percentBox = fmt.Sprintf(" %6.02f%%", percent)
	}

	// counters
	if pb.ShowCounters {
		current := Format(current).To(pb.Units).Width(pb.UnitsWidth)
		if total > 0 {
			totalS := Format(total).To(pb.Units).Width(pb.UnitsWidth)
			countersBox = fmt.Sprintf(" %s / %s ", current, totalS)
		} else {
			countersBox = fmt.Sprintf(" %s / ? ", current)
		}
	}

	// time left
	pb.mu.Lock()
	currentFromStart := current - pb.startValue
	fromStart := time.Now().Sub(pb.startTime)
	lastChangeTime := pb.changeTime
	fromChange := lastChangeTime.Sub(pb.startTime)
	pb.mu.Unlock()

	if pb.ShowElapsedTime {
		timeSpentBox = fmt.Sprintf(" %s ", (fromStart/time.Second)*time.Second)
	}

	select {
	case <-pb.finish:
		if pb.ShowFinalTime {
			var left time.Duration
			left = (fromStart / time.Second) * time.Second
			timeLeftBox = fmt.Sprintf(" %s", left.String())
		}
	default:
		if pb.ShowTimeLeft && currentFromStart > 0 {
			perEntry := fromChange / time.Duration(currentFromStart)
			var left time.Duration
			if total > 0 {
				left = time.Duration(total-currentFromStart) * perEntry
				left -= time.Since(lastChangeTime)
				left = (left / time.Second) * time.Second
			} else {
				left = time.Duration(currentFromStart) * perEntry
				left = (left / time.Second) * time.Second
			}
			if left > 0 {
				timeLeft := Format(int64(left)).To(U_DURATION).String()
				timeLeftBox = fmt.Sprintf(" %s", timeLeft)
			}
		}
	}

	if len(timeLeftBox) < pb.TimeBoxWidth {
		timeLeftBox = fmt.Sprintf("%s%s", strings.Repeat(" ", pb.TimeBoxWidth-len(timeLeftBox)), timeLeftBox)
	}

	// speed
	if pb.ShowSpeed && currentFromStart > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(currentFromStart) / (float64(fromStart) / float64(time.Second))
		speedBox = " " + Format(int64(speed)).To(pb.Units).Width(pb.UnitsWidth).PerSec().String()
	}

	barWidth := escapeAwareRuneCountInString(countersBox + pb.BarStart + pb.BarEnd + percentBox + timeSpentBox + timeLeftBox + speedBox + pb.prefix + pb.postfix)
	// bar
	if pb.ShowBar {
		size := width - barWidth
		if size > 0 {
			if total > 0 {
				curSize := int(math.Ceil((float64(current) / float64(total)) * float64(size)))
				emptySize := size - curSize
				barBox = pb.BarStart
				if emptySize < 0 {
					emptySize = 0
				}
				if curSize > size {
					curSize = size
				}

				cursorLen := escapeAwareRuneCountInString(pb.Current)
				if emptySize <= 0 {
					barBox += strings.Repeat(pb.Current, curSize/cursorLen)
				} else if curSize > 0 {
					cursorEndLen := escapeAwareRuneCountInString(pb.CurrentN)
					cursorRepetitions := (curSize - cursorEndLen) / cursorLen
					barBox += strings.Repeat(pb.Current, cursorRepetitions)
					barBox += pb.CurrentN
				}

				emptyLen := escapeAwareRuneCountInString(pb.Empty)
				barBox += strings.Repeat(pb.Empty, emptySize/emptyLen)
				barBox += pb.BarEnd
			} else {
				pos := size - int(current)%int(size)
				barBox = pb.BarStart
				if pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, pos-1)
				}
				barBox += pb.Current
				if size-pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, size-pos-1)
				}
				barBox += pb.BarEnd
			}
		}
	}

	// check len
	out = pb.prefix + timeSpentBox + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix
	if cl := escapeAwareRuneCountInString(out); cl < width {
		end = strings.Repeat(" ", width-cl)
	}

	// and print!
	pb.mu.Lock()
	defer pb.mu.Unlock()
	pb.lastPrint = out + end
	isFinish := pb.isFinish

	switch {
	case isFinish:
		return
	case pb.Output != nil:
		fmt.Fprint(pb.Output, "\r"+out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

// GetTerminalWidth - returns terminal width for all platforms.
func GetTerminalWidth() (int, error) {
	return terminalWidth()
}

func (pb *ProgressBar) GetWidth() int {
	if pb.ForceWidth {
		return pb.Width
	}

	width := pb.Width
	termWidth, _ := terminalWidth()
	if width == 0 || termWidth <= width {
		width = termWidth
	}

	return width
}

// Write the current state of the progressbar
func (pb *ProgressBar) Update() {
	c := atomic.LoadInt64(&pb.current)
	p := atomic.LoadInt64(&pb.previous)
	t := atomic.LoadInt64(&pb.Total)
	if p != c {
		pb.mu.Lock()
		pb.changeTime = time.Now()
		pb.mu.Unlock()
		atomic.StoreInt64(&pb.previous, c)
	}
	pb.write(t, c)
	if pb.AutoStat {
		if c == 0 {
			pb.startTime = time.Now()
			pb.startValue = 0
		} else if c >= t && pb.isFinish != true {
			pb.Finish()
		}
	}
}

// String return the last bar print
func (pb *ProgressBar) String() string {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.lastPrint
}

// SetTotal atomically sets new total count
func (pb *ProgressBar) SetTotal(total int) *ProgressBar {
	return pb.SetTotal64(int64(total))
}

// SetTotal64 atomically sets new total count
func (pb *ProgressBar) SetTotal64(total int64) *ProgressBar {
	atomic.StoreInt64(&pb.Total, total)
	return pb
}

// Reset bar and set new total count
// Does effect only on finished bar
func (pb *ProgressBar) Reset(total int) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	if pb.isFinish {
		pb.SetTotal(total).Set(0)
		atomic.StoreInt64(&pb.previous, 0)
	}
	return pb
}

// Internal loop for refreshing the progressbar
func (pb *ProgressBar) refresher() {
	for {
		select {
		case <-pb.finish:
			return
		case <-time.After(pb.RefreshRate):
			pb.Update()
		}
	}
}
<commit_msg>1.0.24<commit_after>// Simple console progress bars
package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"
)

// Current version
const Version = "1.0.24"

const (
	// Default refresh rate - 200ms
	DEFAULT_REFRESH_RATE = time.Millisecond * 200
	FORMAT               = "[=>-]"
)

// DEPRECATED
// variables for backward compatibility, from now do not work
// use pb.Format and pb.SetRefreshRate
var (
	DefaultRefreshRate                         = DEFAULT_REFRESH_RATE
	BarStart, BarEnd, Empty, Current, CurrentN string
)

// Create new progress bar object
func New(total int) *ProgressBar {
	return New64(int64(total))
}

// Create new progress bar object using int64 as total
func New64(total int64) *ProgressBar {
	pb := &ProgressBar{
		Total:           total,
		RefreshRate:     DEFAULT_REFRESH_RATE,
		ShowPercent:     true,
		ShowCounters:    true,
		ShowBar:         true,
		ShowTimeLeft:    true,
		ShowElapsedTime: false,
		ShowFinalTime:   true,
		Units:           U_NO,
		ManualUpdate:    false,
		finish:          make(chan struct{}),
	}
	return pb.Format(FORMAT)
}

// Create new object and start
func StartNew(total int) *ProgressBar {
	return New(total).Start()
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current  int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)
	previous int64

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	ShowFinalTime, ShowElapsedTime   bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            Units
	Width                            int
	ForceWidth                       bool
	ManualUpdate                     bool
	AutoStat                         bool

	// Default width for the time box.
	UnitsWidth   int
	TimeBoxWidth int

	finishOnce sync.Once //Guards isFinish
	finish     chan struct{}
	isFinish   bool

	startTime  time.Time
	startValue int64

	changeTime time.Time

	prefix, postfix string

	mu        sync.Mutex
	lastPrint string

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string

	AlwaysUpdate bool
}

// Start print
func (pb *ProgressBar) Start() *ProgressBar {
	pb.startTime = time.Now()
	pb.startValue = atomic.LoadInt64(&pb.current)
	if atomic.LoadInt64(&pb.Total) == 0 {
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
		pb.AutoStat = false
	}
	if !pb.ManualUpdate {
		pb.Update() // Initial printing of the bar before running the bar refresher.
		go pb.refresher()
	}
	return pb
}

// Increment current value
func (pb *ProgressBar) Increment() int {
	return pb.Add(1)
}

// Get current value
func (pb *ProgressBar) Get() int64 {
	c := atomic.LoadInt64(&pb.current)
	return c
}

// Set current value
func (pb *ProgressBar) Set(current int) *ProgressBar {
	return pb.Set64(int64(current))
}

// Set64 sets the current value as int64
func (pb *ProgressBar) Set64(current int64) *ProgressBar {
	atomic.StoreInt64(&pb.current, current)
	return pb
}

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(pb.Add64(int64(add)))
}

func (pb *ProgressBar) Add64(add int64) int64 {
	return atomic.AddInt64(&pb.current, add)
}

// Set prefix string
func (pb *ProgressBar) Prefix(prefix string) *ProgressBar {
	pb.prefix = prefix
	return pb
}

// Set postfix string
func (pb *ProgressBar) Postfix(postfix string) *ProgressBar {
	pb.postfix = postfix
	return pb
}

// Set custom format for bar
// Example: bar.Format("[=>_]")
// Example: bar.Format("[\x00=\x00>\x00-\x00]") // \x00 is the delimiter
func (pb *ProgressBar) Format(format string) *ProgressBar {
	var formatEntries []string
	if utf8.RuneCountInString(format) == 5 {
		formatEntries = strings.Split(format, "")
	} else {
		formatEntries = strings.Split(format, "\x00")
	}
	if len(formatEntries) == 5 {
		pb.BarStart = formatEntries[0]
		pb.BarEnd = formatEntries[4]
		pb.Empty = formatEntries[3]
		pb.Current = formatEntries[1]
		pb.CurrentN = formatEntries[2]
	}
	return pb
}

// Set bar refresh rate
func (pb *ProgressBar) SetRefreshRate(rate time.Duration) *ProgressBar {
	pb.RefreshRate = rate
	return pb
}

// Set units
// bar.SetUnits(U_NO) - by default
// bar.SetUnits(U_BYTES) - for Mb, Kb, etc
func (pb *ProgressBar) SetUnits(units Units) *ProgressBar {
	pb.Units = units
	return pb
}

// Set max width, if width is bigger than terminal width, will be ignored
func (pb *ProgressBar) SetMaxWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = false
	return pb
}

// Set bar width
func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = true
	return pb
}

// End print
func (pb *ProgressBar) Finish() {
	//Protect multiple calls
	pb.finishOnce.Do(func() {
		close(pb.finish)
		pb.write(atomic.LoadInt64(&pb.Total), atomic.LoadInt64(&pb.current))
		pb.mu.Lock()
		defer pb.mu.Unlock()
		switch {
		case pb.Output != nil:
			fmt.Fprintln(pb.Output)
		case !pb.NotPrint:
			fmt.Println()
		}
		pb.isFinish = true
	})
}

// IsFinished return boolean
func (pb *ProgressBar) IsFinished() bool {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.isFinish
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	if pb.Output != nil {
		fmt.Fprintln(pb.Output, str)
	} else {
		fmt.Println(str)
	}
}

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// Create new proxy reader over bar
// Takes io.Reader or io.ReadCloser
func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
	return &Reader{r, pb}
}

func (pb *ProgressBar) write(total, current int64) {
	width := pb.GetWidth()

	var percentBox, countersBox, timeLeftBox, timeSpentBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		var percent float64
		if total > 0 {
			percent = float64(current) / (float64(total) / float64(100))
		} else {
			percent = float64(current) / float64(100)
		}
		percentBox = fmt.Sprintf(" %6.02f%%", percent)
	}

	// counters
	if pb.ShowCounters {
		current := Format(current).To(pb.Units).Width(pb.UnitsWidth)
		if total > 0 {
			totalS := Format(total).To(pb.Units).Width(pb.UnitsWidth)
			countersBox = fmt.Sprintf(" %s / %s ", current, totalS)
		} else {
			countersBox = fmt.Sprintf(" %s / ? ", current)
		}
	}

	// time left
	pb.mu.Lock()
	currentFromStart := current - pb.startValue
	fromStart := time.Now().Sub(pb.startTime)
	lastChangeTime := pb.changeTime
	fromChange := lastChangeTime.Sub(pb.startTime)
	pb.mu.Unlock()

	if pb.ShowElapsedTime {
		timeSpentBox = fmt.Sprintf(" %s ", (fromStart/time.Second)*time.Second)
	}

	select {
	case <-pb.finish:
		if pb.ShowFinalTime {
			var left time.Duration
			left = (fromStart / time.Second) * time.Second
			timeLeftBox = fmt.Sprintf(" %s", left.String())
		}
	default:
		if pb.ShowTimeLeft && currentFromStart > 0 {
			perEntry := fromChange / time.Duration(currentFromStart)
			var left time.Duration
			if total > 0 {
				left = time.Duration(total-currentFromStart) * perEntry
				left -= time.Since(lastChangeTime)
				left = (left / time.Second) * time.Second
			} else {
				left = time.Duration(currentFromStart) * perEntry
				left = (left / time.Second) * time.Second
			}
			if left > 0 {
				timeLeft := Format(int64(left)).To(U_DURATION).String()
				timeLeftBox = fmt.Sprintf(" %s", timeLeft)
			}
		}
	}

	if len(timeLeftBox) < pb.TimeBoxWidth {
		timeLeftBox = fmt.Sprintf("%s%s", strings.Repeat(" ", pb.TimeBoxWidth-len(timeLeftBox)), timeLeftBox)
	}

	// speed
	if pb.ShowSpeed && currentFromStart > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(currentFromStart) / (float64(fromStart) / float64(time.Second))
		speedBox = " " + Format(int64(speed)).To(pb.Units).Width(pb.UnitsWidth).PerSec().String()
	}

	barWidth := escapeAwareRuneCountInString(countersBox + pb.BarStart + pb.BarEnd + percentBox + timeSpentBox + timeLeftBox + speedBox + pb.prefix + pb.postfix)
	// bar
	if pb.ShowBar {
		size := width - barWidth
		if size > 0 {
			if total > 0 {
				curSize := int(math.Ceil((float64(current) / float64(total)) * float64(size)))
				emptySize := size - curSize
				barBox = pb.BarStart
				if emptySize < 0 {
					emptySize = 0
				}
				if curSize > size {
					curSize = size
				}

				cursorLen := escapeAwareRuneCountInString(pb.Current)
				if emptySize <= 0 {
					barBox += strings.Repeat(pb.Current, curSize/cursorLen)
				} else if curSize > 0 {
					cursorEndLen := escapeAwareRuneCountInString(pb.CurrentN)
					cursorRepetitions := (curSize - cursorEndLen) / cursorLen
					barBox += strings.Repeat(pb.Current, cursorRepetitions)
					barBox += pb.CurrentN
				}

				emptyLen := escapeAwareRuneCountInString(pb.Empty)
				barBox += strings.Repeat(pb.Empty, emptySize/emptyLen)
				barBox += pb.BarEnd
			} else {
				pos := size - int(current)%int(size)
				barBox = pb.BarStart
				if pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, pos-1)
				}
				barBox += pb.Current
				if size-pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, size-pos-1)
				}
				barBox += pb.BarEnd
			}
		}
	}

	// check len
	out = pb.prefix + timeSpentBox + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix
	if cl := escapeAwareRuneCountInString(out); cl < width {
		end = strings.Repeat(" ", width-cl)
	}

	// and print!
	pb.mu.Lock()
	defer pb.mu.Unlock()
	pb.lastPrint = out + end
	isFinish := pb.isFinish

	switch {
	case isFinish:
		return
	case pb.Output != nil:
		fmt.Fprint(pb.Output, "\r"+out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

// GetTerminalWidth - returns terminal width for all platforms.
func GetTerminalWidth() (int, error) {
	return terminalWidth()
}

func (pb *ProgressBar) GetWidth() int {
	if pb.ForceWidth {
		return pb.Width
	}

	width := pb.Width
	termWidth, _ := terminalWidth()
	if width == 0 || termWidth <= width {
		width = termWidth
	}

	return width
}

// Write the current state of the progressbar
func (pb *ProgressBar) Update() {
	c := atomic.LoadInt64(&pb.current)
	p := atomic.LoadInt64(&pb.previous)
	t := atomic.LoadInt64(&pb.Total)
	if p != c {
		pb.mu.Lock()
		pb.changeTime = time.Now()
		pb.mu.Unlock()
		atomic.StoreInt64(&pb.previous, c)
	}
	pb.write(t, c)
	if pb.AutoStat {
		if c == 0 {
			pb.startTime = time.Now()
			pb.startValue = 0
		} else if c >= t && pb.isFinish != true {
			pb.Finish()
		}
	}
}

// String return the last bar print
func (pb *ProgressBar) String() string {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.lastPrint
}

// SetTotal atomically sets new total count
func (pb *ProgressBar) SetTotal(total int) *ProgressBar {
	return pb.SetTotal64(int64(total))
}

// SetTotal64 atomically sets new total count
func (pb *ProgressBar) SetTotal64(total int64) *ProgressBar {
	atomic.StoreInt64(&pb.Total, total)
	return pb
}

// Reset bar and set new total count
// Does effect only on finished bar
func (pb *ProgressBar) Reset(total int) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	if pb.isFinish {
		pb.SetTotal(total).Set(0)
		atomic.StoreInt64(&pb.previous, 0)
	}
	return pb
}

// Internal loop for refreshing the progressbar
func (pb *ProgressBar) refresher() {
	for {
		select {
		case <-pb.finish:
			return
		case <-time.After(pb.RefreshRate):
			pb.Update()
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
)

const baseUrl = "http://localhost:8989/api/"

func main() {
	switch os.Args[1] {
	case "refresh":
		refreshSeries()
	case "search":
		search()
	case "list":
		list()
	}
}

type command struct {
	Name string
}

func search() {

}

func list() {
	apiKey := readApiKey()
	resp, err := http.Get(baseUrl + "/series?apikey=" + apiKey)

	check(err)

	fmt.Printf("Response:\n%s", getBody(resp))
}

func refreshSeries() {
	sendCommand("RefreshSeries")
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func readApiKey() string {
	key, err := ioutil.ReadFile("api_key")
	check(err)
	return strings.TrimSpace(string(key))
}

func sendCommand(name string) {
	apiKey := readApiKey()
	endpoint := baseUrl + "command?apikey=" + apiKey

	b := new(bytes.Buffer)
	json.NewEncoder(b).Encode(command{Name: name})
	resp, err := http.Post(endpoint, "application/json", b)

	check(err)

	fmt.Printf("Response:\n%s", getBody(resp))
}

func getBody(resp *http.Response) []byte {
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)

	check(err)

	return body
}
<commit_msg>Cleaning up some code, adding series search method<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
)

const baseUrl = "http://localhost:8989/api/"

func main() {
	switch os.Args[1] {
	case "refresh":
		refreshSeries()
	case "search":
		search(os.Args[2])
	case "list":
		list()
	}
}

type command struct {
	Name string
}

func search(seriesId string) {

}

func list() {
	resp, err := http.Get(getUrl("series"))

	check(err)

	fmt.Printf("Response:\n%s", getBody(resp))
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func refreshSeries() {
	b := new(bytes.Buffer)
	json.NewEncoder(b).Encode(command{Name: "RefreshSeries"})
	resp, err := http.Post(getUrl("command"), "application/json", b)

	check(err)

	fmt.Printf("Response:\n%s", getBody(resp))
}

func getBody(resp *http.Response) []byte {
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)

	check(err)

	return body
}

func getUrl(endpoint string) string {
	apiKey := readApiKey()
	return baseUrl + endpoint + "?apikey=" + apiKey
}

func readApiKey() string {
	key, err := ioutil.ReadFile("api_key")
	check(err)
	return strings.TrimSpace(string(key))
}
<|endoftext|>


<commit_before>package zeroless

import "testing"

func TestPortUnderRange(t *testing.T) {
	port := 1023
	_, err := NewServer(port).Push()

	if err == nil {
		t.Error("Port", port, "is under 1024")
	}
}

func TestPortOnRange(t *testing.T) {
	port := 1024
	_, err := NewServer(port).Push()

	if err != nil {
		t.Error("Port", port, "is not on range")
	}
}

func TestPortAfterRange(t *testing.T) {
	port := 65536
	_, err := NewServer(port).Push()

	if err == nil {
		t.Error("Port", port, "is after 65535")
	}
}
<commit_msg>Renamed some tests and increased the coverage for port is out of range checks.<commit_after>package zeroless

import "testing"

func TestPortBellowRange(t *testing.T) {
	port := 1023
	_, err := NewServer(port).Push()

	if err == nil {
		t.Error("Port", port, "is bellow 1024")
	}

	client := NewClient()
	client.ConnectLocal(port)
	_, err = client.Push()

	if err == nil {
		t.Error("Port", port, "is bellow 1024")
	}
}

func TestPortWithinRange(t *testing.T) {
	port := 1024
	_, err := NewServer(port).Push()

	if err != nil {
		t.Error("Port", port, "is within range")
	}

	client := NewClient()
	client.ConnectLocal(port)
	_, err = client.Push()

	if err != nil {
		t.Error("Port", port, "is within range")
	}
}

func TestPortAboveRange(t *testing.T) {
	port := 65536
	_, err := NewServer(port).Push()

	if err == nil {
		t.Error("Port", port, "is above 65535")
	}

	client := NewClient()
	client.ConnectLocal(port)
	_, err = client.Push()

	if err == nil {
		t.Error("Port", port, "is above 65535")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
//  "io/ioutil"
//  "strconv"
)

const (
	endpoint = "http://pokeapi.co/api/v1"
)

type Game struct {
  Name string `json:"name"`
  Id int `json:"id"`
  Resource_uri string `json:"resource_uri"`
  Created string `json:"created"`
  Modified string `json:"modified"`
  Release_year int `json:"release"year"`
  Generation int `json:"generation"`
}

type Ability struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
}

type Description struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
}

type EggGroup struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
}

type Evolution struct {
	Level       int    `json:"level"`
	Method      string `json:"method"`
	Resouce_uri string `json:"resource_uri"`
	To          string `json:"to"`
}

type Move struct {
	Learn_type   string `json:"learn_type"`
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
	Level        int    `json:"level"`
}

type Sprite struct {
	Name        string `json:"name"`
	Resouce_uri string `json:"resource_uri"`
}

type Type struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
}

type Pokemon struct {
	Name              string        `json:"name"`
	National_id       int           `json:"national_id"`
	Resource_uri      string        `json:"resource_uri"`
	Created           string        `json:"created"`
	Modified          string        `json:"modified"`
	Abilites          []Ability     `json:"abilities"`
	Egg_groups        []EggGroup    `json:"egg_groups"`
	Evolutions        []Evolution   `json:"evolutions"`
	Descriptions      []Description `json:"descriptions"`
	Moves             []Move        `json:"moves"`
	Types             []Type        `json:"types"`
	Catch_rate        int           `json:"catch_rate"`
	Species           string        `json:"species"`
	Hp                int           `json:"hp"`
	Attack            int           `json:"attack"`
	Defense           int           `json:"defense"`
	Sp_atk            int           `json:"sp_atk"`
	Sp_def            int           `json:"sp_def"`
	Speed             int           `json:"speed"`
	Egg_cycles        int           `json:"egg_cycles"`
	Ev_yield          string        `json:"ev_yield"`
	Exp               int           `json:"exp"`
	Growth_rate       string        `json:"growth_rate"`
	Happiness         int           `json:"happiness"`
	Height            string        `json:"height"`
	Male_female_ratio string        `json:"male_female_ratio"`
	Pkdx_id           int           `json:"pkdx_id"`
	Sprites           []Sprite      `json:"sprites"`
	Total             int           `json:"total"`
	Weight            string           `json:"weight"`
}

type Pokedex struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
	Created      string `json:"created"`
	Modified     string `json:"modified"`
	Pokemon      []Pokemon `json:"pokemon"`
}

func main() {
  _, err := getPokemon("bulbasaur")
  if err != nil {
    fmt.Println(err)
  }
}

func getPokemon(identifier string) (Pokemon, error) {
  url := endpoint + "/pokemon/" + identifier
  res, err := http.Get(url)
  if err != nil {
    return Pokemon{}, err
  }
  defer res.Body.Close()
  decoder := json.NewDecoder(res.Body)
  var pokemon Pokemon
  if err := decoder.Decode(&pokemon); err != nil {
    return Pokemon{}, err
  }
  fmt.Println(pokemon.Name, pokemon.Species)
  return pokemon, nil
}

func getGame(identifier string) (Game, error) {
  url := endpoint + "/game/" + identifier
  res, err := http.Get(url)
  if err != nil {
    return Game{}, err
  }
  defer res.Body.Close()
  decoder := json.NewDecoder(res.Body)
  var game Game
  if err := decoder.Decode(&game); err != nil {
    return Game{}, err
  }
  fmt.Println(game.Name,":",game.Generation)
  return game, nil
}
<commit_msg>made cleaner. separated http request part<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
)

const (
	endpoint = "http://pokeapi.co/api/v1"
)

type Game struct {
	Name         string `json:"name"`
	Id           int    `json:"id"`
	Resource_uri string `json:"resource_uri"`
	Created      string `json:"created"`
	Modified     string `json:"modified"`
	Release_year int    `json:"release"year"`
	Generation   int    `json:"generation"`
}

type Ability struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
}

type Description struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
}

type EggGroup struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
}

type Evolution struct {
	Level       int    `json:"level"`
	Method      string `json:"method"`
	Resouce_uri string `json:"resource_uri"`
	To          string `json:"to"`
}

type Move struct {
	Learn_type   string `json:"learn_type"`
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
	Level        int    `json:"level"`
}

type Sprite struct {
	Name        string `json:"name"`
	Resouce_uri string `json:"resource_uri"`
}

type Type struct {
	Name         string `json:"name"`
	Resource_uri string `json:"resource_uri"`
}

type Pokemon struct {
	Name              string        `json:"name"`
	National_id       int           `json:"national_id"`
	Resource_uri      string        `json:"resource_uri"`
	Created           string        `json:"created"`
	Modified          string        `json:"modified"`
	Abilites          []Ability     `json:"abilities"`
	Egg_groups        []EggGroup    `json:"egg_groups"`
	Evolutions        []Evolution   `json:"evolutions"`
	Descriptions      []Description `json:"descriptions"`
	Moves             []Move        `json:"moves"`
	Types             []Type        `json:"types"`
	Catch_rate        int           `json:"catch_rate"`
	Species           string        `json:"species"`
	Hp                int           `json:"hp"`
	Attack            int           `json:"attack"`
	Defense           int           `json:"defense"`
	Sp_atk            int           `json:"sp_atk"`
	Sp_def            int           `json:"sp_def"`
	Speed             int           `json:"speed"`
	Egg_cycles        int           `json:"egg_cycles"`
	Ev_yield          string        `json:"ev_yield"`
	Exp               int           `json:"exp"`
	Growth_rate       string        `json:"growth_rate"`
	Happiness         int           `json:"happiness"`
	Height            string        `json:"height"`
	Male_female_ratio string        `json:"male_female_ratio"`
	Pkdx_id           int           `json:"pkdx_id"`
	Sprites           []Sprite      `json:"sprites"`
	Total             int           `json:"total"`
	Weight            string        `json:"weight"`
}

type Pokedex struct {
	Name         string    `json:"name"`
	Resource_uri string    `json:"resource_uri"`
	Created      string    `json:"created"`
	Modified     string    `json:"modified"`
	Pokemon      []Pokemon `json:"pokemon"`
}

func main() {
	_, err := getPokedex("1")
	if err != nil {
		fmt.Println(err)
	}
}

// This function gets the JSON from the API and populates the value field
// which is passed by reference to it
func endpointRequest(url string, value interface{}) error {
	res, err := http.Get(url)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	decoder := json.NewDecoder(res.Body)
	return decoder.Decode(&value)
}

func getPokedex(identifier string) (pokedex Pokedex, err error) {
	url := endpoint + "/pokedex/" + identifier
	if err = endpointRequest(url, &pokedex); err != nil {
		return Pokedex{}, err
	}
	fmt.Println(pokedex.Name, pokedex.Resource_uri)
	return pokedex, nil
}
func getPokemon(identifier string) (pokemon Pokemon, err error) {
	url := endpoint + "/pokemon/" + identifier
	if err = endpointRequest(url, &pokemon); err != nil {
		return Pokemon{}, err
	}
	fmt.Println(pokemon.Name, pokemon.Species)
	return pokemon, nil
}

func getGame(identifier string) (game Game, err error) {
	url := endpoint + "/game/" + identifier
	if err = endpointRequest(url, &game); err != nil {
		return Game{}, err
	}
	fmt.Println(game.Name, ":", game.Generation)
	return game, nil
}
<|endoftext|>


<commit_before>package gotrail

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
)

func getBuckets() []*s3.Bucket {
	s3client := gets3client()
	var params *s3.ListBucketsInput
	resp, err := s3client.ListBuckets(params)

	if err != nil {
		fmt.Println(err.Error())
	}
	fmt.Println(resp)
	return resp.Buckets
}

func getObjects() {
	s3client := gets3client()
	params := &s3.ListObjectsInput{
		Bucket: aws.String("BucketName"), // Required
		// Delimiter:    aws.String("Delimiter"),
		// EncodingType: aws.String("EncodingType"),
		// Marker:       aws.String("Marker"),
		// MaxKeys:      aws.Int64(1),
		// Prefix:       aws.String("Prefix"),
	}
	resp, err := s3client.ListObjects(params)

	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		fmt.Println(err.Error())
		return
	}

	// Pretty-print the response data.
	fmt.Println(resp)
}
<commit_msg>Adding downloadObjects function<commit_after>package gotrail

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
)

func getBuckets() []*s3.Bucket {
	s3client := gets3client()
	var params *s3.ListBucketsInput
	resp, err := s3client.ListBuckets(params)

	if err != nil {
		fmt.Println(err.Error())
	}
	fmt.Println(resp)
	return resp.Buckets
}

func getObjects() {
	s3client := gets3client()
	params := &s3.GetObjectInput{
		Bucket: aws.String("BucketName"),
		Key:    aws.String("ObjectKey"),
	}
	resp, err := s3client.GetObject(params)

	if err != nil {
		fmt.Println(err.Error())
		return
	}

	// Pretty-print the response data.
	fmt.Println(resp)
}

func downloadObjects() {

}
<|endoftext|>


<commit_before>// +build !containers_image_storage_stub

package storage

import (
	"fmt"
	"testing"

	"github.com/containers/image/docker/reference"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	sha256digestHex = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
	sha256Digest2   = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
)

func TestTransportName(t *testing.T) {
	assert.Equal(t, "containers-storage", Transport.Name())
}

func TestTransportParseStoreReference(t *testing.T) {
	const digest3 = "sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"

	store := newStore(t)

	Transport.SetStore(nil)
	for _, c := range []struct{ input, expectedRef, expectedID string }{
		{"", "", ""}, // Empty input
		// Handling of the store prefix
		// FIXME? Should we be silently discarding input like this?
		{"[unterminated", "", ""},                                    // Unterminated store specifier
		{"[garbage]busybox", "docker.io/library/busybox:latest", ""}, // Store specifier is overridden by the store we pass to ParseStoreReference

		{"UPPERCASEISINVALID", "", ""},                                                   // Invalid single-component name
		{"sha256:" + sha256digestHex, "docker.io/library/sha256:" + sha256digestHex, ""}, // Valid single-component name; the hex part is not an ID unless it has a "@" prefix, so it looks like a tag
		// FIXME: This test is now incorrect, this should not fail _if the image ID matches_
		{sha256digestHex, "", ""},                    // Invalid single-component ID; not an ID without a "@" prefix, so it's parsed as a name, but names aren't allowed to look like IDs
		{"@" + sha256digestHex, "", sha256digestHex}, // Valid single-component ID
		// FIXME: This should fail it seems, everything else uses ref.digest only if ref.name != nil
		// {"@sha256:" + sha256digestHex, "", ""},                                      // Valid single-component digest
		// "aaaa", either a valid image ID prefix, or a short form of docker.io/library/aaaa, untested
		{"sha256:ab", "docker.io/library/sha256:ab", ""},                                   // Valid single-component name, explicit tag
		{"busybox", "docker.io/library/busybox:latest", ""},                                // Valid single-component name, implicit tag
		{"busybox:notlatest", "docker.io/library/busybox:notlatest", ""},                   // Valid single-component name, explicit tag
		{"docker.io/library/busybox:notlatest", "docker.io/library/busybox:notlatest", ""}, // Valid single-component name, everything explicit

		{"UPPERCASEISINVALID@" + sha256digestHex, "", ""},                                                // Invalid name in name@digestOrID
		{"busybox@ab", "", ""},                                                                           // Invalid ID in name@digestOrID
		{"busybox@", "", ""},                                                                             // Empty ID in name@digestOrID
		{"busybox@sha256:ab", "", ""},                                                                    // Invalid digest in name@digestOrID
		{"busybox@sha256:" + sha256digestHex, "docker.io/library/busybox@sha256:" + sha256digestHex, ""}, // Valid name@digest, no tag
		{"busybox@" + sha256digestHex, "docker.io/library/busybox:latest", sha256digestHex},              // Valid name@ID, implicit tag
		// "busybox@aaaa", a valid image ID prefix, untested
		{"busybox:notlatest@" + sha256digestHex, "docker.io/library/busybox:notlatest", sha256digestHex},                     // Valid name@ID, explicit tag
		{"docker.io/library/busybox:notlatest@" + sha256digestHex, "docker.io/library/busybox:notlatest", sha256digestHex},   // Valid name@ID, everything explicit
		{"docker.io/library/busybox:notlatest@" + sha256Digest2, "docker.io/library/busybox:notlatest@" + sha256Digest2, ""}, // Valid name:tag@digest, everything explicit

		{"busybox@sha256:" + sha256digestHex + "@ab", "", ""},                                                                                // Invalid ID in name@digest@ID
		{"busybox@ab@" + sha256digestHex, "", ""},                                                                                // Invalid digest in name@digest@ID
		{"busybox@@" + sha256digestHex, "", ""},                                                                                // Invalid digest in name@digest@ID
		{"busybox@" + sha256Digest2 + "@" + sha256digestHex, "docker.io/library/busybox@" + sha256Digest2, sha256digestHex},                                                         // name@digest@ID
		{"docker.io/library/busybox@" + sha256Digest2 + "@" + sha256digestHex, "docker.io/library/busybox@" + sha256Digest2, sha256digestHex},                                       // name@digest@ID, everything explicit
		{"docker.io/library/busybox:notlatest@sha256:" + sha256digestHex + "@" + sha256digestHex, "docker.io/library/busybox:notlatest@sha256:" + sha256digestHex, sha256digestHex}, // name:tag@digest@ID, everything explicit
		// FIXME: Is this supposed to work? the validation of idOrDigest seems to make this impossible.
		// "busybox@sha256:"+sha256digestHex+"@aaaa", a valid image ID prefix, untested
		// FIXME FIXME: two digests
		{"busybox:notlatest@" + sha256Digest2 + "@" + digest3 + "@" + sha256digestHex, "docker.io/library/busybox:notlatest@" + digest3, sha256digestHex}, // name@digest@ID, with name containing a digest
	} {
		storageRef, err := Transport.ParseStoreReference(store, c.input)
		if c.expectedRef == "" && c.expectedID == "" {
			assert.Error(t, err, c.input)
		} else {
			require.NoError(t, err, c.input)
			assert.Equal(t, store, storageRef.transport.store, c.input)
			assert.Equal(t, c.expectedRef, storageRef.reference, c.input)
			assert.Equal(t, c.expectedID, storageRef.id, c.input)
			if c.expectedRef == "" {
				assert.Nil(t, storageRef.name, c.input)
			} else {
				dockerRef, err := reference.ParseNormalizedNamed(c.expectedRef)
				require.NoError(t, err)
				require.NotNil(t, storageRef.name, c.input)
				assert.Equal(t, dockerRef.String(), storageRef.reference)
				assert.Equal(t, dockerRef.String(), storageRef.DockerReference().String())
			}
		}
	}
}

func TestTransportParseReference(t *testing.T) {
	store := newStore(t)
	driver := store.GraphDriverName()
	root := store.GraphRoot()

	for _, c := range []struct{ prefix, expectedDriver, expectedRoot, expectedRunRoot string }{
		{"", driver, root, ""},                                             // Implicit store location prefix
		{"[unterminated", "", "", ""},                                      // Unterminated store specifier
		{"[]", "", "", ""},                                                 // Empty store specifier
		{"[relative/path]", "", "", ""},                                    // Non-absolute graph root path
		{"[" + driver + "@relative/path]", "", "", ""},                     // Non-absolute graph root path
		{"[@" + root + "suffix2]", "", "", ""},                             // Empty graph driver
		{"[" + driver + "@]", "", "", ""},                                  // Empty root path
		{"[thisisunknown@" + root + "suffix2]", "", "", ""},                // Unknown graph driver
		{"[" + root + "suffix1]", "", "", ""},                              // A valid root path, but no run dir
		{"[" + driver + "@" + root + "suffix3+relative/path]", "", "", ""}, // Non-absolute run dir
		{"[" + driver + "@" + root + "suffix3+" + root + "suffix4]",
			driver,
			root + "suffix3",
			root + "suffix4"}, // A valid root@graph+run set
		{"[" + driver + "@" + root + "suffix3+" + root + "suffix4:options,options,options]",
			driver,
			root + "suffix3",
			root + "suffix4"}, // A valid root@graph+run+options set
	} {
		t.Logf("parsing %q", c.prefix+"busybox")
		ref, err := Transport.ParseReference(c.prefix + "busybox")
		if c.expectedDriver == "" {
			assert.Error(t, err, c.prefix)
		} else {
			require.NoError(t, err, c.prefix)
			storageRef, ok := ref.(*storageReference)
			require.True(t, ok, c.prefix)
			assert.Equal(t, c.expectedDriver, storageRef.transport.store.GraphDriverName(), c.prefix)
			assert.Equal(t, c.expectedRoot, storageRef.transport.store.GraphRoot(), c.prefix)
			if c.expectedRunRoot != "" {
				assert.Equal(t, c.expectedRunRoot, storageRef.transport.store.RunRoot(), c.prefix)
			}
		}
	}
}

func TestTransportValidatePolicyConfigurationScope(t *testing.T) {
	store := newStore(t)
	driver := store.GraphDriverName()
	root := store.GraphRoot()
	storeSpec := fmt.Sprintf("[%s@%s]", driver, root) // As computed in PolicyConfigurationNamespaces

	// Valid inputs
	for _, scope := range []string{
		"[" + root + "suffix1]",                                                                    // driverlessStoreSpec in PolicyConfigurationNamespaces
		"[" + driver + "@" + root + "suffix3]",                                                     // storeSpec in PolicyConfigurationNamespaces
		storeSpec + "@" + sha256digestHex,                                                          // ID only
		storeSpec + "docker.io",                                                                    // Host name only
		storeSpec + "docker.io/library",                                                            // A repository namespace
		storeSpec + "docker.io/library/busybox",                                                    // A repository name
		storeSpec + "docker.io/library/busybox:notlatest",                                          // name:tag
		storeSpec + "docker.io/library/busybox:notlatest@" + sha256digestHex,                       // name@ID
		storeSpec + "docker.io/library/busybox@" + sha256Digest2,                                   // name@digest
		storeSpec + "docker.io/library/busybox@" + sha256Digest2 + "@" + sha256digestHex,           // name@digest@ID
		storeSpec + "docker.io/library/busybox:notlatest@" + sha256Digest2,                         // name:tag@digest
		storeSpec + "docker.io/library/busybox:notlatest@" + sha256Digest2 + "@" + sha256digestHex, // name:tag@digest@ID
	} {
		err := Transport.ValidatePolicyConfigurationScope(scope)
		assert.NoError(t, err, scope)
	}

	// Invalid inputs
	for _, scope := range []string{
		"busybox",                        // Unprefixed reference
		"[unterminated",                  // Unterminated store specifier
		"[]",                             // Empty store specifier
		"[relative/path]",                // Non-absolute graph root path
		"[" + driver + "@relative/path]", // Non-absolute graph root path
		// "[thisisunknown@" + root + "suffix2]", // Unknown graph driver FIXME: validate against storage.ListGraphDrivers() once that's available
		storeSpec + "@", // An incomplete two-component name

		storeSpec + "docker.io/library/busybox@sha256:ab",                    // Invalid digest in name@digest
		storeSpec + "docker.io/library/busybox@ab",                           // Invalid ID in name@ID
		storeSpec + "docker.io/library/busybox@",                             // Empty ID/digest in name@ID
		storeSpec + "docker.io/library/busybox@@" + sha256digestHex,          // Empty digest in name@digest@ID
		storeSpec + "docker.io/library/busybox@ab@" + sha256digestHex,        // Invalid digest in name@digest@ID
		storeSpec + "docker.io/library/busybox@sha256:ab@" + sha256digestHex, // Invalid digest in name@digest@ID
		storeSpec + "docker.io/library/busybox@" + sha256Digest2 + "@",       // Empty ID in name@digest@ID
		storeSpec + "docker.io/library/busybox@" + sha256Digest2 + "@ab",     // Invalid ID in name@digest@ID
	} {
		err := Transport.ValidatePolicyConfigurationScope(scope)
		assert.Error(t, err, scope)
	}
}
<commit_msg>Add a test for verboseName<commit_after>// +build !containers_image_storage_stub

package storage

import (
	"fmt"
	"testing"

	"github.com/containers/image/docker/reference"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	sha256digestHex = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
	sha256Digest2   = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
)

func TestTransportName(t *testing.T) {
	assert.Equal(t, "containers-storage", Transport.Name())
}

func TestTransportParseStoreReference(t *testing.T) {
	const digest3 = "sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"

	store := newStore(t)

	Transport.SetStore(nil)
	for _, c := range []struct{ input, expectedRef, expectedID string }{
		{"", "", ""}, // Empty input
		// Handling of the store prefix
		// FIXME? Should we be silently discarding input like this?
		{"[unterminated", "", ""},                                    // Unterminated store specifier
		{"[garbage]busybox", "docker.io/library/busybox:latest", ""}, // Store specifier is overridden by the store we pass to ParseStoreReference

		{"UPPERCASEISINVALID", "", ""},                                                   // Invalid single-component name
		{"sha256:" + sha256digestHex, "docker.io/library/sha256:" + sha256digestHex, ""}, // Valid single-component name; the hex part is not an ID unless it has a "@" prefix, so it looks like a tag
		// FIXME: This test is now incorrect, this should not fail _if the image ID matches_
		{sha256digestHex, "", ""},                    // Invalid single-component ID; not an ID without a "@" prefix, so it's parsed as a name, but names aren't allowed to look like IDs
		{"@" + sha256digestHex, "", sha256digestHex}, // Valid single-component ID
		// FIXME: This should fail it seems, everything else uses ref.digest only if ref.name != nil
		// {"@sha256:" + sha256digestHex, "", ""},                                      // Valid single-component digest
		// "aaaa", either a valid image ID prefix, or a short form of docker.io/library/aaaa, untested
		{"sha256:ab", "docker.io/library/sha256:ab", ""},                                   // Valid single-component name, explicit tag
		{"busybox", "docker.io/library/busybox:latest", ""},                                // Valid single-component name, implicit tag
		{"busybox:notlatest", "docker.io/library/busybox:notlatest", ""},                   // Valid single-component name, explicit tag
		{"docker.io/library/busybox:notlatest", "docker.io/library/busybox:notlatest", ""}, // Valid single-component name, everything explicit

		{"UPPERCASEISINVALID@" + sha256digestHex, "", ""},                                                // Invalid name in name@digestOrID
		{"busybox@ab", "", ""},                                                                           // Invalid ID in name@digestOrID
		{"busybox@", "", ""},                                                                             // Empty ID in name@digestOrID
		{"busybox@sha256:ab", "", ""},                                                                    // Invalid digest in name@digestOrID
		{"busybox@sha256:" + sha256digestHex, "docker.io/library/busybox@sha256:" + sha256digestHex, ""}, // Valid name@digest, no tag
		{"busybox@" + sha256digestHex, "docker.io/library/busybox:latest", sha256digestHex},              // Valid name@ID, implicit tag
		// "busybox@aaaa", a valid image ID prefix, untested
		{"busybox:notlatest@" + sha256digestHex, "docker.io/library/busybox:notlatest", sha256digestHex},                     // Valid name@ID, explicit tag
		{"docker.io/library/busybox:notlatest@" + sha256digestHex, "docker.io/library/busybox:notlatest", sha256digestHex},   // Valid name@ID, everything explicit
		{"docker.io/library/busybox:notlatest@" + sha256Digest2, "docker.io/library/busybox:notlatest@" + sha256Digest2, ""}, // Valid name:tag@digest, everything explicit

		{"busybox@sha256:" + sha256digestHex + "@ab", "", ""},                                                                                // Invalid ID in name@digest@ID
		{"busybox@ab@" + sha256digestHex, "", ""},                                                                                // Invalid digest in name@digest@ID
		{"busybox@@" + sha256digestHex, "", ""},                                                                                // Invalid digest in name@digest@ID
		{"busybox@" + sha256Digest2 + "@" + sha256digestHex, "docker.io/library/busybox@" + sha256Digest2, sha256digestHex},                                                         // name@digest@ID
		{"docker.io/library/busybox@" + sha256Digest2 + "@" + sha256digestHex, "docker.io/library/busybox@" + sha256Digest2, sha256digestHex},                                       // name@digest@ID, everything explicit
		{"docker.io/library/busybox:notlatest@sha256:" + sha256digestHex + "@" + sha256digestHex, "docker.io/library/busybox:notlatest@sha256:" + sha256digestHex, sha256digestHex}, // name:tag@digest@ID, everything explicit
		// FIXME: Is this supposed to work? the validation of idOrDigest seems to make this impossible.
		// "busybox@sha256:"+sha256digestHex+"@aaaa", a valid image ID prefix, untested
		// FIXME FIXME: two digests
		{"busybox:notlatest@" + sha256Digest2 + "@" + digest3 + "@" + sha256digestHex, "docker.io/library/busybox:notlatest@" + digest3, sha256digestHex}, // name@digest@ID, with name containing a digest
	} {
		storageRef, err := Transport.ParseStoreReference(store, c.input)
		if c.expectedRef == "" && c.expectedID == "" {
			assert.Error(t, err, c.input)
		} else {
			require.NoError(t, err, c.input)
			assert.Equal(t, store, storageRef.transport.store, c.input)
			assert.Equal(t, c.expectedRef, storageRef.reference, c.input)
			assert.Equal(t, c.expectedID, storageRef.id, c.input)
			if c.expectedRef == "" {
				assert.Nil(t, storageRef.name, c.input)
			} else {
				dockerRef, err := reference.ParseNormalizedNamed(c.expectedRef)
				require.NoError(t, err)
				require.NotNil(t, storageRef.name, c.input)
				assert.Equal(t, dockerRef.String(), storageRef.reference)
				assert.Equal(t, dockerRef.String(), storageRef.DockerReference().String())
			}
		}
	}
}

func TestTransportParseReference(t *testing.T) {
	store := newStore(t)
	driver := store.GraphDriverName()
	root := store.GraphRoot()

	for _, c := range []struct{ prefix, expectedDriver, expectedRoot, expectedRunRoot string }{
		{"", driver, root, ""},                                             // Implicit store location prefix
		{"[unterminated", "", "", ""},                                      // Unterminated store specifier
		{"[]", "", "", ""},                                                 // Empty store specifier
		{"[relative/path]", "", "", ""},                                    // Non-absolute graph root path
		{"[" + driver + "@relative/path]", "", "", ""},                     // Non-absolute graph root path
		{"[@" + root + "suffix2]", "", "", ""},                             // Empty graph driver
		{"[" + driver + "@]", "", "", ""},                                  // Empty root path
		{"[thisisunknown@" + root + "suffix2]", "", "", ""},                // Unknown graph driver
		{"[" + root + "suffix1]", "", "", ""},                              // A valid root path, but no run dir
		{"[" + driver + "@" + root + "suffix3+relative/path]", "", "", ""}, // Non-absolute run dir
		{"[" + driver + "@" + root + "suffix3+" + root + "suffix4]",
			driver,
			root + "suffix3",
			root + "suffix4"}, // A valid root@graph+run set
		{"[" + driver + "@" + root + "suffix3+" + root + "suffix4:options,options,options]",
			driver,
			root + "suffix3",
			root + "suffix4"}, // A valid root@graph+run+options set
	} {
		t.Logf("parsing %q", c.prefix+"busybox")
		ref, err := Transport.ParseReference(c.prefix + "busybox")
		if c.expectedDriver == "" {
			assert.Error(t, err, c.prefix)
		} else {
			require.NoError(t, err, c.prefix)
			storageRef, ok := ref.(*storageReference)
			require.True(t, ok, c.prefix)
			assert.Equal(t, c.expectedDriver, storageRef.transport.store.GraphDriverName(), c.prefix)
			assert.Equal(t, c.expectedRoot, storageRef.transport.store.GraphRoot(), c.prefix)
			if c.expectedRunRoot != "" {
				assert.Equal(t, c.expectedRunRoot, storageRef.transport.store.RunRoot(), c.prefix)
			}
		}
	}
}

func TestTransportValidatePolicyConfigurationScope(t *testing.T) {
	store := newStore(t)
	driver := store.GraphDriverName()
	root := store.GraphRoot()
	storeSpec := fmt.Sprintf("[%s@%s]", driver, root) // As computed in PolicyConfigurationNamespaces

	// Valid inputs
	for _, scope := range []string{
		"[" + root + "suffix1]",                                                                    // driverlessStoreSpec in PolicyConfigurationNamespaces
		"[" + driver + "@" + root + "suffix3]",                                                     // storeSpec in PolicyConfigurationNamespaces
		storeSpec + "@" + sha256digestHex,                                                          // ID only
		storeSpec + "docker.io",                                                                    // Host name only
		storeSpec + "docker.io/library",                                                            // A repository namespace
		storeSpec + "docker.io/library/busybox",                                                    // A repository name
		storeSpec + "docker.io/library/busybox:notlatest",                                          // name:tag
		storeSpec + "docker.io/library/busybox:notlatest@" + sha256digestHex,                       // name@ID
		storeSpec + "docker.io/library/busybox@" + sha256Digest2,                                   // name@digest
		storeSpec + "docker.io/library/busybox@" + sha256Digest2 + "@" + sha256digestHex,           // name@digest@ID
		storeSpec + "docker.io/library/busybox:notlatest@" + sha256Digest2,                         // name:tag@digest
		storeSpec + "docker.io/library/busybox:notlatest@" + sha256Digest2 + "@" + sha256digestHex, // name:tag@digest@ID
	} {
		err := Transport.ValidatePolicyConfigurationScope(scope)
		assert.NoError(t, err, scope)
	}

	// Invalid inputs
	for _, scope := range []string{
		"busybox",                        // Unprefixed reference
		"[unterminated",                  // Unterminated store specifier
		"[]",                             // Empty store specifier
		"[relative/path]",                // Non-absolute graph root path
		"[" + driver + "@relative/path]", // Non-absolute graph root path
		// "[thisisunknown@" + root + "suffix2]", // Unknown graph driver FIXME: validate against storage.ListGraphDrivers() once that's available
		storeSpec + "@", // An incomplete two-component name

		storeSpec + "docker.io/library/busybox@sha256:ab",                    // Invalid digest in name@digest
		storeSpec + "docker.io/library/busybox@ab",                           // Invalid ID in name@ID
		storeSpec + "docker.io/library/busybox@",                             // Empty ID/digest in name@ID
		storeSpec + "docker.io/library/busybox@@" + sha256digestHex,          // Empty digest in name@digest@ID
		storeSpec + "docker.io/library/busybox@ab@" + sha256digestHex,        // Invalid digest in name@digest@ID
		storeSpec + "docker.io/library/busybox@sha256:ab@" + sha256digestHex, // Invalid digest in name@digest@ID
		storeSpec + "docker.io/library/busybox@" + sha256Digest2 + "@",       // Empty ID in name@digest@ID
		storeSpec + "docker.io/library/busybox@" + sha256Digest2 + "@ab",     // Invalid ID in name@digest@ID
	} {
		err := Transport.ValidatePolicyConfigurationScope(scope)
		assert.Error(t, err, scope)
	}
}

func TestVerboseName(t *testing.T) {
	newStore(t)

	assert.Equal(t, "", verboseName(nil))

	for _, c := range []struct{ input, expected string }{
		{"sha256:" + sha256digestHex, "@sha256:" + sha256digestHex},
		{"busybox", "docker.io/library/busybox"}, // The normalization actually happens already in ParseAnyReference
		{"docker.io/library/busybox", "docker.io/library/busybox"},
		{"docker.io/library/busybox:latest", "docker.io/library/busybox:latest"},
		{"docker.io/library/busybox:notlatest", "docker.io/library/busybox:notlatest"},
		{"docker.io/library/busybox@sha256:" + sha256digestHex, "docker.io/library/busybox@sha256:" + sha256digestHex},
		{"docker.io/library/busybox:notlatest@sha256:" + sha256digestHex, "docker.io/library/busybox:notlatest@sha256:" + sha256digestHex},
	} {
		ref, err := reference.ParseAnyReference(c.input)
		require.NoError(t, err, c.input)
		name := verboseName(ref)
		assert.Equal(t, c.expected, name, c.input)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/4180122/distbayes/distmlMatlab"
	"github.com/arcaneiceman/GoVector/govec"
	//"github.com/gonum/matrix/mat64"
	//"io/ioutil"
	"net"
	"os"
	//"strconv"
	//"strings"
	"encoding/gob"
	"time"
)

const BUFFSIZE = 200000

//10485760

var (
	cnum      int = 0
	name      string
	inputargs []string
	myaddr    *net.TCPAddr
	svaddr    *net.TCPAddr
	model     distmlMatlab.MatModel
	logger    *govec.GoLog
	X         string
	Y         string
	Xt        string
	Yt        string
	l         *net.TCPListener
	gmodel    distmlMatlab.MatGlobalModel
	gempty    distmlMatlab.MatGlobalModel
	committed bool
	isjoining bool
)

type message struct {
	Id       int
	NodeIp   string
	NodeName string
	Type     string
	Model    distmlMatlab.MatModel
	GModel   distmlMatlab.MatGlobalModel
}

type response struct {
	Resp  string
	Error string
}

func main() {
	//Parsing inputargs
	parseArgs()

	//Hacky solution to the Matlab problem (Mathworks, please fix this!)
	// see: https://www.mathworks.com/matlabcentral/answers/305877-what-is-the-primary-message-table-for-module-77
	// and  https://github.com/JuliaInterop/MATLAB.jl/issues/47
	distmlMatlab.Hack()

	//Initialize stuff
	model = distmlMatlab.NewModel(X, Y)

	//Initialize TCP Connection and listener
	l, _ = net.ListenTCP("tcp", myaddr)
	fmt.Printf("Node initialized as %v.\n", name)
	go listener()

	isjoining = true
	for isjoining {
		requestJoin()
	}

	committed = false

	//Main function of this server
	for {
		//parseUserInput()
		time.Sleep(time.Duration(2 * time.Second))
		if !committed {
			requestCommit()
		}
	}
}

func listener() {
	for {
		conn, err := l.AcceptTCP()
		checkError(err)
		go connHandler(conn)
	}
}

func connHandler(conn *net.TCPConn) {
	var msg message
	enc := gob.NewEncoder(conn)
	dec := gob.NewDecoder(conn)
	err := dec.Decode(&msg)
	checkError(err)
	switch msg.Type {
	case "test_request":
		// server is asking me to test
		enc.Encode(response{"OK", ""})
		go testModel(msg.Id, msg.Model)
	case "global_grant":
		// server is sending global model
		enc.Encode(response{"OK", ""})
		gmodel = msg.GModel
		fmt.Printf("\n <-- Pulled global model from server.\nEnter command: ")
		//go testGlobal(msg.GModel)
	default:
		// respond to ping
		enc.Encode(response{"NO", "Unknown Command"})
	}
	conn.Close()
}

func parseUserInput() {
	var ident string
	reader := bufio.NewReader(os.Stdin)
	fmt.Print("Enter command: ")
	text, _ := reader.ReadString('\n')
	//Windows adds its own strange carriage return, the following lines fix it
	if text[len(text)-2] == '\r' {
		ident = text[0 : len(text)-2]
	} else {
		ident = text[0 : len(text)-1]
	}
	switch ident {
	case "read":
		//x = readData(inputargs[3])
		//y = readData(inputargs[4])
		fmt.Printf(" --- Local data updated.\n")
	case "train":
		model = distmlMatlab.NewModel(X, Y) //GOOD
		fmt.Printf(" --- Local model error on local data is: %v.\n", model.Weight)
	case "push":
		requestCommit()
	case "pull":
		requestGlobal()
	case "valid":
		acc, _ := distmlMatlab.GetErrorGlobal(X, Y, gmodel)
		fmt.Printf(" --- Global model error on local data is: %v.\n", acc)
	case "test":
		acc := distmlMatlab.GetError(Xt, Yt, model)
		fmt.Printf(" --- Local model error on test data is: %v.\n", acc)
	case "testg":
		acc, _ := distmlMatlab.GetErrorGlobal(Xt, Yt, gmodel)
		fmt.Printf(" --- Global model error on test data is: %v.\n", acc)
	case "who":
		fmt.Printf("%v\n", name)
	default:
		fmt.Printf(" Command not recognized: %v.\n\n", ident)
		fmt.Printf("  Choose from the following commands\n")
		fmt.Printf("  read  -- Read data from disk\n")
		fmt.Printf("  push  -- Push trained model to server\n")
		fmt.Printf("  pull  -- Obtain global model from server\n")
		fmt.Printf("  train -- Train model from data (reports error)\n")
		fmt.Printf("  valid -- Validate global model with local data\n")
		fmt.Printf("  test  -- Test local model with test data\n")
		fmt.Printf("  testg -- Test global model with test data\n")
		fmt.Printf("  who   -- Print node name\n\n")
	}
}

func requestJoin() {
	msg := message{cnum, myaddr.String(), name, "join_request", model, gempty}
	fmt.Printf(" --> Asking server to join.")
	tcpSend(msg)
}

func requestCommit() {
	cnum++
	msg := message{cnum, myaddr.String(), name, "commit_request", model, gempty}
	fmt.Printf(" --> Pushing local model to server.")
	tcpSend(msg)
}

func requestGlobal() {
	msg := message{cnum, myaddr.String(), name, "global_request", model, gempty}
	fmt.Printf(" --> Requesting global model from server.")
	tcpSend(msg)
}

func testModel(id int, testmodel distmlMatlab.MatModel) {
	//func testModel(id int, testmodel bclass.Model) {
	fmt.Printf("\n <-- Received test requset.\nEnter command: ")
	distmlMatlab.TestModel(X, Y, &testmodel)
	msg := message{id, myaddr.String(), name, "test_complete", testmodel, gempty}
	fmt.Printf("\n --> Sending completed test requset.")
	tcpSend(msg)
	fmt.Printf("Enter command: ")
}

func tcpSend(msg message) {
	conn, err := net.DialTCP("tcp", nil, svaddr)
	checkError(err)
	enc := gob.NewEncoder(conn)
	dec := gob.NewDecoder(conn)
	err = enc.Encode(&msg)
	checkError(err)
	var r response
	err = dec.Decode(&r)
	checkError(err)
	if r.Resp == "OK" {
		fmt.Printf(" [OK]\n")
		if msg.Type == "commit_request" {
			committed = true
		}
		if msg.Type == "join_request" {
			isjoining = false
		}
	} else if r.Resp == "NO" {
		fmt.Printf(" [%s]\n *** Request was denied by server: %v.\nEnter command: ", r.Resp, r.Error)
	} else {
		fmt.Printf(" [%s]\n *** Something strange Happened: %v.\nEnter command: ", r.Resp, r.Error)
	}
}

func parseArgs() {
	flag.Parse()
	inputargs = flag.Args()
	var err error
	if len(inputargs) < 2 {
		fmt.Printf("Not enough inputs.\n")
		return
	}
	name = inputargs[0]
	myaddr, err = net.ResolveTCPAddr("tcp", inputargs[1])
	checkError(err)
	svaddr, err = net.ResolveTCPAddr("tcp", inputargs[2])
	checkError(err)
	X = inputargs[3]
	Y = inputargs[4]
	Xt = "C:/work/src/github.com/4180122/distbayes/testdata/xv.txt"
	Yt = "C:/work/src/github.com/4180122/distbayes/testdata/yv.txt"
	logger = govec.Initialize(inputargs[0], inputargs[5])
}

func checkError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error())
		os.Exit(1)
	}
}
<commit_msg>Updated client state machine<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/4180122/distbayes/distmlMatlab"
	"github.com/arcaneiceman/GoVector/govec"
	//"github.com/gonum/matrix/mat64"
	//"io/ioutil"
	"net"
	"os"
	//"strconv"
	//"strings"
	"encoding/gob"
	"time"
)

const BUFFSIZE = 200000

//10485760

var (
	cnum      int = 0
	name      string
	inputargs []string
	myaddr    *net.TCPAddr
	svaddr    *net.TCPAddr
	model     distmlMatlab.MatModel
	logger    *govec.GoLog
	X         string
	Y         string
	Xt        string
	Yt        string
	l         *net.TCPListener
	gmodel    distmlMatlab.MatGlobalModel
	gempty    distmlMatlab.MatGlobalModel
	committed bool
	isjoining bool
	istesting bool
)

type message struct {
	Id       int
	NodeIp   string
	NodeName string
	Type     string
	Model    distmlMatlab.MatModel
	GModel   distmlMatlab.MatGlobalModel
}

type response struct {
	Resp  string
	Error string
}

func main() {
	//Parsing inputargs
	parseArgs()

	//Hacky solution to the Matlab problem (Mathworks, please fix this!)
	// see: https://www.mathworks.com/matlabcentral/answers/305877-what-is-the-primary-message-table-for-module-77
	// and  https://github.com/JuliaInterop/MATLAB.jl/issues/47
	distmlMatlab.Hack()

	//Initialize stuff
	model = distmlMatlab.NewModel(X, Y)

	//Initialize TCP Connection and listener
	l, _ = net.ListenTCP("tcp", myaddr)
	fmt.Printf("Node initialized as %v.\n", name)
	go listener()

	isjoining = true
	for isjoining {
		requestJoin()
	}

	committed = false
	istesting = false

	//Main function of this server
	for {
		//parseUserInput()
		time.Sleep(time.Duration(2 * time.Second))
		if !committed && !istesting {
			requestCommit()
		}
	}
}

func listener() {
	for {
		conn, err := l.AcceptTCP()
		checkError(err)
		go connHandler(conn)
	}
}

func connHandler(conn *net.TCPConn) {
	var msg message
	enc := gob.NewEncoder(conn)
	dec := gob.NewDecoder(conn)
	err := dec.Decode(&msg)
	checkError(err)
	switch msg.Type {
	case "test_request":
		// server is asking me to test
		enc.Encode(response{"OK", ""})
		go testModel(msg.Id, msg.Model)
	case "global_grant":
		// server is sending global model
		enc.Encode(response{"OK", ""})
		gmodel = msg.GModel
		fmt.Printf("\n <-- Pulled global model from server.\nEnter command: ")
		//go testGlobal(msg.GModel)
	default:
		// respond to ping
		enc.Encode(response{"NO", "Unknown Command"})
	}
	conn.Close()
}

func parseUserInput() {
	var ident string
	reader := bufio.NewReader(os.Stdin)
	fmt.Print("Enter command: ")
	text, _ := reader.ReadString('\n')
	//Windows adds its own strange carriage return, the following lines fix it
	if text[len(text)-2] == '\r' {
		ident = text[0 : len(text)-2]
	} else {
		ident = text[0 : len(text)-1]
	}
	switch ident {
	case "read":
		//x = readData(inputargs[3])
		//y = readData(inputargs[4])
		fmt.Printf(" --- Local data updated.\n")
	case "train":
		model = distmlMatlab.NewModel(X, Y) //GOOD
		fmt.Printf(" --- Local model error on local data is: %v.\n", model.Weight)
	case "push":
		requestCommit()
	case "pull":
		requestGlobal()
	case "valid":
		acc, _ := distmlMatlab.GetErrorGlobal(X, Y, gmodel)
		fmt.Printf(" --- Global model error on local data is: %v.\n", acc)
	case "test":
		acc := distmlMatlab.GetError(Xt, Yt, model)
		fmt.Printf(" --- Local model error on test data is: %v.\n", acc)
	case "testg":
		acc, _ := distmlMatlab.GetErrorGlobal(Xt, Yt, gmodel)
		fmt.Printf(" --- Global model error on test data is: %v.\n", acc)
	case "who":
		fmt.Printf("%v\n", name)
	default:
		fmt.Printf(" Command not recognized: %v.\n\n", ident)
		fmt.Printf("  Choose from the following commands\n")
		fmt.Printf("  read  -- Read data from disk\n")
		fmt.Printf("  push  -- Push trained model to server\n")
		fmt.Printf("  pull  -- Obtain global model from server\n")
		fmt.Printf("  train -- Train model from data (reports error)\n")
		fmt.Printf("  valid -- Validate global model with local data\n")
		fmt.Printf("  test  -- Test local model with test data\n")
		fmt.Printf("  testg -- Test global model with test data\n")
		fmt.Printf("  who   -- Print node name\n\n")
	}
}

func requestJoin() {
	msg := message{cnum, myaddr.String(), name, "join_request", model, gempty}
	fmt.Printf(" --> Asking server to join.")
	tcpSend(msg)
}

func requestCommit() {
	cnum++
	msg := message{cnum, myaddr.String(), name, "commit_request", model, gempty}
	fmt.Printf(" --> Pushing local model to server.")
	tcpSend(msg)
}

func requestGlobal() {
	msg := message{cnum, myaddr.String(), name, "global_request", model, gempty}
	fmt.Printf(" --> Requesting global model from server.")
	tcpSend(msg)
}

func testModel(id int, testmodel distmlMatlab.MatModel) {
	//func testModel(id int, testmodel bclass.Model) {
	fmt.Printf("\n <-- Received test requset.\nEnter command: ")
	istesting = true
	distmlMatlab.TestModel(X, Y, &testmodel)
	msg := message{id, myaddr.String(), name, "test_complete", testmodel, gempty}
	fmt.Printf("\n --> Sending completed test requset.")
	tcpSend(msg)
	istesting = false
	fmt.Printf("Enter command: ")
}

func tcpSend(msg message) {
	conn, err := net.DialTCP("tcp", nil, svaddr)
	checkError(err)
	enc := gob.NewEncoder(conn)
	dec := gob.NewDecoder(conn)
	err = enc.Encode(&msg)
	checkError(err)
	var r response
	err = dec.Decode(&r)
	checkError(err)
	if r.Resp == "OK" {
		fmt.Printf(" [OK]\n")
		if msg.Type == "commit_request" {
			committed = true
		}
		if msg.Type == "join_request" {
			isjoining = false
		}
	} else if r.Resp == "NO" {
		fmt.Printf(" [%s]\n *** Request was denied by server: %v.\nEnter command: ", r.Resp, r.Error)
	} else {
		fmt.Printf(" [%s]\n *** Something strange Happened: %v.\nEnter command: ", r.Resp, r.Error)
	}
}

func parseArgs() {
	flag.Parse()
	inputargs = flag.Args()
	var err error
	if len(inputargs) < 2 {
		fmt.Printf("Not enough inputs.\n")
		return
	}
	name = inputargs[0]
	myaddr, err = net.ResolveTCPAddr("tcp", inputargs[1])
	checkError(err)
	svaddr, err = net.ResolveTCPAddr("tcp", inputargs[2])
	checkError(err)
	X = inputargs[3]
	Y = inputargs[4]
	Xt = "C:/work/src/github.com/4180122/distbayes/testdata/xv.txt"
	Yt = "C:/work/src/github.com/4180122/distbayes/testdata/yv.txt"
	logger = govec.Initialize(inputargs[0], inputargs[5])
}

func checkError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error())
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package cloud

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"testing"
)

func TestStatusCategoryService_GetList(t *testing.T) {
	setup()
	defer teardown()
	testAPIEndpoint := "/rest/api/3/statuscategory"

	raw, err := os.ReadFile("../testing/mock-data/all_statuscategories.json")
	if err != nil {
		t.Error(err.Error())
	}
	testMux.HandleFunc(testAPIEndpoint, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testRequestURL(t, r, testAPIEndpoint)
		fmt.Fprint(w, string(raw))
	})

	statusCategory, _, err := testClient.StatusCategory.GetList(context.Background())
	if statusCategory == nil {
		t.Error("Expected statusCategory list. StatusCategory list is nil")
	}
	if err != nil {
		t.Errorf("Error given: %s", err)
	}
}

func TestStatusCategoryService_Get(t *testing.T) {
	setup()
	defer teardown()
	testAPIEndpoint := "/rest/api/3/statuscategory/1"

	raw, err := os.ReadFile("../testing/mock-data/status_category.json")
	if err != nil {
		t.Error(err.Error())
	}
	testMux.HandleFunc(testAPIEndpoint, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testRequestURL(t, r, testAPIEndpoint)
		fmt.Fprint(w, string(raw))
	})

	statusCategory, _, err := testClient.StatusCategory.Get(context.Background(), "1")

	if err != nil {
		t.Errorf("Error given: %s", err)
	} else if statusCategory == nil {
		t.Error("Expected status category. StatusCategory is nil")
	}
}
<commit_msg>Cloud/Status category: Added two additional testing checks<commit_after>package cloud

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"testing"
)

func TestStatusCategoryService_GetList(t *testing.T) {
	setup()
	defer teardown()
	testAPIEndpoint := "/rest/api/3/statuscategory"

	raw, err := os.ReadFile("../testing/mock-data/all_statuscategories.json")
	if err != nil {
		t.Error(err.Error())
	}

	testMux.HandleFunc(testAPIEndpoint, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testRequestURL(t, r, testAPIEndpoint)
		fmt.Fprint(w, string(raw))
	})

	statusCategory, _, err := testClient.StatusCategory.GetList(context.Background())
	if statusCategory == nil {
		t.Error("Expected statusCategory list. StatusCategory list is nil")
	}
	if l := len(statusCategory); l != 4 {
		t.Errorf("Expected 4 statusCategory list items. Got %d", l)
	}
	if err != nil {
		t.Errorf("Error given: %s", err)
	}
}

func TestStatusCategoryService_Get(t *testing.T) {
	setup()
	defer teardown()
	testAPIEndpoint := "/rest/api/3/statuscategory/1"

	raw, err := os.ReadFile("../testing/mock-data/status_category.json")
	if err != nil {
		t.Error(err.Error())
	}

	testMux.HandleFunc(testAPIEndpoint, func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, http.MethodGet)
		testRequestURL(t, r, testAPIEndpoint)
		fmt.Fprint(w, string(raw))
	})

	statusCategory, _, err := testClient.StatusCategory.Get(context.Background(), "1")
	if err != nil {
		t.Errorf("Error given: %s", err)
	} else if statusCategory == nil {
		t.Error("Expected status category. StatusCategory is nil")

		// Checking testdata
	} else if statusCategory.ColorName != "medium-gray" {
		t.Errorf("Expected statusCategory.ColorName to be 'medium-gray'. Got '%s'", statusCategory.ColorName)
	}
}
<|endoftext|>


<commit_before>package tests

import (
	"errors"
	"koding/db/mongodb/modelhelper"
	"socialapi/config"
	"socialapi/models"
	"socialapi/request"
	"socialapi/rest"
	"socialapi/workers/integration/webhook"
	"socialapi/workers/integration/webhook/api"
	"testing"
	"time"

	"github.com/koding/runner"
	. "github.com/smartystreets/goconvey/convey"
)

func newPushRequest(body string) *api.PushRequest {
	wr := &api.PushRequest{
		Message: webhook.Message{
			Body: body,
		},
	}

	return wr
}

func newBotChannelRequest(nick, groupName string) *api.BotChannelRequest {
	return &api.BotChannelRequest{
		GroupName: groupName,
		Username:  nick,
	}
}

func TestWebhook(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	appConfig := config.MustRead(r.Conf.Path)
	modelhelper.Initialize(appConfig.Mongo)
	defer modelhelper.Close()

	Convey("We should be able to successfully push message", t, func() {
		channelIntegration, topicChannel := webhook.CreateTestChannelIntegration(t)

		account, err := models.CreateAccountInBothDbsWithNick("sinan")
		So(err, ShouldBeNil)

		_, err = topicChannel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		err = rest.DoPushRequest(newPushRequest(models.RandomName()), channelIntegration.Token)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, channelIntegration.GroupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		resp, err := rest.GetHistory(channelIntegration.ChannelId,
			&request.Query{
				AccountId: account.Id,
			},
			ses.ClientId,
		)

		So(err, ShouldBeNil)
		So(len(resp.MessageList), ShouldEqual, 1)
	})

	Convey("We should be able to successfully fetch bot channel of the user", t, func() {
		account, err := models.CreateAccountInBothDbsWithNick("sinan")
		So(err, ShouldBeNil)
		groupName := models.RandomGroupName()
		channel := models.CreateTypedGroupedChannelWithTest(account.Id, models.Channel_TYPE_GROUP, groupName)
		_, err = channel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, groupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		channelId, err := rest.DoBotChannelRequest(ses.ClientId)

		So(err, ShouldBeNil)
		So(channelId, ShouldNotEqual, 0)
	})

	Convey("We should be able to successfully receive github push messages via middleware", t, func() {
		channelIntegration, topicChannel := webhook.CreateTestChannelIntegration(t)

		account, err := models.CreateAccountInBothDbsWithNick(models.RandomName())
		So(err, ShouldBeNil)

		channel := models.CreateTypedGroupedChannelWithTest(account.Id, models.Channel_TYPE_GROUP, channelIntegration.GroupName)

		_, err = channel.AddParticipant(account.Id)
		So(err, ShouldBeNil)
		_, err = topicChannel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		err = rest.DoGithubPush(githubPushEventData, channelIntegration.Token)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, channelIntegration.GroupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		tick := time.Tick(time.Millisecond * 200)
		deadLine := time.After(10 * time.Second)
		for {
			select {
			case <-tick:
				resp, err := rest.GetHistory(topicChannel.Id,
					&request.Query{},
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				if len(resp.MessageList) > 0 {
					So(len(resp.MessageList), ShouldEqual, 1)
					So(resp.MessageList[0].Message.Body, ShouldStartWith, "[canthefason](https://github.com/canthefason) [pushed]")
					return
				}
			case <-deadLine:
				So(errors.New("Could not fetch messages"), ShouldBeNil)
			}
		}

	})

	Convey("We should be able to successfully receive pivotal push messages via middleware", t, func() {
		channelIntegration, topicChannel := webhook.CreateTestChannelIntegration(t)

		account, err := models.CreateAccountInBothDbsWithNick(models.RandomName())
		So(err, ShouldBeNil)

		channel := models.CreateTypedGroupedChannelWithTest(account.Id, models.Channel_TYPE_GROUP, channelIntegration.GroupName)

		_, err = channel.AddParticipant(account.Id)
		So(err, ShouldBeNil)
		_, err = topicChannel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		err = rest.DoPivotalPush("POST", pivotalEventData, channelIntegration.Token)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, channelIntegration.GroupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		tick := time.Tick(time.Millisecond * 200)
		deadLine := time.After(10 * time.Second)
		for {
			select {
			case <-tick:
				resp, err := rest.GetHistory(topicChannel.Id,
					&request.Query{},
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				if len(resp.MessageList) > 0 {
					So(len(resp.MessageList), ShouldEqual, 1)
					So(resp.MessageList[0].Message.Body, ShouldStartWith, "[pivotal-project] Mehmet Ali Savas started this feature")
					return
				}
			case <-deadLine:
				So(errors.New("Could not fetch messages"), ShouldBeNil)
			}
		}

	})
}
<commit_msg>social/tests: payload test of integration is added<commit_after>package tests

import (
	"errors"
	"koding/db/mongodb/modelhelper"
	"socialapi/config"
	"socialapi/models"
	"socialapi/request"
	"socialapi/rest"
	"socialapi/workers/integration/webhook"
	"socialapi/workers/integration/webhook/api"
	"testing"
	"time"

	"github.com/koding/runner"
	. "github.com/smartystreets/goconvey/convey"
)

func newPushRequest(body string) *api.PushRequest {
	wr := &api.PushRequest{
		Message: webhook.Message{
			Body: body,
		},
	}

	return wr
}

func newBotChannelRequest(nick, groupName string) *api.BotChannelRequest {
	return &api.BotChannelRequest{
		GroupName: groupName,
		Username:  nick,
	}
}

func TestWebhook(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	appConfig := config.MustRead(r.Conf.Path)
	modelhelper.Initialize(appConfig.Mongo)
	defer modelhelper.Close()

	Convey("We should be able to successfully push message", t, func() {
		channelIntegration, topicChannel := webhook.CreateTestChannelIntegration(t)

		account, err := models.CreateAccountInBothDbsWithNick("sinan")
		So(err, ShouldBeNil)

		_, err = topicChannel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		err = rest.DoPushRequest(newPushRequest(models.RandomName()), channelIntegration.Token)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, channelIntegration.GroupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		resp, err := rest.GetHistory(channelIntegration.ChannelId,
			&request.Query{
				AccountId: account.Id,
			},
			ses.ClientId,
		)

		So(err, ShouldBeNil)
		So(len(resp.MessageList), ShouldEqual, 1)
	})

	Convey("We should be able to successfully fetch bot channel of the user", t, func() {
		account, err := models.CreateAccountInBothDbsWithNick("sinan")
		So(err, ShouldBeNil)
		groupName := models.RandomGroupName()
		channel := models.CreateTypedGroupedChannelWithTest(account.Id, models.Channel_TYPE_GROUP, groupName)
		_, err = channel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, groupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		channelId, err := rest.DoBotChannelRequest(ses.ClientId)

		So(err, ShouldBeNil)
		So(channelId, ShouldNotEqual, 0)
	})

	Convey("We should be able to successfully receive github push messages via middleware", t, func() {
		channelIntegration, topicChannel := webhook.CreateTestChannelIntegration(t)

		account, err := models.CreateAccountInBothDbsWithNick(models.RandomName())
		So(err, ShouldBeNil)

		channel := models.CreateTypedGroupedChannelWithTest(account.Id, models.Channel_TYPE_GROUP, channelIntegration.GroupName)

		_, err = channel.AddParticipant(account.Id)
		So(err, ShouldBeNil)
		_, err = topicChannel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		err = rest.DoGithubPush(githubPushEventData, channelIntegration.Token)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, channelIntegration.GroupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		tick := time.Tick(time.Millisecond * 200)
		deadLine := time.After(10 * time.Second)
		for {
			select {
			case <-tick:
				resp, err := rest.GetHistory(topicChannel.Id,
					&request.Query{},
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				if len(resp.MessageList) > 0 {
					So(len(resp.MessageList), ShouldEqual, 1)
					So(resp.MessageList[0].Message.Body, ShouldStartWith, "[canthefason](https://github.com/canthefason) [pushed]")
					return
				}
			case <-deadLine:
				So(errors.New("Could not fetch messages"), ShouldBeNil)
			}
		}

	})

	Convey("We should be able to successfully receive payload of github push messages via middleware", t, func() {
		channelIntegration, topicChannel := webhook.CreateTestChannelIntegration(t)

		account, err := models.CreateAccountInBothDbsWithNick(models.RandomName())
		So(err, ShouldBeNil)

		channel := models.CreateTypedGroupedChannelWithTest(account.Id, models.Channel_TYPE_GROUP, channelIntegration.GroupName)

		_, err = channel.AddParticipant(account.Id)
		So(err, ShouldBeNil)
		_, err = topicChannel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		err = rest.DoGithubPush(githubPushEventData, channelIntegration.Token)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, channelIntegration.GroupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		tick := time.Tick(time.Millisecond * 200)
		deadLine := time.After(10 * time.Second)
		for {
			select {
			case <-tick:
				resp, err := rest.GetHistory(topicChannel.Id,
					&request.Query{},
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				if len(resp.MessageList) > 0 {
					So(len(resp.MessageList), ShouldEqual, 1)
					So(*resp.MessageList[0].Message.Payload["eventType"], ShouldEqual, "push")
					return
				}
			case <-deadLine:
				So(errors.New("Could not fetch messages"), ShouldBeNil)
			}
		}

	})

	Convey("We should be able to successfully receive pivotal push messages via middleware", t, func() {
		channelIntegration, topicChannel := webhook.CreateTestChannelIntegration(t)

		account, err := models.CreateAccountInBothDbsWithNick(models.RandomName())
		So(err, ShouldBeNil)

		channel := models.CreateTypedGroupedChannelWithTest(account.Id, models.Channel_TYPE_GROUP, channelIntegration.GroupName)

		_, err = channel.AddParticipant(account.Id)
		So(err, ShouldBeNil)
		_, err = topicChannel.AddParticipant(account.Id)
		So(err, ShouldBeNil)

		err = rest.DoPivotalPush("POST", pivotalEventData, channelIntegration.Token)
		So(err, ShouldBeNil)

		ses, err := models.FetchOrCreateSession(account.Nick, channelIntegration.GroupName)
		So(err, ShouldBeNil)
		So(ses, ShouldNotBeNil)

		tick := time.Tick(time.Millisecond * 200)
		deadLine := time.After(10 * time.Second)
		for {
			select {
			case <-tick:
				resp, err := rest.GetHistory(topicChannel.Id,
					&request.Query{},
					ses.ClientId,
				)
				So(err, ShouldBeNil)
				if len(resp.MessageList) > 0 {
					So(len(resp.MessageList), ShouldEqual, 1)
					So(resp.MessageList[0].Message.Body, ShouldStartWith, "[pivotal-project] Mehmet Ali Savas started this feature")
					return
				}
			case <-deadLine:
				So(errors.New("Could not fetch messages"), ShouldBeNil)
			}
		}

	})
}
<|endoftext|>


<commit_before>package gorm

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"reflect"
)

func getRealValue(value reflect.Value, field string) interface{} {
	result := reflect.Indirect(value).FieldByName(field).Interface()
	if r, ok := result.(driver.Valuer); ok {
		result, _ = r.Value()
	}
	return result
}

func equalAsString(a interface{}, b interface{}) bool {
	return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b)
}

func Preload(scope *Scope) {
	fields := scope.Fields()
	isSlice := scope.IndirectValue().Kind() == reflect.Slice

	if scope.Search.preload != nil {
		for key, conditions := range scope.Search.preload {
			for _, field := range fields {
				if field.Name == key && field.Relationship != nil {
					results := makeSlice(field.Struct.Type)
					relation := field.Relationship
					primaryName := scope.PrimaryField().Name
					associationPrimaryKey := scope.New(results).PrimaryField().Name

					switch relation.Kind {
					case "has_one":
						if primaryKeys := scope.getColumnAsArray(primaryName); len(primaryKeys) > 0 {
							condition := fmt.Sprintf("%v IN (?)", scope.Quote(relation.ForeignDBName))
							scope.NewDB().Where(condition, primaryKeys).Find(results, conditions...)

							resultValues := reflect.Indirect(reflect.ValueOf(results))
							for i := 0; i < resultValues.Len(); i++ {
								result := resultValues.Index(i)
								if isSlice {
									value := getRealValue(result, relation.ForeignFieldName)
									objects := scope.IndirectValue()
									for j := 0; j < objects.Len(); j++ {
										if equalAsString(getRealValue(objects.Index(j), primaryName), value) {
											reflect.Indirect(objects.Index(j)).FieldByName(field.Name).Set(result)
											break
										}
									}
								} else {
									scope.SetColumn(field, result)
								}
							}
						}
					case "has_many":
						if primaryKeys := scope.getColumnAsArray(primaryName); len(primaryKeys) > 0 {
							condition := fmt.Sprintf("%v IN (?)", scope.Quote(relation.ForeignDBName))
							scope.NewDB().Where(condition, primaryKeys).Find(results, conditions...)
							resultValues := reflect.Indirect(reflect.ValueOf(results))
							if isSlice {
								for i := 0; i < resultValues.Len(); i++ {
									result := resultValues.Index(i)
									value := getRealValue(result, relation.ForeignFieldName)
									objects := scope.IndirectValue()
									for j := 0; j < objects.Len(); j++ {
										object := reflect.Indirect(objects.Index(j))
										if equalAsString(getRealValue(object, primaryName), value) {
											f := object.FieldByName(field.Name)
											f.Set(reflect.Append(f, result))
											break
										}
									}
								}
							} else {
								scope.SetColumn(field, resultValues)
							}
						}
					case "belongs_to":
						if primaryKeys := scope.getColumnAsArray(relation.ForeignFieldName); len(primaryKeys) > 0 {
							scope.NewDB().Where(primaryKeys).Find(results, conditions...)
							resultValues := reflect.Indirect(reflect.ValueOf(results))
							for i := 0; i < resultValues.Len(); i++ {
								result := resultValues.Index(i)
								if isSlice {
									value := getRealValue(result, associationPrimaryKey)
									objects := scope.IndirectValue()
									for j := 0; j < objects.Len(); j++ {
										object := reflect.Indirect(objects.Index(j))
										if equalAsString(getRealValue(object, relation.ForeignFieldName), value) {
											object.FieldByName(field.Name).Set(result)
										}
									}
								} else {
									scope.SetColumn(field, result)
								}
							}
						}
					case "many_to_many":
						scope.Err(errors.New("not supported relation"))
					default:
						scope.Err(errors.New("not supported relation"))
					}
					break
				}
			}
		}
	}
}

func makeSlice(typ reflect.Type) interface{} {
	if typ.Kind() == reflect.Slice {
		typ = typ.Elem()
	}
	sliceType := reflect.SliceOf(typ)
	slice := reflect.New(sliceType)
	slice.Elem().Set(reflect.MakeSlice(sliceType, 0, 0))
	return slice.Interface()
}

func (scope *Scope) getColumnAsArray(column string) (primaryKeys []interface{}) {
	values := scope.IndirectValue()
	switch values.Kind() {
	case reflect.Slice:
		primaryKeyMap := map[interface{}]bool{}
		for i := 0; i < values.Len(); i++ {
			primaryKeyMap[reflect.Indirect(values.Index(i)).FieldByName(column).Interface()] = true
		}
		for key := range primaryKeyMap {
			primaryKeys = append(primaryKeys, key)
		}
	case reflect.Struct:
		return []interface{}{values.FieldByName(column).Interface()}
	}
	return
}
<commit_msg>Only load Fields when defined preload<commit_after>package gorm

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"reflect"
)

func getRealValue(value reflect.Value, field string) interface{} {
	result := reflect.Indirect(value).FieldByName(field).Interface()
	if r, ok := result.(driver.Valuer); ok {
		result, _ = r.Value()
	}
	return result
}

func equalAsString(a interface{}, b interface{}) bool {
	return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b)
}

func Preload(scope *Scope) {
	if scope.Search.preload != nil {
		fields := scope.Fields()
		isSlice := scope.IndirectValue().Kind() == reflect.Slice

		for key, conditions := range scope.Search.preload {
			for _, field := range fields {
				if field.Name == key && field.Relationship != nil {
					results := makeSlice(field.Struct.Type)
					relation := field.Relationship
					primaryName := scope.PrimaryField().Name
					associationPrimaryKey := scope.New(results).PrimaryField().Name

					switch relation.Kind {
					case "has_one":
						if primaryKeys := scope.getColumnAsArray(primaryName); len(primaryKeys) > 0 {
							condition := fmt.Sprintf("%v IN (?)", scope.Quote(relation.ForeignDBName))
							scope.NewDB().Where(condition, primaryKeys).Find(results, conditions...)

							resultValues := reflect.Indirect(reflect.ValueOf(results))
							for i := 0; i < resultValues.Len(); i++ {
								result := resultValues.Index(i)
								if isSlice {
									value := getRealValue(result, relation.ForeignFieldName)
									objects := scope.IndirectValue()
									for j := 0; j < objects.Len(); j++ {
										if equalAsString(getRealValue(objects.Index(j), primaryName), value) {
											reflect.Indirect(objects.Index(j)).FieldByName(field.Name).Set(result)
											break
										}
									}
								} else {
									scope.SetColumn(field, result)
								}
							}
						}
					case "has_many":
						if primaryKeys := scope.getColumnAsArray(primaryName); len(primaryKeys) > 0 {
							condition := fmt.Sprintf("%v IN (?)", scope.Quote(relation.ForeignDBName))
							scope.NewDB().Where(condition, primaryKeys).Find(results, conditions...)
							resultValues := reflect.Indirect(reflect.ValueOf(results))
							if isSlice {
								for i := 0; i < resultValues.Len(); i++ {
									result := resultValues.Index(i)
									value := getRealValue(result, relation.ForeignFieldName)
									objects := scope.IndirectValue()
									for j := 0; j < objects.Len(); j++ {
										object := reflect.Indirect(objects.Index(j))
										if equalAsString(getRealValue(object, primaryName), value) {
											f := object.FieldByName(field.Name)
											f.Set(reflect.Append(f, result))
											break
										}
									}
								}
							} else {
								scope.SetColumn(field, resultValues)
							}
						}
					case "belongs_to":
						if primaryKeys := scope.getColumnAsArray(relation.ForeignFieldName); len(primaryKeys) > 0 {
							scope.NewDB().Where(primaryKeys).Find(results, conditions...)
							resultValues := reflect.Indirect(reflect.ValueOf(results))
							for i := 0; i < resultValues.Len(); i++ {
								result := resultValues.Index(i)
								if isSlice {
									value := getRealValue(result, associationPrimaryKey)
									objects := scope.IndirectValue()
									for j := 0; j < objects.Len(); j++ {
										object := reflect.Indirect(objects.Index(j))
										if equalAsString(getRealValue(object, relation.ForeignFieldName), value) {
											object.FieldByName(field.Name).Set(result)
										}
									}
								} else {
									scope.SetColumn(field, result)
								}
							}
						}
					case "many_to_many":
						scope.Err(errors.New("not supported relation"))
					default:
						scope.Err(errors.New("not supported relation"))
					}
					break
				}
			}
		}
	}
}

func makeSlice(typ reflect.Type) interface{} {
	if typ.Kind() == reflect.Slice {
		typ = typ.Elem()
	}
	sliceType := reflect.SliceOf(typ)
	slice := reflect.New(sliceType)
	slice.Elem().Set(reflect.MakeSlice(sliceType, 0, 0))
	return slice.Interface()
}

func (scope *Scope) getColumnAsArray(column string) (primaryKeys []interface{}) {
	values := scope.IndirectValue()
	switch values.Kind() {
	case reflect.Slice:
		primaryKeyMap := map[interface{}]bool{}
		for i := 0; i < values.Len(); i++ {
			primaryKeyMap[reflect.Indirect(values.Index(i)).FieldByName(column).Interface()] = true
		}
		for key := range primaryKeyMap {
			primaryKeys = append(primaryKeys, key)
		}
	case reflect.Struct:
		return []interface{}{values.FieldByName(column).Interface()}
	}
	return
}
<|endoftext|>