<commit_before>// Package lint provides abstractions on top of go/analysis.
package lint

import (
	"errors"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/token"
	"strconv"
	"strings"

	"golang.org/x/tools/go/analysis"
)

type Documentation struct {
	Title      string
	Text       string
	Since      string
	NonDefault bool
	Options    []string
}

func (doc *Documentation) String() string {
	b := &strings.Builder{}
	fmt.Fprintf(b, "%s\n\n", doc.Title)
	if doc.Text != "" {
		fmt.Fprintf(b, "%s\n\n", doc.Text)
	}
	fmt.Fprint(b, "Available since\n    ")
	if doc.Since == "" {
		fmt.Fprint(b, "unreleased")
	} else {
		fmt.Fprintf(b, "%s", doc.Since)
	}
	if doc.NonDefault {
		fmt.Fprint(b, ", non-default")
	}
	fmt.Fprint(b, "\n")
	if len(doc.Options) > 0 {
		fmt.Fprintf(b, "\nOptions\n")
		for _, opt := range doc.Options {
			fmt.Fprintf(b, "    %s", opt)
		}
		fmt.Fprint(b, "\n")
	}
	return b.String()
}

func newVersionFlag() flag.Getter {
	tags := build.Default.ReleaseTags
	v := tags[len(tags)-1][2:]
	version := new(VersionFlag)
	if err := version.Set(v); err != nil {
		panic(fmt.Sprintf("internal error: %s", err))
	}
	return version
}

type VersionFlag int

func (v *VersionFlag) String() string {
	return fmt.Sprintf("1.%d", *v)
}

func (v *VersionFlag) Set(s string) error {
	if len(s) < 3 {
		return errors.New("invalid Go version")
	}
	if s[0] != '1' {
		return errors.New("invalid Go version")
	}
	if s[1] != '.' {
		return errors.New("invalid Go version")
	}
	i, err := strconv.Atoi(s[2:])
	*v = VersionFlag(i)
	return err
}

func (v *VersionFlag) Get() interface{} {
	return int(*v)
}

func InitializeAnalyzers(docs map[string]*Documentation, analyzers map[string]*analysis.Analyzer) map[string]*analysis.Analyzer {
	out := make(map[string]*analysis.Analyzer, len(analyzers))
	for k, v := range analyzers {
		vc := *v
		out[k] = &vc

		vc.Name = k
		doc, ok := docs[k]
		if !ok {
			panic(fmt.Sprintf("missing documentation for check %s", k))
		}
		vc.Doc = doc.String()
		if vc.Flags.Usage == nil {
			fs := flag.NewFlagSet("", flag.PanicOnError)
			fs.Var(newVersionFlag(), "go", "Target Go version")
			vc.Flags = *fs
		}
	}
	return out
}

// ExhaustiveTypeSwitch panics when called. It can be used to ensure
// that type switches are exhaustive.
func ExhaustiveTypeSwitch(v interface{}) {
	panic(fmt.Sprintf("internal error: unhandled case %T", v))
}

// A directive is a comment of the form '//lint:<command>
// [arguments...]'. It represents instructions to the static analysis
// tool.
type Directive struct {
	Command   string
	Arguments []string
	Directive *ast.Comment
	Node      ast.Node
}

func parseDirective(s string) (cmd string, args []string) {
	if !strings.HasPrefix(s, "//lint:") {
		return "", nil
	}
	s = strings.TrimPrefix(s, "//lint:")
	fields := strings.Split(s, " ")
	return fields[0], fields[1:]
}

func ParseDirectives(files []*ast.File, fset *token.FileSet) []Directive {
	var dirs []Directive
	for _, f := range files {
		// OPT(dh): in our old code, we skip all the commentmap work if we
		// couldn't find any directives, benchmark if that's actually
		// worth doing
		cm := ast.NewCommentMap(fset, f, f.Comments)
		for node, cgs := range cm {
			for _, cg := range cgs {
				for _, c := range cg.List {
					if !strings.HasPrefix(c.Text, "//lint:") {
						continue
					}
					cmd, args := parseDirective(c.Text)
					d := Directive{
						Command:   cmd,
						Arguments: args,
						Directive: c,
						Node:      node,
					}
					dirs = append(dirs, d)
				}
			}
		}
	}
	return dirs
}
<commit_msg>analysis/lint: include link to online documentation in output of -explain<commit_after>// Package lint provides abstractions on top of go/analysis.
package lint

import (
	"errors"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/token"
	"strconv"
	"strings"

	"golang.org/x/tools/go/analysis"
)

type Documentation struct {
	Title      string
	Text       string
	Since      string
	NonDefault bool
	Options    []string
}

func (doc *Documentation) String() string {
	b := &strings.Builder{}
	fmt.Fprintf(b, "%s\n\n", doc.Title)
	if doc.Text != "" {
		fmt.Fprintf(b, "%s\n\n", doc.Text)
	}
	fmt.Fprint(b, "Available since\n    ")
	if doc.Since == "" {
		fmt.Fprint(b, "unreleased")
	} else {
		fmt.Fprintf(b, "%s", doc.Since)
	}
	if doc.NonDefault {
		fmt.Fprint(b, ", non-default")
	}
	fmt.Fprint(b, "\n")
	if len(doc.Options) > 0 {
		fmt.Fprintf(b, "\nOptions\n")
		for _, opt := range doc.Options {
			fmt.Fprintf(b, "    %s", opt)
		}
		fmt.Fprint(b, "\n")
	}
	return b.String()
}

func newVersionFlag() flag.Getter {
	tags := build.Default.ReleaseTags
	v := tags[len(tags)-1][2:]
	version := new(VersionFlag)
	if err := version.Set(v); err != nil {
		panic(fmt.Sprintf("internal error: %s", err))
	}
	return version
}

type VersionFlag int

func (v *VersionFlag) String() string {
	return fmt.Sprintf("1.%d", *v)
}

func (v *VersionFlag) Set(s string) error {
	if len(s) < 3 {
		return errors.New("invalid Go version")
	}
	if s[0] != '1' {
		return errors.New("invalid Go version")
	}
	if s[1] != '.' {
		return errors.New("invalid Go version")
	}
	i, err := strconv.Atoi(s[2:])
	*v = VersionFlag(i)
	return err
}

func (v *VersionFlag) Get() interface{} {
	return int(*v)
}

func InitializeAnalyzers(docs map[string]*Documentation, analyzers map[string]*analysis.Analyzer) map[string]*analysis.Analyzer {
	out := make(map[string]*analysis.Analyzer, len(analyzers))
	for k, v := range analyzers {
		vc := *v
		out[k] = &vc

		vc.Name = k
		doc, ok := docs[k]
		if !ok {
			panic(fmt.Sprintf("missing documentation for check %s", k))
		}
		vc.Doc = fmt.Sprintf("%s\nOnline documentation\n    https://staticcheck.io/docs/checks#%s", doc.String(), k)
		if vc.Flags.Usage == nil {
			fs := flag.NewFlagSet("", flag.PanicOnError)
			fs.Var(newVersionFlag(), "go", "Target Go version")
			vc.Flags = *fs
		}
	}
	return out
}

// ExhaustiveTypeSwitch panics when called. It can be used to ensure
// that type switches are exhaustive.
func ExhaustiveTypeSwitch(v interface{}) {
	panic(fmt.Sprintf("internal error: unhandled case %T", v))
}

// A directive is a comment of the form '//lint:<command>
// [arguments...]'. It represents instructions to the static analysis
// tool.
type Directive struct {
	Command   string
	Arguments []string
	Directive *ast.Comment
	Node      ast.Node
}

func parseDirective(s string) (cmd string, args []string) {
	if !strings.HasPrefix(s, "//lint:") {
		return "", nil
	}
	s = strings.TrimPrefix(s, "//lint:")
	fields := strings.Split(s, " ")
	return fields[0], fields[1:]
}

func ParseDirectives(files []*ast.File, fset *token.FileSet) []Directive {
	var dirs []Directive
	for _, f := range files {
		// OPT(dh): in our old code, we skip all the commentmap work if we
		// couldn't find any directives, benchmark if that's actually
		// worth doing
		cm := ast.NewCommentMap(fset, f, f.Comments)
		for node, cgs := range cm {
			for _, cg := range cgs {
				for _, c := range cg.List {
					if !strings.HasPrefix(c.Text, "//lint:") {
						continue
					}
					cmd, args := parseDirective(c.Text)
					d := Directive{
						Command:   cmd,
						Arguments: args,
						Directive: c,
						Node:      node,
					}
					dirs = append(dirs, d)
				}
			}
		}
	}
	return dirs
}
<|endoftext|>


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

package main

import (
  "flag"
  "fmt"
  "log"
  "net/http"
)

var (
  argPort = flag.Int("port", 8080, "The port to listen to for incomming HTTP requests")
)

func main() {
  flag.Parse()
  log.Print("Starting HTTP server on port ", *argPort)

  // Run a HTTP server that serves static files from current directory.
  // TODO(bryk): Disable directory listing.
  http.Handle("/", http.FileServer(http.Dir("./")))
  err := http.ListenAndServe(fmt.Sprintf(":%d", *argPort), nil)

  if err != nil {
    log.Fatal("HTTP server error: ", err)
  }
}
<commit_msg>Use pflag and glog packages in go server code<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
  "flag"
  "fmt"
  "net/http"

  "github.com/golang/glog"
  "github.com/spf13/pflag"
)

var (
  argPort = pflag.Int("port", 8080, "The port to listen to for incomming HTTP requests")
)

func main() {
  pflag.CommandLine.AddGoFlagSet(flag.CommandLine)
  pflag.Parse()
  glog.Info("Starting HTTP server on port ", *argPort)
  defer glog.Flush();

  // Run a HTTP server that serves static files from current directory.
  // TODO(bryk): Disable directory listing.
  http.Handle("/", http.FileServer(http.Dir("./")))
  glog.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", *argPort), nil))
}
<|endoftext|>


<commit_before>package core

import (
	"reflect"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/elves/elvish/edit/tty"
	"github.com/elves/elvish/edit/ui"
	"github.com/elves/elvish/styled"
	"github.com/elves/elvish/sys"
)

var (
	ka     = ui.Key{Rune: 'a'}
	kb     = ui.Key{Rune: 'b'}
	kc     = ui.Key{Rune: 'c'}
	kEnter = ui.Key{Rune: ui.Enter}

	keysABCEnter = []ui.Key{ka, kb, kc, kEnter}
	eventsABC    = []tty.Event{
		tty.KeyEvent(ka), tty.KeyEvent(kb), tty.KeyEvent(kc)}
	eventsABCEnter = []tty.Event{
		tty.KeyEvent(ka), tty.KeyEvent(kb),
		tty.KeyEvent(kc), tty.KeyEvent(kEnter)}
)

func TestReadCode_PassesInputEventsToMode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	m := &fakeMode{maxKeys: 3}
	ed.state.Mode = m

	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}

	ed.ReadCode()

	wantKeysHandled := []ui.Key{
		ui.Key{Rune: 'a'}, ui.Key{Rune: 'b'}, ui.Key{Rune: 'c'},
	}
	if !reflect.DeepEqual(m.keysHandled, wantKeysHandled) {
		t.Errorf("Mode gets keys %v, want %v", m.keysHandled, wantKeysHandled)
	}
}

func TestReadCode_CallsBeforeReadlineOnce(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)

	called := 0
	ed.config.BeforeReadline = []func(){func() { called++ }}

	// Causes basicMode to quit
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed.ReadCode()

	if called != 1 {
		t.Errorf("BeforeReadline hook called %d times, want 1", called)
	}
}

func TestReadCode_CallsAfterReadlineOnceWithCode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)

	called := 0
	code := ""
	ed.config.AfterReadline = []func(string){func(s string) {
		called++
		code = s
	}}

	// Causes basicMode to write state.Code and then quit
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed.ReadCode()

	if called != 1 {
		t.Errorf("AfterReadline hook called %d times, want 1", called)
	}
	if code != "abc" {
		t.Errorf("AfterReadline hook called with %q, want %q", code, "abc")
	}
}

func TestReadCode_RespectsMaxHeight(t *testing.T) {
	maxHeight := 5

	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	// Will fill more than maxHeight but less than terminal height
	ed.state.Code = strings.Repeat("a", 80*10)
	ed.state.Dot = len(ed.state.Code)

	codeCh, _ := readCodeAsync(ed)

	buf1 := <-terminal.bufCh
	// Make sure that normally the height does exceed maxHeight.
	if h := len(buf1.Lines); h <= maxHeight {
		t.Errorf("Buffer height is %d, should > %d", h, maxHeight)
	}

	ed.configMutex.Lock()
	ed.config.RenderConfig.MaxHeight = maxHeight
	ed.configMutex.Unlock()

	ed.loop.Redraw(false)
	buf2 := <-terminal.bufCh
	if h := len(buf2.Lines); h > maxHeight {
		t.Errorf("Buffer height is %d, should <= %d", h, maxHeight)
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

var bufChTimeout = 1 * time.Second

func TestReadCode_RendersHighlightedCode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	ed.config.RenderConfig.Highlighter = func(code string) (styled.Text, []error) {
		return styled.Text{
			styled.Segment{styled.Style{Foreground: "red"}, code}}, nil
	}

	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}
	codeCh, _ := readCodeAsync(ed)

	wantBuf := ui.NewBufferBuilder(80).
		WriteString("abc", "31" /* SGR for red foreground */).
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("Did not see buffer containing highlighted code")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func TestReadCode_RendersErrorFromHighlighter(t *testing.T) {
	// TODO
}

func TestReadCode_RendersPrompt(t *testing.T) {
	// TODO
}

func TestReadCode_RendersRprompt(t *testing.T) {
	// TODO
}

func TestReadCode_SupportsPersistentRprompt(t *testing.T) {
	// TODO
}

func TestReadCode_QuitsOnSIGHUP(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	codeCh, _ := readCodeAsync(ed)
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("a").
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGHUP")
	}

	sigs.ch <- syscall.SIGHUP

	select {
	case <-codeCh:
		// TODO: Test that ReadCode returns with io.EOF
	case <-time.After(time.Second):
		t.Errorf("SIGHUP did not cause ReadCode to return")
	}
}

func TestReadCode_ResetsOnSIGHUP(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	codeCh, _ := readCodeAsync(ed)
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("a").
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGINT")
	}

	sigs.ch <- syscall.SIGINT

	wantBuf = ui.NewBufferBuilder(80).Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("Terminal state is not reset after SIGINT")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func TestReadCode_RedrawsOnSIGWINCH(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	ed.state.Code = "1234567890"
	ed.state.Dot = len(ed.state.Code)

	codeCh, _ := readCodeAsync(ed)
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("1234567890").
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGWINCH")
	}

	terminal.setSize(24, 4)
	sigs.ch <- sys.SIGWINCH

	wantBuf = ui.NewBufferBuilder(4).WriteUnstyled("1234567890").
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("Terminal is not redrawn after SIGWINCH")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func readCodeAsync(ed *Editor) (<-chan string, <-chan error) {
	codeCh := make(chan string, 1)
	errCh := make(chan error, 1)
	go func() {
		code, err := ed.ReadCode()
		codeCh <- code
		errCh <- err
	}()
	return codeCh, errCh
}

var checkBufferTimeout = time.Second

func checkBuffer(ch <-chan *ui.Buffer, want *ui.Buffer) bool {
	for {
		select {
		case buf := <-ch:
			if reflect.DeepEqual(buf, want) {
				return true
			}
		case <-time.After(checkBufferTimeout):
			return false
		}
	}
}
<commit_msg>newedit/core: Remove unused code in editor_test.go.<commit_after>package core

import (
	"reflect"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/elves/elvish/edit/tty"
	"github.com/elves/elvish/edit/ui"
	"github.com/elves/elvish/styled"
	"github.com/elves/elvish/sys"
)

func TestReadCode_PassesInputEventsToMode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	m := &fakeMode{maxKeys: 3}
	ed.state.Mode = m

	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}

	ed.ReadCode()

	wantKeysHandled := []ui.Key{
		ui.Key{Rune: 'a'}, ui.Key{Rune: 'b'}, ui.Key{Rune: 'c'},
	}
	if !reflect.DeepEqual(m.keysHandled, wantKeysHandled) {
		t.Errorf("Mode gets keys %v, want %v", m.keysHandled, wantKeysHandled)
	}
}

func TestReadCode_CallsBeforeReadlineOnce(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)

	called := 0
	ed.config.BeforeReadline = []func(){func() { called++ }}

	// Causes basicMode to quit
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed.ReadCode()

	if called != 1 {
		t.Errorf("BeforeReadline hook called %d times, want 1", called)
	}
}

func TestReadCode_CallsAfterReadlineOnceWithCode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)

	called := 0
	code := ""
	ed.config.AfterReadline = []func(string){func(s string) {
		called++
		code = s
	}}

	// Causes basicMode to write state.Code and then quit
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed.ReadCode()

	if called != 1 {
		t.Errorf("AfterReadline hook called %d times, want 1", called)
	}
	if code != "abc" {
		t.Errorf("AfterReadline hook called with %q, want %q", code, "abc")
	}
}

func TestReadCode_RespectsMaxHeight(t *testing.T) {
	maxHeight := 5

	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	// Will fill more than maxHeight but less than terminal height
	ed.state.Code = strings.Repeat("a", 80*10)
	ed.state.Dot = len(ed.state.Code)

	codeCh, _ := readCodeAsync(ed)

	buf1 := <-terminal.bufCh
	// Make sure that normally the height does exceed maxHeight.
	if h := len(buf1.Lines); h <= maxHeight {
		t.Errorf("Buffer height is %d, should > %d", h, maxHeight)
	}

	ed.configMutex.Lock()
	ed.config.RenderConfig.MaxHeight = maxHeight
	ed.configMutex.Unlock()

	ed.loop.Redraw(false)
	buf2 := <-terminal.bufCh
	if h := len(buf2.Lines); h > maxHeight {
		t.Errorf("Buffer height is %d, should <= %d", h, maxHeight)
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

var bufChTimeout = 1 * time.Second

func TestReadCode_RendersHighlightedCode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	ed.config.RenderConfig.Highlighter = func(code string) (styled.Text, []error) {
		return styled.Text{
			styled.Segment{styled.Style{Foreground: "red"}, code}}, nil
	}

	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}
	codeCh, _ := readCodeAsync(ed)

	wantBuf := ui.NewBufferBuilder(80).
		WriteString("abc", "31" /* SGR for red foreground */).
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("Did not see buffer containing highlighted code")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func TestReadCode_RendersErrorFromHighlighter(t *testing.T) {
	// TODO
}

func TestReadCode_RendersPrompt(t *testing.T) {
	// TODO
}

func TestReadCode_RendersRprompt(t *testing.T) {
	// TODO
}

func TestReadCode_SupportsPersistentRprompt(t *testing.T) {
	// TODO
}

func TestReadCode_QuitsOnSIGHUP(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	codeCh, _ := readCodeAsync(ed)
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("a").
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGHUP")
	}

	sigs.ch <- syscall.SIGHUP

	select {
	case <-codeCh:
		// TODO: Test that ReadCode returns with io.EOF
	case <-time.After(time.Second):
		t.Errorf("SIGHUP did not cause ReadCode to return")
	}
}

func TestReadCode_ResetsOnSIGHUP(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	codeCh, _ := readCodeAsync(ed)
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("a").
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGINT")
	}

	sigs.ch <- syscall.SIGINT

	wantBuf = ui.NewBufferBuilder(80).Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("Terminal state is not reset after SIGINT")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func TestReadCode_RedrawsOnSIGWINCH(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	ed.state.Code = "1234567890"
	ed.state.Dot = len(ed.state.Code)

	codeCh, _ := readCodeAsync(ed)
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("1234567890").
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGWINCH")
	}

	terminal.setSize(24, 4)
	sigs.ch <- sys.SIGWINCH

	wantBuf = ui.NewBufferBuilder(4).WriteUnstyled("1234567890").
		SetDotToCursor().Buffer()
	if !checkBuffer(terminal.bufCh, wantBuf) {
		t.Errorf("Terminal is not redrawn after SIGWINCH")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func readCodeAsync(ed *Editor) (<-chan string, <-chan error) {
	codeCh := make(chan string, 1)
	errCh := make(chan error, 1)
	go func() {
		code, err := ed.ReadCode()
		codeCh <- code
		errCh <- err
	}()
	return codeCh, errCh
}

var checkBufferTimeout = time.Second

func checkBuffer(ch <-chan *ui.Buffer, want *ui.Buffer) bool {
	for {
		select {
		case buf := <-ch:
			if reflect.DeepEqual(buf, want) {
				return true
			}
		case <-time.After(checkBufferTimeout):
			return false
		}
	}
}
<|endoftext|>


<commit_before>package parsemail

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"mime/multipart"
	"net/mail"
	"strings"
	"time"
)

const contentTypeMultipartMixed = "multipart/mixed"
const contentTypeMultipartAlternative = "multipart/alternative"
const contentTypeMultipartRelated = "multipart/related"
const contentTypeTextHtml = "text/html"
const contentTypeTextPlain = "text/plain"

// Parse an email message read from io.Reader into parsemail.Email struct
func Parse(r io.Reader) (email Email, err error) {
	msg, err := mail.ReadMessage(r)
	if err != nil {
		return
	}

	email, err = createEmailFromHeader(msg.Header)
	if err != nil {
		return
	}

	contentType, params, err := parseContentType(msg.Header.Get("Content-Type"))
	if err != nil {
		return
	}

	switch contentType {
	case contentTypeMultipartMixed:
		email.TextBody, email.HTMLBody, email.Attachments, email.EmbeddedFiles, err = parseMultipartMixed(msg.Body, params["boundary"])
	case contentTypeMultipartAlternative:
		email.TextBody, email.HTMLBody, email.EmbeddedFiles, err = parseMultipartAlternative(msg.Body, params["boundary"])
	case contentTypeTextPlain:
		message, _ := ioutil.ReadAll(msg.Body)
		email.TextBody = strings.TrimSuffix(string(message[:]), "\n")
	case contentTypeTextHtml:
		message, _ := ioutil.ReadAll(msg.Body)
		email.HTMLBody = strings.TrimSuffix(string(message[:]), "\n")
	default:
		err = fmt.Errorf("Unknown top level mime type: %s", contentType)
	}

	return
}

func createEmailFromHeader(header mail.Header) (email Email, err error) {
	hp := headerParser{header: &header}

	email.Subject = decodeMimeSentence(header.Get("Subject"))
	email.From = hp.parseAddressList(header.Get("From"))
	email.Sender = hp.parseAddress(header.Get("Sender"))
	email.ReplyTo = hp.parseAddressList(header.Get("Reply-To"))
	email.To = hp.parseAddressList(header.Get("To"))
	email.Cc = hp.parseAddressList(header.Get("Cc"))
	email.Bcc = hp.parseAddressList(header.Get("Bcc"))
	email.Date = hp.parseTime(header.Get("Date"))
	email.ResentFrom = hp.parseAddressList(header.Get("Resent-From"))
	email.ResentSender = hp.parseAddress(header.Get("Resent-Sender"))
	email.ResentTo = hp.parseAddressList(header.Get("Resent-To"))
	email.ResentCc = hp.parseAddressList(header.Get("Resent-Cc"))
	email.ResentBcc = hp.parseAddressList(header.Get("Resent-Bcc"))
	email.ResentMessageID = hp.parseMessageId(header.Get("Resent-Message-ID"))
	email.MessageID = hp.parseMessageId(header.Get("Message-ID"))
	email.InReplyTo = hp.parseMessageIdList(header.Get("In-Reply-To"))
	email.References = hp.parseMessageIdList(header.Get("References"))
	email.ResentDate = hp.parseTime(header.Get("Resent-Date"))

	if hp.err != nil {
		err = hp.err
		return
	}

	//decode whole header for easier access to extra fields
	//todo: should we decode? aren't only standard fields mime encoded?
	email.Header, err = decodeHeaderMime(header)
	if err != nil {
		return
	}

	return
}

func parseContentType(contentTypeHeader string) (contentType string, params map[string]string, err error) {
	if contentTypeHeader == "" {
		contentType = contentTypeTextPlain
		return
	}

	return mime.ParseMediaType(contentTypeHeader)
}

func parseMultipartRelated(msg io.Reader, boundary string) (textBody, htmlBody string, embeddedFiles []EmbeddedFile, err error) {
	pmr := multipart.NewReader(msg, boundary)
	for {
		part, err := pmr.NextPart()

		if err == io.EOF {
			break
		} else if err != nil {
			return textBody, htmlBody, embeddedFiles, err
		}

		contentType, params, err := mime.ParseMediaType(part.Header.Get("Content-Type"))
		if err != nil {
			return textBody, htmlBody, embeddedFiles, err
		}

		switch contentType {
		case contentTypeTextPlain:
			ppContent, err := ioutil.ReadAll(part)
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			textBody += strings.TrimSuffix(string(ppContent[:]), "\n")
		case contentTypeTextHtml:
			ppContent, err := ioutil.ReadAll(part)
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			htmlBody += strings.TrimSuffix(string(ppContent[:]), "\n")
		case contentTypeMultipartAlternative:
			tb, hb, ef, err := parseMultipartAlternative(part, params["boundary"])
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			htmlBody += hb
			textBody += tb
			embeddedFiles = append(embeddedFiles, ef...)
		default:
			if isEmbeddedFile(part) {
				ef, err := decodeEmbeddedFile(part)
				if err != nil {
					return textBody, htmlBody, embeddedFiles, err
				}

				embeddedFiles = append(embeddedFiles, ef)
			} else {
				return textBody, htmlBody, embeddedFiles, fmt.Errorf("Can't process multipart/related inner mime type: %s", contentType)
			}
		}
	}

	return textBody, htmlBody, embeddedFiles, err
}

func parseMultipartAlternative(msg io.Reader, boundary string) (textBody, htmlBody string, embeddedFiles []EmbeddedFile, err error) {
	pmr := multipart.NewReader(msg, boundary)
	for {
		part, err := pmr.NextPart()

		if err == io.EOF {
			break
		} else if err != nil {
			return textBody, htmlBody, embeddedFiles, err
		}

		contentType, params, err := mime.ParseMediaType(part.Header.Get("Content-Type"))
		if err != nil {
			return textBody, htmlBody, embeddedFiles, err
		}

		switch contentType {
		case contentTypeTextPlain:
			ppContent, err := ioutil.ReadAll(part)
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			textBody += strings.TrimSuffix(string(ppContent[:]), "\n")
		case contentTypeTextHtml:
			ppContent, err := ioutil.ReadAll(part)
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			htmlBody += strings.TrimSuffix(string(ppContent[:]), "\n")
		case contentTypeMultipartRelated:
			tb, hb, ef, err := parseMultipartAlternative(part, params["boundary"])
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			htmlBody += hb
			textBody += tb
			embeddedFiles = append(embeddedFiles, ef...)
		default:
			if isEmbeddedFile(part) {
				ef, err := decodeEmbeddedFile(part)
				if err != nil {
					return textBody, htmlBody, embeddedFiles, err
				}

				embeddedFiles = append(embeddedFiles, ef)
			} else {
				return textBody, htmlBody, embeddedFiles, fmt.Errorf("Can't process multipart/alternative inner mime type: %s", contentType)
			}
		}
	}

	return textBody, htmlBody, embeddedFiles, err
}

func parseMultipartMixed(msg io.Reader, boundary string) (textBody, htmlBody string, attachments []Attachment, embeddedFiles []EmbeddedFile, err error) {
	mr := multipart.NewReader(msg, boundary)
	for {
		part, err := mr.NextPart()
		if err == io.EOF {
			break
		} else if err != nil {
			return textBody, htmlBody, attachments, embeddedFiles, err
		}

		contentType, params, err := mime.ParseMediaType(part.Header.Get("Content-Type"))
		if err != nil {
			return textBody, htmlBody, attachments, embeddedFiles, err
		}

		if contentType == contentTypeMultipartAlternative {
			textBody, htmlBody, embeddedFiles, err = parseMultipartAlternative(part, params["boundary"])
			if err != nil {
				return textBody, htmlBody, attachments, embeddedFiles, err
			}
		} else if contentType == contentTypeMultipartRelated {
			textBody, htmlBody, embeddedFiles, err = parseMultipartRelated(part, params["boundary"])
			if err != nil {
				return textBody, htmlBody, attachments, embeddedFiles, err
			}
		} else if isAttachment(part) {
			at, err := decodeAttachment(part)
			if err != nil {
				return textBody, htmlBody, attachments, embeddedFiles, err
			}

			attachments = append(attachments, at)
		} else {
			return textBody, htmlBody, attachments, embeddedFiles, fmt.Errorf("Unknown multipart/mixed nested mime type: %s", contentType)
		}
	}

	return textBody, htmlBody, attachments, embeddedFiles, err
}

func decodeMimeSentence(s string) string {
	result := []string{}
	ss := strings.Split(s, " ")

	for _, word := range ss {
		dec := new(mime.WordDecoder)
		w, err := dec.Decode(word)
		if err != nil {
			if len(result) == 0 {
				w = word
			} else {
				w = " " + word
			}
		}

		result = append(result, w)
	}

	return strings.Join(result, "")
}

func decodeHeaderMime(header mail.Header) (mail.Header, error) {
	parsedHeader := map[string][]string{}

	for headerName, headerData := range header {

		parsedHeaderData := []string{}
		for _, headerValue := range headerData {
			parsedHeaderData = append(parsedHeaderData, decodeMimeSentence(headerValue))
		}

		parsedHeader[headerName] = parsedHeaderData
	}

	return mail.Header(parsedHeader), nil
}

func decodePartData(part *multipart.Part) (io.Reader, error) {
	encoding := part.Header.Get("Content-Transfer-Encoding")

	if strings.EqualFold(encoding, "base64") {
		dr := base64.NewDecoder(base64.StdEncoding, part)
		dd, err := ioutil.ReadAll(dr)
		if err != nil {
			return nil, err
		}

		return bytes.NewReader(dd), nil
	}

	return nil, fmt.Errorf("Unknown encoding: %s", encoding)
}

func isEmbeddedFile(part *multipart.Part) bool {
	return part.Header.Get("Content-Transfer-Encoding") != ""
}

func decodeEmbeddedFile(part *multipart.Part) (ef EmbeddedFile, err error) {
	cid := decodeMimeSentence(part.Header.Get("Content-Id"))
	decoded, err := decodePartData(part)
	if err != nil {
		return
	}

	ef.CID = strings.Trim(cid, "<>")
	ef.Data = decoded
	ef.ContentType = part.Header.Get("Content-Type")

	return
}

func isAttachment(part *multipart.Part) bool {
	return part.FileName() != ""
}

func decodeAttachment(part *multipart.Part) (at Attachment, err error) {
	filename := decodeMimeSentence(part.FileName())
	decoded, err := decodePartData(part)
	if err != nil {
		return
	}

	at.Filename = filename
	at.Data = decoded
	at.ContentType = strings.Split(part.Header.Get("Content-Type"), ";")[0]

	return
}

type headerParser struct {
	header *mail.Header
	err    error
}

func (hp headerParser) parseAddress(s string) (ma *mail.Address) {
	if hp.err != nil {
		return nil
	}

	if strings.Trim(s, " \n") != "" {
		ma, hp.err = mail.ParseAddress(s)

		return ma
	}

	return nil
}

func (hp headerParser) parseAddressList(s string) (ma []*mail.Address) {
	if hp.err != nil {
		return
	}

	if strings.Trim(s, " \n") != "" {
		ma, hp.err = mail.ParseAddressList(s)
		return
	}

	return
}

func (hp headerParser) parseTime(s string) (t time.Time) {
	if hp.err != nil || s == "" {
		return
	}

	t, hp.err = time.Parse(time.RFC1123Z, s)
	if hp.err == nil {
		return t
	}

	t, hp.err = time.Parse("Mon, 2 Jan 2006 15:04:05 -0700", s)

	return
}

func (hp headerParser) parseMessageId(s string) string {
	if hp.err != nil {
		return ""
	}

	return strings.Trim(s, "<> ")
}

func (hp headerParser) parseMessageIdList(s string) (result []string) {
	if hp.err != nil {
		return
	}

	for _, p := range strings.Split(s, " ") {
		if strings.Trim(p, " \n") != "" {
			result = append(result, hp.parseMessageId(p))
		}
	}

	return
}

// Attachment with filename, content type and data (as a io.Reader)
type Attachment struct {
	Filename    string
	ContentType string
	Data        io.Reader
}

// EmbeddedFile with content id, content type and data (as a io.Reader)
type EmbeddedFile struct {
	CID         string
	ContentType string
	Data        io.Reader
}

// Email with fields for all the headers defined in RFC5322 with it's attachments and
type Email struct {
	Header mail.Header

	Subject    string
	Sender     *mail.Address
	From       []*mail.Address
	ReplyTo    []*mail.Address
	To         []*mail.Address
	Cc         []*mail.Address
	Bcc        []*mail.Address
	Date       time.Time
	MessageID  string
	InReplyTo  []string
	References []string

	ResentFrom      []*mail.Address
	ResentSender    *mail.Address
	ResentTo        []*mail.Address
	ResentDate      time.Time
	ResentCc        []*mail.Address
	ResentBcc       []*mail.Address
	ResentMessageID string

	HTMLBody string
	TextBody string

	Attachments   []Attachment
	EmbeddedFiles []EmbeddedFile
}
<commit_msg>check if alternative part contains related part<commit_after>package parsemail

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"mime/multipart"
	"net/mail"
	"strings"
	"time"
)

const contentTypeMultipartMixed = "multipart/mixed"
const contentTypeMultipartAlternative = "multipart/alternative"
const contentTypeMultipartRelated = "multipart/related"
const contentTypeTextHtml = "text/html"
const contentTypeTextPlain = "text/plain"

// Parse an email message read from io.Reader into parsemail.Email struct
func Parse(r io.Reader) (email Email, err error) {
	msg, err := mail.ReadMessage(r)
	if err != nil {
		return
	}

	email, err = createEmailFromHeader(msg.Header)
	if err != nil {
		return
	}

	contentType, params, err := parseContentType(msg.Header.Get("Content-Type"))
	if err != nil {
		return
	}

	switch contentType {
	case contentTypeMultipartMixed:
		email.TextBody, email.HTMLBody, email.Attachments, email.EmbeddedFiles, err = parseMultipartMixed(msg.Body, params["boundary"])
	case contentTypeMultipartAlternative:
		email.TextBody, email.HTMLBody, email.EmbeddedFiles, err = parseMultipartAlternative(msg.Body, params["boundary"])
	case contentTypeTextPlain:
		message, _ := ioutil.ReadAll(msg.Body)
		email.TextBody = strings.TrimSuffix(string(message[:]), "\n")
	case contentTypeTextHtml:
		message, _ := ioutil.ReadAll(msg.Body)
		email.HTMLBody = strings.TrimSuffix(string(message[:]), "\n")
	default:
		err = fmt.Errorf("Unknown top level mime type: %s", contentType)
	}

	return
}

func createEmailFromHeader(header mail.Header) (email Email, err error) {
	hp := headerParser{header: &header}

	email.Subject = decodeMimeSentence(header.Get("Subject"))
	email.From = hp.parseAddressList(header.Get("From"))
	email.Sender = hp.parseAddress(header.Get("Sender"))
	email.ReplyTo = hp.parseAddressList(header.Get("Reply-To"))
	email.To = hp.parseAddressList(header.Get("To"))
	email.Cc = hp.parseAddressList(header.Get("Cc"))
	email.Bcc = hp.parseAddressList(header.Get("Bcc"))
	email.Date = hp.parseTime(header.Get("Date"))
	email.ResentFrom = hp.parseAddressList(header.Get("Resent-From"))
	email.ResentSender = hp.parseAddress(header.Get("Resent-Sender"))
	email.ResentTo = hp.parseAddressList(header.Get("Resent-To"))
	email.ResentCc = hp.parseAddressList(header.Get("Resent-Cc"))
	email.ResentBcc = hp.parseAddressList(header.Get("Resent-Bcc"))
	email.ResentMessageID = hp.parseMessageId(header.Get("Resent-Message-ID"))
	email.MessageID = hp.parseMessageId(header.Get("Message-ID"))
	email.InReplyTo = hp.parseMessageIdList(header.Get("In-Reply-To"))
	email.References = hp.parseMessageIdList(header.Get("References"))
	email.ResentDate = hp.parseTime(header.Get("Resent-Date"))

	if hp.err != nil {
		err = hp.err
		return
	}

	//decode whole header for easier access to extra fields
	//todo: should we decode? aren't only standard fields mime encoded?
	email.Header, err = decodeHeaderMime(header)
	if err != nil {
		return
	}

	return
}

func parseContentType(contentTypeHeader string) (contentType string, params map[string]string, err error) {
	if contentTypeHeader == "" {
		contentType = contentTypeTextPlain
		return
	}

	return mime.ParseMediaType(contentTypeHeader)
}

func parseMultipartRelated(msg io.Reader, boundary string) (textBody, htmlBody string, embeddedFiles []EmbeddedFile, err error) {
	pmr := multipart.NewReader(msg, boundary)
	for {
		part, err := pmr.NextPart()

		if err == io.EOF {
			break
		} else if err != nil {
			return textBody, htmlBody, embeddedFiles, err
		}

		contentType, params, err := mime.ParseMediaType(part.Header.Get("Content-Type"))
		if err != nil {
			return textBody, htmlBody, embeddedFiles, err
		}

		switch contentType {
		case contentTypeTextPlain:
			ppContent, err := ioutil.ReadAll(part)
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			textBody += strings.TrimSuffix(string(ppContent[:]), "\n")
		case contentTypeTextHtml:
			ppContent, err := ioutil.ReadAll(part)
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			htmlBody += strings.TrimSuffix(string(ppContent[:]), "\n")
		case contentTypeMultipartAlternative:
			tb, hb, ef, err := parseMultipartAlternative(part, params["boundary"])
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			htmlBody += hb
			textBody += tb
			embeddedFiles = append(embeddedFiles, ef...)
		default:
			if isEmbeddedFile(part) {
				ef, err := decodeEmbeddedFile(part)
				if err != nil {
					return textBody, htmlBody, embeddedFiles, err
				}

				embeddedFiles = append(embeddedFiles, ef)
			} else {
				return textBody, htmlBody, embeddedFiles, fmt.Errorf("Can't process multipart/related inner mime type: %s", contentType)
			}
		}
	}

	return textBody, htmlBody, embeddedFiles, err
}

func parseMultipartAlternative(msg io.Reader, boundary string) (textBody, htmlBody string, embeddedFiles []EmbeddedFile, err error) {
	pmr := multipart.NewReader(msg, boundary)
	for {
		part, err := pmr.NextPart()

		if err == io.EOF {
			break
		} else if err != nil {
			return textBody, htmlBody, embeddedFiles, err
		}

		contentType, params, err := mime.ParseMediaType(part.Header.Get("Content-Type"))
		if err != nil {
			return textBody, htmlBody, embeddedFiles, err
		}

		switch contentType {
		case contentTypeTextPlain:
			ppContent, err := ioutil.ReadAll(part)
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			textBody += strings.TrimSuffix(string(ppContent[:]), "\n")
		case contentTypeTextHtml:
			ppContent, err := ioutil.ReadAll(part)
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			htmlBody += strings.TrimSuffix(string(ppContent[:]), "\n")
		case contentTypeMultipartRelated:
			tb, hb, ef, err := parseMultipartRelated(part, params["boundary"])
			if err != nil {
				return textBody, htmlBody, embeddedFiles, err
			}

			htmlBody += hb
			textBody += tb
			embeddedFiles = append(embeddedFiles, ef...)
		default:
			if isEmbeddedFile(part) {
				ef, err := decodeEmbeddedFile(part)
				if err != nil {
					return textBody, htmlBody, embeddedFiles, err
				}

				embeddedFiles = append(embeddedFiles, ef)
			} else {
				return textBody, htmlBody, embeddedFiles, fmt.Errorf("Can't process multipart/alternative inner mime type: %s", contentType)
			}
		}
	}

	return textBody, htmlBody, embeddedFiles, err
}

func parseMultipartMixed(msg io.Reader, boundary string) (textBody, htmlBody string, attachments []Attachment, embeddedFiles []EmbeddedFile, err error) {
	mr := multipart.NewReader(msg, boundary)
	for {
		part, err := mr.NextPart()
		if err == io.EOF {
			break
		} else if err != nil {
			return textBody, htmlBody, attachments, embeddedFiles, err
		}

		contentType, params, err := mime.ParseMediaType(part.Header.Get("Content-Type"))
		if err != nil {
			return textBody, htmlBody, attachments, embeddedFiles, err
		}

		if contentType == contentTypeMultipartAlternative {
			textBody, htmlBody, embeddedFiles, err = parseMultipartAlternative(part, params["boundary"])
			if err != nil {
				return textBody, htmlBody, attachments, embeddedFiles, err
			}
		} else if contentType == contentTypeMultipartRelated {
			textBody, htmlBody, embeddedFiles, err = parseMultipartRelated(part, params["boundary"])
			if err != nil {
				return textBody, htmlBody, attachments, embeddedFiles, err
			}
		} else if isAttachment(part) {
			at, err := decodeAttachment(part)
			if err != nil {
				return textBody, htmlBody, attachments, embeddedFiles, err
			}

			attachments = append(attachments, at)
		} else {
			return textBody, htmlBody, attachments, embeddedFiles, fmt.Errorf("Unknown multipart/mixed nested mime type: %s", contentType)
		}
	}

	return textBody, htmlBody, attachments, embeddedFiles, err
}

func decodeMimeSentence(s string) string {
	result := []string{}
	ss := strings.Split(s, " ")

	for _, word := range ss {
		dec := new(mime.WordDecoder)
		w, err := dec.Decode(word)
		if err != nil {
			if len(result) == 0 {
				w = word
			} else {
				w = " " + word
			}
		}

		result = append(result, w)
	}

	return strings.Join(result, "")
}

func decodeHeaderMime(header mail.Header) (mail.Header, error) {
	parsedHeader := map[string][]string{}

	for headerName, headerData := range header {

		parsedHeaderData := []string{}
		for _, headerValue := range headerData {
			parsedHeaderData = append(parsedHeaderData, decodeMimeSentence(headerValue))
		}

		parsedHeader[headerName] = parsedHeaderData
	}

	return mail.Header(parsedHeader), nil
}

func decodePartData(part *multipart.Part) (io.Reader, error) {
	encoding := part.Header.Get("Content-Transfer-Encoding")

	if strings.EqualFold(encoding, "base64") {
		dr := base64.NewDecoder(base64.StdEncoding, part)
		dd, err := ioutil.ReadAll(dr)
		if err != nil {
			return nil, err
		}

		return bytes.NewReader(dd), nil
	}

	return nil, fmt.Errorf("Unknown encoding: %s", encoding)
}

func isEmbeddedFile(part *multipart.Part) bool {
	return part.Header.Get("Content-Transfer-Encoding") != ""
}

func decodeEmbeddedFile(part *multipart.Part) (ef EmbeddedFile, err error) {
	cid := decodeMimeSentence(part.Header.Get("Content-Id"))
	decoded, err := decodePartData(part)
	if err != nil {
		return
	}

	ef.CID = strings.Trim(cid, "<>")
	ef.Data = decoded
	ef.ContentType = part.Header.Get("Content-Type")

	return
}

func isAttachment(part *multipart.Part) bool {
	return part.FileName() != ""
}

func decodeAttachment(part *multipart.Part) (at Attachment, err error) {
	filename := decodeMimeSentence(part.FileName())
	decoded, err := decodePartData(part)
	if err != nil {
		return
	}

	at.Filename = filename
	at.Data = decoded
	at.ContentType = strings.Split(part.Header.Get("Content-Type"), ";")[0]

	return
}

type headerParser struct {
	header *mail.Header
	err    error
}

func (hp headerParser) parseAddress(s string) (ma *mail.Address) {
	if hp.err != nil {
		return nil
	}

	if strings.Trim(s, " \n") != "" {
		ma, hp.err = mail.ParseAddress(s)

		return ma
	}

	return nil
}

func (hp headerParser) parseAddressList(s string) (ma []*mail.Address) {
	if hp.err != nil {
		return
	}

	if strings.Trim(s, " \n") != "" {
		ma, hp.err = mail.ParseAddressList(s)
		return
	}

	return
}

func (hp headerParser) parseTime(s string) (t time.Time) {
	if hp.err != nil || s == "" {
		return
	}

	t, hp.err = time.Parse(time.RFC1123Z, s)
	if hp.err == nil {
		return t
	}

	t, hp.err = time.Parse("Mon, 2 Jan 2006 15:04:05 -0700", s)

	return
}

func (hp headerParser) parseMessageId(s string) string {
	if hp.err != nil {
		return ""
	}

	return strings.Trim(s, "<> ")
}

func (hp headerParser) parseMessageIdList(s string) (result []string) {
	if hp.err != nil {
		return
	}

	for _, p := range strings.Split(s, " ") {
		if strings.Trim(p, " \n") != "" {
			result = append(result, hp.parseMessageId(p))
		}
	}

	return
}

// Attachment with filename, content type and data (as a io.Reader)
type Attachment struct {
	Filename    string
	ContentType string
	Data        io.Reader
}

// EmbeddedFile with content id, content type and data (as a io.Reader)
type EmbeddedFile struct {
	CID         string
	ContentType string
	Data        io.Reader
}

// Email with fields for all the headers defined in RFC5322 with it's attachments and
type Email struct {
	Header mail.Header

	Subject    string
	Sender     *mail.Address
	From       []*mail.Address
	ReplyTo    []*mail.Address
	To         []*mail.Address
	Cc         []*mail.Address
	Bcc        []*mail.Address
	Date       time.Time
	MessageID  string
	InReplyTo  []string
	References []string

	ResentFrom      []*mail.Address
	ResentSender    *mail.Address
	ResentTo        []*mail.Address
	ResentDate      time.Time
	ResentCc        []*mail.Address
	ResentBcc       []*mail.Address
	ResentMessageID string

	HTMLBody string
	TextBody string

	Attachments   []Attachment
	EmbeddedFiles []EmbeddedFile
}
<|endoftext|>


<commit_before>package easyhmac

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

var (
	testmessage = "eyJwIjoiZEdWemRDQnRaWE56WVdkbCIsInMiOiJZbTVpYVUxUmRrNURkVEUzY0U4clFVRkVaV1Z5Wnk5bmFYZFZjV3RPTUVweVFYTlRkRWRpUVUxcVVUMD0ifQ=="
	Secret      = "test"
)

func TestEncode(t *testing.T) {

	// Initialize SignedMessage struct with secret
	key := SignedMessage{}

	// Add payload data
	key.Payload = []byte("test message")

	// Create HMAC signature
	key.Sign()

	// Marshal message to JSON and encode in url-safe base64
	signedkey, err := key.Encode()

	assert.NoError(t, err, "should be no error")

	assert.Equal(t, testmessage, signedkey, "they should be equal")

}

func TestDecode(t *testing.T) {

	// Initialize SignedMessage struct with secret
	key := SignedMessage{}

	// Decode message
	err := message.Decode(testmessage)

	assert.NoError(t, err, "should be no error")

	check := message.Verify()

	assert.True(t, check, "should be true")

	assert.Equal(t, "test message", key.Payload, "they should be equal")

}
<commit_msg>add test<commit_after>package easyhmac

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

var testmessage = "eyJwIjoiZEdWemRDQnRaWE56WVdkbCIsInMiOiJZbTVpYVUxUmRrNURkVEUzY0U4clFVRkVaV1Z5Wnk5bmFYZFZjV3RPTUVweVFYTlRkRWRpUVUxcVVUMD0ifQ=="

func TestEncode(t *testing.T) {

	Secret = "test"

	// Initialize SignedMessage struct with secret
	key := SignedMessage{}

	// Add payload data
	key.Payload = []byte("test message")

	// Create HMAC signature
	key.Sign()

	// Marshal message to JSON and encode in url-safe base64
	signedkey, err := key.Encode()

	assert.NoError(t, err, "should be no error")

	assert.Equal(t, testmessage, signedkey, "they should be equal")

}

func TestDecode(t *testing.T) {

	Secret = "test"

	// Initialize SignedMessage struct with secret
	key := SignedMessage{}

	// Decode message
	err := key.Decode(testmessage)

	assert.NoError(t, err, "should be no error")

	check := key.Verify()

	assert.True(t, check, "should be true")

	assert.Equal(t, "test message", key.Payload, "they should be equal")

}
<|endoftext|>


<commit_before>package gtnd

type SearchParam struct {
	Keyword string

}
<commit_msg>add parameter<commit_after>package gtnd

type SearchParam struct {
	Keyword string
	EventId int64
	NickName string
	Count int64
	Ymd int64
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

// testExit is a channel that is used during testing only to trigger a return from the alternate
// function.
var testExit chan struct{}

// alternate runs a command with alternating parameters inserted in place of the placeholder. Each
// time a SIGUSR1 is received, a new command is run with the next parameter, and a SIGINT is sent to
// the previous command after the overlap duration has elapsed. The alternate logs are written to
// stderr, and the command logs are written to cmdStdout and cmdStderr.
func alternate(command string, placeholder string, params []string, overlap time.Duration,
	stderr, cmdStdout, cmdStderr io.Writer) {

	log.SetPrefix("alternate | ")
	log.SetOutput(stderr)
	log.SetFlags(0)

	log.Printf("Starting with command %q, placeholder %q, params = %q, overlap = %v\n",
		command, placeholder, params, overlap)

	// When a command exits, cmdExit receives the parameter this command was run with.
	cmdExit := make(chan string)

	// When the overlap duration has elapsed, overlapEnd receives an empty struct.
	overlapEnd := make(chan struct{})

	next := make(chan os.Signal, 1)
	// Buffer a fake signal to run the first command.
	next <- syscall.Signal(0)
	// Then listen to USR1 signals dispatched by the user.
	signal.Notify(next, syscall.SIGUSR1)

	m := newManager(params)

	for {
		select {
		case <-testExit:
			log.Println("Test exit channel received a value, exiting alternate")
			killAllCmds(params, m)
			return

		case param := <-cmdExit:
			log.Printf("Command with parameter %q exited\n", param)
			m.unsetCmd(param)
			if !m.hasCmds() {
				log.Println("All commands have exited, exiting alternate")
				return
			}

		case <-overlapEnd:
			finishRotation(m)

		case signal := <-next:
			first := signal != syscall.SIGUSR1
			nextParam := m.nextParam()

			if !first {
				log.Printf("Received signal USR1, rotating to next parameter %q", nextParam)
			}

			if m.nextCmd() != nil {
				log.Printf("A command with parameter %q is already running, cannot run again",
					nextParam)
				break
			}

			s := strings.Replace(command, placeholder, nextParam, 1)
			nextCmd := cmd(s, cmdStdout, cmdStderr)
			m.setCmd(nextParam, nextCmd)

			if err := run(nextCmd, nextParam, cmdExit); err != nil {
				log.Printf("Failed to run the command with parameter %q, error: %v\n",
					err.Error())
				m.unsetCmd(nextParam)
				break
			}

			if first {
				// There is no previous command to interrupt.
				m.rotate()
				break
			}

			if overlap == 0 {
				finishRotation(m)
			} else {
				log.Printf("Waiting %v before sending interrupt to command with parameter %q\n",
					overlap, m.currentParam())
				go func() {
					time.Sleep(overlap)
					overlapEnd <- struct{}{}
				}()
			}
		}
	}
}

func finishRotation(m *manager) {
	if m.nextCmd() == nil {
		// The next command is not running, cancel the rotation without interrupting the current
		// command.
		return
	}
	if c := m.currentCmd(); c != nil {
		currentParam := m.currentParam()
		log.Printf("Sending interrupt to command with parameter %q\n", currentParam)
		if err := interruptCmd(c); err != nil {
			log.Printf("Failed to send interrupt to command with parameter %q, error: %v\n",
				currentParam, err)
		}
	}
	m.rotate()
}

func interruptCmd(c *exec.Cmd) error {
	if c == nil {
		return errors.New("interruptCmd error: cmd is nil")
	}
	p := c.Process
	if p == nil {
		return errors.New("interruptCmd error: cmd.Process is nil")
	}
	p.Signal(os.Interrupt)
	return nil
}

func cmd(s string, stdout, stderr io.Writer) *exec.Cmd {
	f := strings.Fields(s)
	c := exec.Command(f[0], f[1:]...)
	c.Stdout = stdout
	c.Stderr = stderr
	return c
}

func run(c *exec.Cmd, param string, exit chan string) error {
	log.Printf("Running command with parameter %q\n", param)
	if err := c.Start(); err != nil {
		return err
	}
	go func() {
		c.Wait()
		exit <- param
	}()
	return nil
}

func killAllCmds(params []string, m *manager) {
	for _, param := range params {
		if c := m.cmd(param); c != nil {
			if p := c.Process; p != nil {
				log.Printf("Killing command with parameter %q\n", param)
				p.Signal(os.Kill)
			}
		}
	}
}
<commit_msg>Update comments and logging<commit_after>package main

import (
	"errors"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

// testExit is a channel that is used during testing only to trigger a return from the alternate
// function.
var testExit chan struct{}

// alternate runs a command with alternating parameters inserted in place of the placeholder. Each
// time a SIGUSR1 is received, a new command is run with the next parameter, and a SIGINT is sent to
// the previous command after the overlap duration has elapsed. The alternate logs are written to
// stderr, and the command logs are written to cmdStdout and cmdStderr.
func alternate(command string, placeholder string, params []string, overlap time.Duration,
	stderr, cmdStdout, cmdStderr io.Writer) {

	log.SetPrefix("alternate | ")
	log.SetOutput(stderr)
	log.SetFlags(0)

	log.Printf("Starting with command %q, placeholder %q, params = %q, overlap = %v\n",
		command, placeholder, params, overlap)

	// When a command exits, cmdExit receives the parameter this command was run with.
	cmdExit := make(chan string)

	// When the overlap duration has elapsed, overlapEnd receives an empty struct.
	overlapEnd := make(chan struct{})

	next := make(chan os.Signal, 1)
	// Buffer a fake signal to run the first command.
	next <- syscall.Signal(0)
	// Then listen to USR1 signals dispatched by the user.
	signal.Notify(next, syscall.SIGUSR1)

	m := newManager(params)

	for {
		select {
		case <-testExit:
			log.Println("Test exit channel received a value, exiting alternate")
			killAllCmds(params, m)
			return

		case param := <-cmdExit:
			log.Printf("Command with parameter %q exited\n", param)
			m.unsetCmd(param)
			if !m.hasCmds() {
				log.Println("All commands have exited, exiting alternate")
				return
			}

		case <-overlapEnd:
			finishRotation(m)

		case signal := <-next:
			first := signal != syscall.SIGUSR1
			nextParam := m.nextParam()

			if !first {
				log.Printf("Received signal USR1, rotating to next parameter %q", nextParam)
			}

			if m.nextCmd() != nil {
				log.Printf("A command with parameter %q is already running, cannot run again",
					nextParam)
				break
			}

			s := strings.Replace(command, placeholder, nextParam, 1)
			nextCmd := cmd(s, cmdStdout, cmdStderr)
			m.setCmd(nextParam, nextCmd)

			log.Printf("Running command with parameter %q\n", nextParam)
			if err := run(nextCmd, nextParam, cmdExit); err != nil {
				log.Printf("Failed to run the command with parameter %q, error: %v\n",
					err.Error())
				m.unsetCmd(nextParam)
				break
			}

			if first {
				// There is no previous command to interrupt.
				m.rotate()
				break
			}

			if overlap == 0 {
				finishRotation(m)
			} else {
				log.Printf("Waiting %v before sending interrupt to command with parameter %q\n",
					overlap, m.currentParam())
				go func() {
					time.Sleep(overlap)
					overlapEnd <- struct{}{}
				}()
			}
		}
	}
}

func finishRotation(m *manager) {
	if m.nextCmd() == nil {
		// The next command is not running. Cancel the rotation without interrupting the current
		// command.
		return
	}
	if c := m.currentCmd(); c != nil {
		currentParam := m.currentParam()
		log.Printf("Sending interrupt to command with parameter %q\n", currentParam)
		if err := interruptCmd(c); err != nil {
			log.Printf("Failed to send interrupt to command with parameter %q, error: %v\n",
				currentParam, err)
		}
	}
	m.rotate()
}

// interruptCmd sends SIGINT to a command.
func interruptCmd(c *exec.Cmd) error {
	if c == nil {
		return errors.New("interruptCmd error: cmd is nil")
	}
	p := c.Process
	if p == nil {
		return errors.New("interruptCmd error: cmd.Process is nil")
	}
	p.Signal(os.Interrupt)
	return nil
}

// cmd returns a command built from the given string. The command will print to the given stdout and
// stderr.
func cmd(s string, stdout, stderr io.Writer) *exec.Cmd {
	f := strings.Fields(s)
	c := exec.Command(f[0], f[1:]...)
	c.Stdout = stdout
	c.Stderr = stderr
	return c
}

// run runs a command without blocking. After the command exits, it sends exitMsg on the exit
// channel.
func run(c *exec.Cmd, exitMsg string, exit chan string) error {
	if err := c.Start(); err != nil {
		return err
	}
	go func() {
		c.Wait()
		exit <- exitMsg
	}()
	return nil
}

// killAllCmds sends a SIGKILL to all commands in the manager.
func killAllCmds(params []string, m *manager) {
	for _, param := range params {
		if c := m.cmd(param); c != nil {
			if p := c.Process; p != nil {
				log.Printf("Killing command with parameter %q\n", param)
				p.Signal(os.Kill)
			}
		}
	}
}
<|endoftext|>


<commit_before>// +build !providerless

/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package azure

import (
	"net/http"
	"testing"
	"time"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
	"github.com/Azure/go-autorest/autorest/to"
	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"

	"k8s.io/apimachinery/pkg/types"
	cloudprovider "k8s.io/cloud-provider"
	azcache "k8s.io/legacy-cloud-providers/azure/cache"
	"k8s.io/legacy-cloud-providers/azure/clients/vmclient/mockvmclient"
	"k8s.io/legacy-cloud-providers/azure/retry"
)

var (
	fakeCacheTTL = 2 * time.Second
)

func TestStandardAttachDisk(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	testCases := []struct {
		desc        string
		nodeName    types.NodeName
		expectedErr bool
	}{
		{
			desc:        "an error shall be returned if there's no corresponding vms",
			nodeName:    "vm2",
			expectedErr: true,
		},
		{
			desc:        "no error shall be returned if everything's good",
			nodeName:    "vm1",
			expectedErr: false,
		},
	}

	for i, test := range testCases {
		testCloud := GetTestCloud(ctrl)
		vmSet := testCloud.vmSet
		expectedVMs := setTestVirtualMachines(testCloud, map[string]string{"vm1": "PowerState/Running"}, false)
		mockVMsClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
		for _, vm := range expectedVMs {
			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
		}
		mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, "vm2", gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
		mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()

		err := vmSet.AttachDisk(true, "",
			"uri", test.nodeName, 0, compute.CachingTypesReadOnly, "", false)
		assert.Equal(t, test.expectedErr, err != nil, "TestCase[%d]: %s, err: %v", i, test.desc, err)
	}
}

func TestStandardDetachDisk(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	testCases := []struct {
		desc          string
		nodeName      types.NodeName
		diskName      string
		expectedError bool
	}{
		{
			desc:          "no error shall be returned if there's no corresponding vm",
			nodeName:      "vm2",
			expectedError: false,
		},
		{
			desc:          "no error shall be returned if there's no corresponding disk",
			nodeName:      "vm1",
			diskName:      "disk2",
			expectedError: false,
		},
		{
			desc:          "no error shall be returned if there's a corresponding disk",
			nodeName:      "vm1",
			diskName:      "disk1",
			expectedError: false,
		},
	}

	for i, test := range testCases {
		testCloud := GetTestCloud(ctrl)
		vmSet := testCloud.vmSet
		expectedVMs := setTestVirtualMachines(testCloud, map[string]string{"vm1": "PowerState/Running"}, false)
		mockVMsClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
		for _, vm := range expectedVMs {
			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
		}
		mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, "vm2", gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
		mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()

		err := vmSet.DetachDisk(test.diskName, "", test.nodeName)
		assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)
	}
}

func TestGetDataDisks(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var testCases = []struct {
		desc              string
		nodeName          types.NodeName
		expectedDataDisks []compute.DataDisk
		expectedError     bool
		crt               azcache.AzureCacheReadType
	}{
		{
			desc:              "an error shall be returned if there's no corresponding vm",
			nodeName:          "vm2",
			expectedDataDisks: nil,
			expectedError:     true,
			crt:               azcache.CacheReadTypeDefault,
		},
		{
			desc:     "correct list of data disks shall be returned if everything is good",
			nodeName: "vm1",
			expectedDataDisks: []compute.DataDisk{
				{
					Lun:  to.Int32Ptr(0),
					Name: to.StringPtr("disk1"),
				},
			},
			expectedError: false,
			crt:           azcache.CacheReadTypeDefault,
		},
		{
			desc:     "correct list of data disks shall be returned if everything is good",
			nodeName: "vm1",
			expectedDataDisks: []compute.DataDisk{
				{
					Lun:  to.Int32Ptr(0),
					Name: to.StringPtr("disk1"),
				},
			},
			expectedError: false,
			crt:           azcache.CacheReadTypeUnsafe,
		},
	}
	for i, test := range testCases {
		testCloud := GetTestCloud(ctrl)
		vmSet := testCloud.vmSet
		expectedVMs := setTestVirtualMachines(testCloud, map[string]string{"vm1": "PowerState/Running"}, false)
		mockVMsClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
		for _, vm := range expectedVMs {
			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
		}
		mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, gomock.Not("vm1"), gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
		mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()

		dataDisks, err := vmSet.GetDataDisks(test.nodeName, test.crt)
		assert.Equal(t, test.expectedDataDisks, dataDisks, "TestCase[%d]: %s", i, test.desc)
		assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)

		if test.crt == azcache.CacheReadTypeUnsafe {
			time.Sleep(fakeCacheTTL)
			dataDisks, err := vmSet.GetDataDisks(test.nodeName, test.crt)
			assert.Equal(t, test.expectedDataDisks, dataDisks, "TestCase[%d]: %s", i, test.desc)
			assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)
		}
	}
}
<commit_msg>Improves unittest CC for azure_controller_standard<commit_after>// +build !providerless

/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package azure

import (
	"net/http"
	"testing"
	"time"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
	"github.com/Azure/go-autorest/autorest/to"
	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"

	"k8s.io/apimachinery/pkg/types"
	cloudprovider "k8s.io/cloud-provider"
	azcache "k8s.io/legacy-cloud-providers/azure/cache"
	"k8s.io/legacy-cloud-providers/azure/clients/vmclient/mockvmclient"
	"k8s.io/legacy-cloud-providers/azure/retry"
)

var (
	fakeCacheTTL = 2 * time.Second
)

func TestStandardAttachDisk(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	testCases := []struct {
		desc          string
		nodeName      types.NodeName
		isManagedDisk bool
		isAttachFail  bool
		expectedErr   bool
	}{
		{
			desc:          "an error shall be returned if there's no corresponding vms",
			nodeName:      "vm2",
			isManagedDisk: true,
			expectedErr:   true,
		},
		{
			desc:          "no error shall be returned if everything's good",
			nodeName:      "vm1",
			isManagedDisk: true,
			expectedErr:   false,
		},
		{
			desc:          "no error shall be returned if everything's good with non managed disk",
			nodeName:      "vm1",
			isManagedDisk: false,
			expectedErr:   false,
		},
		{
			desc:          "an error shall be returned if update attach disk failed",
			nodeName:      "vm1",
			isManagedDisk: true,
			isAttachFail:  true,
			expectedErr:   true,
		},
	}

	for i, test := range testCases {
		testCloud := GetTestCloud(ctrl)
		vmSet := testCloud.vmSet
		expectedVMs := setTestVirtualMachines(testCloud, map[string]string{"vm1": "PowerState/Running"}, false)
		mockVMsClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
		for _, vm := range expectedVMs {
			vm.StorageProfile = &compute.StorageProfile{
				OsDisk: &compute.OSDisk{
					Name: to.StringPtr("osdisk1"),
					ManagedDisk: &compute.ManagedDiskParameters{
						ID: to.StringPtr("ManagedID"),
						DiskEncryptionSet: &compute.DiskEncryptionSetParameters{
							ID: to.StringPtr("DiskEncryptionSetID"),
						},
					},
				},
				DataDisks: &[]compute.DataDisk{},
			}
			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
		}
		mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, "vm2", gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
		if test.isAttachFail {
			mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(&retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
		} else {
			mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
		}

		err := vmSet.AttachDisk(test.isManagedDisk, "",
			"uri", test.nodeName, 0, compute.CachingTypesReadOnly, "", false)
		assert.Equal(t, test.expectedErr, err != nil, "TestCase[%d]: %s, err: %v", i, test.desc, err)
	}
}

func TestStandardDetachDisk(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	testCases := []struct {
		desc          string
		nodeName      types.NodeName
		diskName      string
		isDetachFail  bool
		expectedError bool
	}{
		{
			desc:          "no error shall be returned if there's no corresponding vm",
			nodeName:      "vm2",
			expectedError: false,
		},
		{
			desc:          "no error shall be returned if there's no corresponding disk",
			nodeName:      "vm1",
			diskName:      "disk2",
			expectedError: false,
		},
		{
			desc:          "no error shall be returned if there's a corresponding disk",
			nodeName:      "vm1",
			diskName:      "disk1",
			expectedError: false,
		},
		{
			desc:          "an error shall be returned if detach disk failed",
			nodeName:      "vm1",
			isDetachFail:  true,
			expectedError: true,
		},
	}

	for i, test := range testCases {
		testCloud := GetTestCloud(ctrl)
		vmSet := testCloud.vmSet
		expectedVMs := setTestVirtualMachines(testCloud, map[string]string{"vm1": "PowerState/Running"}, false)
		mockVMsClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
		for _, vm := range expectedVMs {
			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
		}
		mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, "vm2", gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
		if test.isDetachFail {
			mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(&retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
		} else {
			mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
		}

		err := vmSet.DetachDisk(test.diskName, "", test.nodeName)
		assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)
	}
}

func TestGetDataDisks(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var testCases = []struct {
		desc              string
		nodeName          types.NodeName
		isDataDiskNull    bool
		expectedDataDisks []compute.DataDisk
		expectedError     bool
		crt               azcache.AzureCacheReadType
	}{
		{
			desc:              "an error shall be returned if there's no corresponding vm",
			nodeName:          "vm2",
			expectedDataDisks: nil,
			expectedError:     true,
			crt:               azcache.CacheReadTypeDefault,
		},
		{
			desc:     "correct list of data disks shall be returned if everything is good",
			nodeName: "vm1",
			expectedDataDisks: []compute.DataDisk{
				{
					Lun:  to.Int32Ptr(0),
					Name: to.StringPtr("disk1"),
				},
			},
			expectedError: false,
			crt:           azcache.CacheReadTypeDefault,
		},
		{
			desc:     "correct list of data disks shall be returned if everything is good",
			nodeName: "vm1",
			expectedDataDisks: []compute.DataDisk{
				{
					Lun:  to.Int32Ptr(0),
					Name: to.StringPtr("disk1"),
				},
			},
			expectedError: false,
			crt:           azcache.CacheReadTypeUnsafe,
		},
		{
			desc:              "nil shall be returned if DataDisk is null",
			nodeName:          "vm1",
			isDataDiskNull:    true,
			expectedDataDisks: nil,
			expectedError:     false,
			crt:               azcache.CacheReadTypeDefault,
		},
	}
	for i, test := range testCases {
		testCloud := GetTestCloud(ctrl)
		vmSet := testCloud.vmSet
		expectedVMs := setTestVirtualMachines(testCloud, map[string]string{"vm1": "PowerState/Running"}, false)
		mockVMsClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
		for _, vm := range expectedVMs {
			if test.isDataDiskNull {
				vm.StorageProfile = &compute.StorageProfile{}
			}
			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
		}
		mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, gomock.Not("vm1"), gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
		mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()

		dataDisks, err := vmSet.GetDataDisks(test.nodeName, test.crt)
		assert.Equal(t, test.expectedDataDisks, dataDisks, "TestCase[%d]: %s", i, test.desc)
		assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)

		if test.crt == azcache.CacheReadTypeUnsafe {
			time.Sleep(fakeCacheTTL)
			dataDisks, err := vmSet.GetDataDisks(test.nodeName, test.crt)
			assert.Equal(t, test.expectedDataDisks, dataDisks, "TestCase[%d]: %s", i, test.desc)
			assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)
		}
	}
}
<|endoftext|>


<commit_before>package screenshot

import (
	"errors"
	"github.com/kbinani/screenshot/internal/util"
	win "github.com/lxn/win"
	"image"
	"syscall"
	"unsafe"
)

var (
	libUser32, _               = syscall.LoadLibrary("user32.dll")
	funcGetDesktopWindow, _    = syscall.GetProcAddress(syscall.Handle(libUser32), "GetDesktopWindow")
	funcEnumDisplayMonitors, _ = syscall.GetProcAddress(syscall.Handle(libUser32), "EnumDisplayMonitors")
)

func Capture(x, y, width, height int) (*image.RGBA, error) {
	rect := image.Rect(0, 0, width, height)
	img, err := util.CreateImage(rect)
	if err != nil {
		return nil, err
	}

	hwnd := getDesktopWindow()
	hdc := win.GetDC(hwnd)
	if hdc == 0 {
		return nil, errors.New("GetDC failed")
	}
	defer win.ReleaseDC(hwnd, hdc)

	memory_device := win.CreateCompatibleDC(hdc)
	if memory_device == 0 {
		return nil, errors.New("CreateCompatibleDC failed")
	}
	defer win.DeleteDC(memory_device)

	bitmap := win.CreateCompatibleBitmap(hdc, int32(width), int32(height))
	if bitmap == 0 {
		return nil, errors.New("CreateCompatibleBitmap failed")
	}
	defer win.DeleteObject(win.HGDIOBJ(bitmap))

	var header win.BITMAPINFOHEADER
	header.BiSize = uint32(unsafe.Sizeof(header))
	header.BiPlanes = 1
	header.BiBitCount = 32
	header.BiWidth = int32(width)
	header.BiHeight = int32(-height)
	header.BiCompression = win.BI_RGB
	header.BiSizeImage = 0

	// GetDIBits balks at using Go memory on some systems. The MSDN example uses
	// GlobalAlloc, so we'll do that too. See:
	// https://docs.microsoft.com/en-gb/windows/desktop/gdi/capturing-an-image
	bitmapDataSize := uintptr(((int64(width)*int64(header.BiBitCount) + 31) / 32) * 4 * int64(height))
	hmem := win.GlobalAlloc(win.GMEM_MOVEABLE, bitmapDataSize)
	defer win.GlobalFree(hmem)
	memptr := win.GlobalLock(hmem)
	defer win.GlobalUnlock(hmem)

	old := win.SelectObject(memory_device, win.HGDIOBJ(bitmap))
	if old == 0 {
		return nil, errors.New("SelectObject failed")
	}
	defer win.SelectObject(memory_device, old)

	if !win.BitBlt(memory_device, 0, 0, int32(width), int32(height), hdc, int32(x), int32(y), win.SRCCOPY) {
		return nil, errors.New("BitBlt failed")
	}

	if win.GetDIBits(hdc, bitmap, 0, uint32(height), (*uint8)(memptr), (*win.BITMAPINFO)(unsafe.Pointer(&header)), win.DIB_RGB_COLORS) == 0 {
		return nil, errors.New("GetDIBits failed")
	}

	i := 0
	src := uintptr(memptr)
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			v0 := *(*uint8)(unsafe.Pointer(src))
			v1 := *(*uint8)(unsafe.Pointer(src + 1))
			v2 := *(*uint8)(unsafe.Pointer(src + 2))

			// BGRA => RGBA, and set A to 255
			img.Pix[i], img.Pix[i+1], img.Pix[i+2], img.Pix[i+3] = v2, v1, v0, 255

			i += 4
			src += 4
		}
	}

	return img, nil
}

func NumActiveDisplays() int {
	var count int = 0
	enumDisplayMonitors(win.HDC(0), nil, syscall.NewCallback(countupMonitorCallback), uintptr(unsafe.Pointer(&count)))
	return count
}

func GetDisplayBounds(displayIndex int) image.Rectangle {
	var ctx getMonitorBoundsContext
	ctx.Index = displayIndex
	ctx.Count = 0
	enumDisplayMonitors(win.HDC(0), nil, syscall.NewCallback(getMonitorBoundsCallback), uintptr(unsafe.Pointer(&ctx)))
	return image.Rect(
		int(ctx.Rect.Left), int(ctx.Rect.Top),
		int(ctx.Rect.Right), int(ctx.Rect.Bottom))
}

func getDesktopWindow() win.HWND {
	ret, _, _ := syscall.Syscall(funcGetDesktopWindow, 0, 0, 0, 0)
	return win.HWND(ret)
}

func enumDisplayMonitors(hdc win.HDC, lprcClip *win.RECT, lpfnEnum uintptr, dwData uintptr) bool {
	ret, _, _ := syscall.Syscall6(funcEnumDisplayMonitors, 4,
		uintptr(hdc),
		uintptr(unsafe.Pointer(lprcClip)),
		lpfnEnum,
		dwData,
		0,
		0)
	return int(ret) != 0
}

func countupMonitorCallback(hMonitor win.HMONITOR, hdcMonitor win.HDC, lprcMonitor *win.RECT, dwData uintptr) uintptr {
	var count *int
	count = (*int)(unsafe.Pointer(dwData))
	*count = *count + 1
	return uintptr(1)
}

type getMonitorBoundsContext struct {
	Index int
	Rect  win.RECT
	Count int
}

func getMonitorBoundsCallback(hMonitor win.HMONITOR, hdcMonitor win.HDC, lprcMonitor *win.RECT, dwData uintptr) uintptr {
	var ctx *getMonitorBoundsContext
	ctx = (*getMonitorBoundsContext)(unsafe.Pointer(dwData))
	if ctx.Count == ctx.Index {
		ctx.Rect = *lprcMonitor
		return uintptr(0)
	} else {
		ctx.Count = ctx.Count + 1
		return uintptr(1)
	}
}
<commit_msg>bugfix(windows): Use real dimensions of a screen to support scaling<commit_after>package screenshot

import (
	"errors"
	"github.com/kbinani/screenshot/internal/util"
	win "github.com/lxn/win"
	"image"
	"syscall"
	"unsafe"
)

var (
	libUser32, _               = syscall.LoadLibrary("user32.dll")
	funcGetDesktopWindow, _    = syscall.GetProcAddress(syscall.Handle(libUser32), "GetDesktopWindow")
	funcEnumDisplayMonitors, _ = syscall.GetProcAddress(syscall.Handle(libUser32), "EnumDisplayMonitors")
	funcGetMonitorInfo, _      = syscall.GetProcAddress(syscall.Handle(libUser32), "GetMonitorInfoW")
	funcEnumDisplaySettings, _ = syscall.GetProcAddress(syscall.Handle(libUser32), "EnumDisplaySettingsW")
)

func Capture(x, y, width, height int) (*image.RGBA, error) {
	rect := image.Rect(0, 0, width, height)
	img, err := util.CreateImage(rect)
	if err != nil {
		return nil, err
	}

	hwnd := getDesktopWindow()
	hdc := win.GetDC(hwnd)
	if hdc == 0 {
		return nil, errors.New("GetDC failed")
	}
	defer win.ReleaseDC(hwnd, hdc)

	memory_device := win.CreateCompatibleDC(hdc)
	if memory_device == 0 {
		return nil, errors.New("CreateCompatibleDC failed")
	}
	defer win.DeleteDC(memory_device)

	bitmap := win.CreateCompatibleBitmap(hdc, int32(width), int32(height))
	if bitmap == 0 {
		return nil, errors.New("CreateCompatibleBitmap failed")
	}
	defer win.DeleteObject(win.HGDIOBJ(bitmap))

	var header win.BITMAPINFOHEADER
	header.BiSize = uint32(unsafe.Sizeof(header))
	header.BiPlanes = 1
	header.BiBitCount = 32
	header.BiWidth = int32(width)
	header.BiHeight = int32(-height)
	header.BiCompression = win.BI_RGB
	header.BiSizeImage = 0

	// GetDIBits balks at using Go memory on some systems. The MSDN example uses
	// GlobalAlloc, so we'll do that too. See:
	// https://docs.microsoft.com/en-gb/windows/desktop/gdi/capturing-an-image
	bitmapDataSize := uintptr(((int64(width)*int64(header.BiBitCount) + 31) / 32) * 4 * int64(height))
	hmem := win.GlobalAlloc(win.GMEM_MOVEABLE, bitmapDataSize)
	defer win.GlobalFree(hmem)
	memptr := win.GlobalLock(hmem)
	defer win.GlobalUnlock(hmem)

	old := win.SelectObject(memory_device, win.HGDIOBJ(bitmap))
	if old == 0 {
		return nil, errors.New("SelectObject failed")
	}
	defer win.SelectObject(memory_device, old)

	if !win.BitBlt(memory_device, 0, 0, int32(width), int32(height), hdc, int32(x), int32(y), win.SRCCOPY) {
		return nil, errors.New("BitBlt failed")
	}

	if win.GetDIBits(hdc, bitmap, 0, uint32(height), (*uint8)(memptr), (*win.BITMAPINFO)(unsafe.Pointer(&header)), win.DIB_RGB_COLORS) == 0 {
		return nil, errors.New("GetDIBits failed")
	}

	i := 0
	src := uintptr(memptr)
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			v0 := *(*uint8)(unsafe.Pointer(src))
			v1 := *(*uint8)(unsafe.Pointer(src + 1))
			v2 := *(*uint8)(unsafe.Pointer(src + 2))

			// BGRA => RGBA, and set A to 255
			img.Pix[i], img.Pix[i+1], img.Pix[i+2], img.Pix[i+3] = v2, v1, v0, 255

			i += 4
			src += 4
		}
	}

	return img, nil
}

func NumActiveDisplays() int {
	var count int = 0
	enumDisplayMonitors(win.HDC(0), nil, syscall.NewCallback(countupMonitorCallback), uintptr(unsafe.Pointer(&count)))
	return count
}

func GetDisplayBounds(displayIndex int) image.Rectangle {
	var ctx getMonitorBoundsContext
	ctx.Index = displayIndex
	ctx.Count = 0
	enumDisplayMonitors(win.HDC(0), nil, syscall.NewCallback(getMonitorBoundsCallback), uintptr(unsafe.Pointer(&ctx)))
	return image.Rect(
		int(ctx.Rect.Left), int(ctx.Rect.Top),
		int(ctx.Rect.Right), int(ctx.Rect.Bottom))
}

func getDesktopWindow() win.HWND {
	ret, _, _ := syscall.Syscall(funcGetDesktopWindow, 0, 0, 0, 0)
	return win.HWND(ret)
}

func enumDisplayMonitors(hdc win.HDC, lprcClip *win.RECT, lpfnEnum uintptr, dwData uintptr) bool {
	ret, _, _ := syscall.Syscall6(funcEnumDisplayMonitors, 4,
		uintptr(hdc),
		uintptr(unsafe.Pointer(lprcClip)),
		lpfnEnum,
		dwData,
		0,
		0)
	return int(ret) != 0
}

func countupMonitorCallback(hMonitor win.HMONITOR, hdcMonitor win.HDC, lprcMonitor *win.RECT, dwData uintptr) uintptr {
	var count *int
	count = (*int)(unsafe.Pointer(dwData))
	*count = *count + 1
	return uintptr(1)
}

type getMonitorBoundsContext struct {
	Index int
	Rect  win.RECT
	Count int
}

func getMonitorBoundsCallback(hMonitor win.HMONITOR, hdcMonitor win.HDC, lprcMonitor *win.RECT, dwData uintptr) uintptr {
	var ctx *getMonitorBoundsContext
	ctx = (*getMonitorBoundsContext)(unsafe.Pointer(dwData))
	if ctx.Count != ctx.Index {
		ctx.Count = ctx.Count + 1
		return uintptr(1)
	}

	if realSize := getMonitorRealSize(hMonitor); realSize != nil {
		ctx.Rect = *realSize
	} else {
		ctx.Rect = *lprcMonitor
	}

	return uintptr(0)
}

type _MONITORINFOEX struct {
	win.MONITORINFO
	DeviceName [win.CCHDEVICENAME]uint16
}

const _ENUM_CURRENT_SETTINGS = 0xFFFFFFFF

type _DEVMODE struct {
	_            [68]byte
	DmSize       uint16
	_            [6]byte
	DmPosition   win.POINT
	_            [86]byte
	DmPelsWidth  uint32
	DmPelsHeight uint32
	_            [40]byte
}

// getMonitorRealSize makes a call to GetMonitorInfo
// to obtain the device name for the monitor handle
// provided to the method.
//
// With the device name, EnumDisplaySettings is called to
// obtain the current configuration for the monitor, this
// information includes the real resolution of the monitor
// rather than the scaled version based on DPI.
//
// If either handle calls fail, it will return a nil
// allowing the calling method to use the bounds information
// returned by EnumDisplayMonitors which may be affected
// by DPI.
func getMonitorRealSize(hMonitor win.HMONITOR) *win.RECT {
	info := _MONITORINFOEX{}
	info.CbSize = uint32(unsafe.Sizeof(info))

	ret, _, _ := syscall.Syscall(funcGetMonitorInfo, 2, uintptr(hMonitor), uintptr(unsafe.Pointer(&info)), 0)
	if ret == 0 {
		return nil
	}

	devMode := _DEVMODE{}
	devMode.DmSize = uint16(unsafe.Sizeof(devMode))

	if ret, _, _ := syscall.Syscall(funcEnumDisplaySettings, 3, uintptr(unsafe.Pointer(&info.DeviceName[0])), _ENUM_CURRENT_SETTINGS, uintptr(unsafe.Pointer(&devMode))); ret == 0 {
		return nil
	}

	return &win.RECT{
		Left:   devMode.DmPosition.X,
		Right:  devMode.DmPosition.X + int32(devMode.DmPelsWidth),
		Top:    devMode.DmPosition.Y,
		Bottom: devMode.DmPosition.Y + int32(devMode.DmPelsHeight),
	}
}
<|endoftext|>


<commit_before>package mockfilestorage

import (
	"os"
	"path"

	"github.com/Nivl/go-rest-tools/storage/filestorage"
	"github.com/stretchr/testify/mock"
)

var (
	// StringType represent a string argument
	StringType = mock.AnythingOfType("string")

	// ReaderType represents an io.Reader argument
	ReaderType = mock.AnythingOfType("io.Reader")

	// UpdatableAttrType represents a *filestorage.UpdatableFileAttributes argument
	UpdatableAttrType = mock.AnythingOfType("*filestorage.UpdatableFileAttributes")
)

// ExpectWriteIfNotExist is an helper that expects a WriteIfNotExist to
// succeed with the provided params
func (s *FileStorage) ExpectWriteIfNotExist(isNew bool, url string) *mock.Call {
	return s.On("WriteIfNotExist", ReaderType, StringType).Return(isNew, url, nil)
}

// ExpectRead is an helper that expects a Read
func (s *FileStorage) ExpectRead(cwd, filename string) *mock.Call {
	filePath := path.Join(cwd, "fixtures", filename)
	return s.On("Read", StringType).Return(os.Open(filePath))
}

// ExpectExists is an helper that expects Exists() to return true
func (s *FileStorage) ExpectExists() *mock.Call {
	return s.On("Exists", StringType).Return(true, nil)
}

// ExpectNotExists is an helper that expects Exists() to return false
func (s *FileStorage) ExpectNotExists() *mock.Call {
	return s.On("Exists", StringType).Return(false, nil)
}

// ExpectURL is an helper that expects URL() to return given param
func (s *FileStorage) ExpectURL(url string) *mock.Call {
	return s.On("URL", StringType).Return(url, nil)
}

// ExpectSetAttributes is an helper that expects SetAttributes() to work,
// and returns an empty content
func (s *FileStorage) ExpectSetAttributes() *mock.Call {
	attrs := &filestorage.FileAttributes{}
	return s.On("SetAttributes", StringType, UpdatableAttrType).Return(attrs, nil)
}

// ExpectSetAttributesRet is an helper that expects SetAttributes() to
// return the provided object
func (s *FileStorage) ExpectSetAttributesRet(attrs *filestorage.FileAttributes) *mock.Call {
	return s.On("SetAttributes", StringType, UpdatableAttrType).Return(attrs, nil)
}

// ExpectAttributes is an helper that expects Attributes() to
// return the provided object
func (s *FileStorage) ExpectAttributes(attrs *filestorage.FileAttributes) *mock.Call {
	return s.On("Attributes", StringType).Return(attrs, nil)
}

// ExpectDelete is an helper that expects Delete() to succeed
func (s *FileStorage) ExpectDelete() *mock.Call {
	return s.On("Delete", StringType).Return(nil)
}
<commit_msg>fix(tests): fix mock.ExpectWriteIfNotExist to accept 'anything' as reader<commit_after>package mockfilestorage

import (
	"os"
	"path"

	"github.com/Nivl/go-rest-tools/storage/filestorage"
	"github.com/stretchr/testify/mock"
)

var (
	// StringType represent a string argument
	StringType = mock.AnythingOfType("string")

	// UpdatableAttrType represents a *filestorage.UpdatableFileAttributes argument
	UpdatableAttrType = mock.AnythingOfType("*filestorage.UpdatableFileAttributes")
)

// ExpectWriteIfNotExist is an helper that expects a WriteIfNotExist to
// succeed with the provided params
func (s *FileStorage) ExpectWriteIfNotExist(isNew bool, url string) *mock.Call {
	return s.On("WriteIfNotExist", mock.Anything, StringType).Return(isNew, url, nil)
}

// ExpectRead is an helper that expects a Read
func (s *FileStorage) ExpectRead(cwd, filename string) *mock.Call {
	filePath := path.Join(cwd, "fixtures", filename)
	return s.On("Read", StringType).Return(os.Open(filePath))
}

// ExpectExists is an helper that expects Exists() to return true
func (s *FileStorage) ExpectExists() *mock.Call {
	return s.On("Exists", StringType).Return(true, nil)
}

// ExpectNotExists is an helper that expects Exists() to return false
func (s *FileStorage) ExpectNotExists() *mock.Call {
	return s.On("Exists", StringType).Return(false, nil)
}

// ExpectURL is an helper that expects URL() to return given param
func (s *FileStorage) ExpectURL(url string) *mock.Call {
	return s.On("URL", StringType).Return(url, nil)
}

// ExpectSetAttributes is an helper that expects SetAttributes() to work,
// and returns an empty content
func (s *FileStorage) ExpectSetAttributes() *mock.Call {
	attrs := &filestorage.FileAttributes{}
	return s.On("SetAttributes", StringType, UpdatableAttrType).Return(attrs, nil)
}

// ExpectSetAttributesRet is an helper that expects SetAttributes() to
// return the provided object
func (s *FileStorage) ExpectSetAttributesRet(attrs *filestorage.FileAttributes) *mock.Call {
	return s.On("SetAttributes", StringType, UpdatableAttrType).Return(attrs, nil)
}

// ExpectAttributes is an helper that expects Attributes() to
// return the provided object
func (s *FileStorage) ExpectAttributes(attrs *filestorage.FileAttributes) *mock.Call {
	return s.On("Attributes", StringType).Return(attrs, nil)
}

// ExpectDelete is an helper that expects Delete() to succeed
func (s *FileStorage) ExpectDelete() *mock.Call {
	return s.On("Delete", StringType).Return(nil)
}
<|endoftext|>


<commit_before>package analytics

//
// dependencies
//

import "github.com/jehiah/go-strftime"
import "github.com/nu7hatch/gouuid"
import . "encoding/json"
import "net/http"
import "bytes"
import "time"
import "log"

//
// Library version
//

const Version = "0.0.1"

//
// Default API end-point
//

const api = "https://api.segment.io"

//
// Segment.io client
//

type client struct {
	debug      bool
	key        string
	url        string
	flushAt    int
	flushAfter time.Duration
	buffer     []*interface{}
}

//
// Message context library
//

type contextLibrary struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

//
// Message context
//

type context struct {
	Library contextLibrary `json:"library"`
}

//
// Identify message
//

type identify struct {
	Action    string      `json:"action"`
	Traits    interface{} `json:"trailts"`
	Timestamp string      `json:"timestamp"`
}

//
// Alias message
//

type alias struct {
	Action     string `json:"action"`
	PreviousId string `json:"previousId"`
	Timestamp  string `json:"timestamp"`
}

//
// Track message
//

type track struct {
	Action     string      `json:"action"`
	Event      string      `json:"event"`
	Properties interface{} `json:"properties"`
	Timestamp  string      `json:"timestamp"`
}

//
// Group message
//

type group struct {
	Action    string      `json:"action"`
	GroupId   string      `json:"groupId"`
	Traits    interface{} `json:"trailts"`
	Timestamp string      `json:"timestamp"`
}

//
// Page message
//

type page struct {
	Action     string      `json:"action"`
	Category   string      `json:"category"`
	Name       string      `json:"name"`
	Properties interface{} `json:"properties"`
	Timestamp  string      `json:"timestamp"`
}

//
// Batch message
//

type batch struct {
	Context   context        `json:"context"`
	RequestId string         `json:"requestId"`
	Messages  []*interface{} `json:"batch"`
}

//
// Return a new Segment.io client
// with the given write key.
//

func Client(key string) *client {
	c := &client{
		key:    key,
		url:    api,
		buffer: make([]*interface{}, 0),
	}

	c.FlushAt(500)
	c.FlushAfter(10 * time.Second)

	return c
}

//
// Set buffer max.
//

func (c *client) FlushAt(n int) {
	c.flushAt = n
}

//
// Set buffer flush interal.
//

func (c *client) FlushAfter(interval time.Duration) {
	c.flushAfter = interval

	go func() {
		for {
			time.Sleep(interval)
			c.log("interval %v reached", interval)
			c.flush()
		}
	}()
}

//
// Enable debug mode.
//

func (c *client) Debug() {
	c.debug = true
}

//
// Set target url
//

func (c *client) URL(url string) {
	c.url = url
}

//
// Return formatted timestamp.
//

func timestamp() string {
	return strftime.Format("%Y-%m-%dT%H:%M:%S%z", time.Now())
}

// Return a batch message primed
// with context properties
//

func createBatch(msgs []*interface{}) (*batch, error) {
	uid, err := uuid.NewV4()

	if err != nil {
		return nil, err
	}

	batch := &batch{
		RequestId: uid.String(),
		Messages:  msgs,
		Context: context{
			Library: contextLibrary{
				Name:    "analytics-go",
				Version: Version,
			},
		},
	}

	return batch, nil
}

//
// Flush the buffered messages.
//

func (c *client) flush() error {
	if len(c.buffer) == 0 {
		c.log("no messages to flush")
		return nil
	}

	c.log("flushing %d messages", len(c.buffer))
	batch, err := createBatch(c.buffer)

	if err != nil {
		return err
	}

	json, err := Marshal(batch)

	if err != nil {
		return err
	}

	c.buffer = nil

	client := &http.Client{}
	req, err := http.NewRequest("POST", c.url+"/v1/batch", bytes.NewBuffer(json))

	if err != nil {
		return err
	}

	req.Header.Add("User-Agent", "analytics-go (version: "+Version+")")
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Content-Length", string(len(json)))
	req.SetBasicAuth(c.key, "")

	_, err = client.Do(req)

	return err
}

//
// Buffer the given message and flush
// when the buffer exceeds .flushAt.
//

func (c *client) bufferMessage(msg interface{}) error {
	c.buffer = append(c.buffer, &msg)

	c.log("buffer (%d/%d) %v", len(c.buffer), c.flushAt, msg)

	if len(c.buffer) >= c.flushAt {
		return c.flush()
	}

	return nil
}

//
// Log in debug mode.
//

func (c *client) log(format string, v ...interface{}) {
	if c.debug {
		log.Printf(format, v...)
	}
}

//
// Buffer an alias message
//

func (c *client) Alias(previousId string) error {
	return c.bufferMessage(&alias{"Alias", previousId, timestamp()})
}

//
// Buffer a page message
//

func (c *client) Page(name string, category string, properties interface{}) error {
	return c.bufferMessage(&page{"Page", name, category, properties, timestamp()})
}

//
// Buffer a screen message
//

func (c *client) Screen(name string, category string, properties interface{}) error {
	return c.bufferMessage(&page{"Screen", name, category, properties, timestamp()})
}

//
// Buffer a group message
//

func (c *client) Group(id string, traits interface{}) error {
	return c.bufferMessage(&group{"Group", id, traits, timestamp()})
}

//
// Buffer an identify message
//

func (c *client) Identify(traits interface{}) error {
	return c.bufferMessage(&identify{"Identify", traits, timestamp()})
}

//
// Buffer a track message
//

func (c *client) Track(event string, properties interface{}) error {
	return c.bufferMessage(&track{"Track", event, properties, timestamp()})
}
<commit_msg>fix multiple .flushAfter() calls<commit_after>package analytics

//
// dependencies
//

import "github.com/jehiah/go-strftime"
import "github.com/nu7hatch/gouuid"
import . "encoding/json"
import "net/http"
import "bytes"
import "time"
import "log"

//
// Library version
//

const Version = "0.0.1"

//
// Default API end-point
//

const api = "https://api.segment.io"

//
// Segment.io client
//

type client struct {
	debug         bool
	running       bool
	key           string
	url           string
	flushCount    int
	flushInterval time.Duration
	buffer        []*interface{}
}

//
// Message context library
//

type contextLibrary struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

//
// Message context
//

type context struct {
	Library contextLibrary `json:"library"`
}

//
// Identify message
//

type identify struct {
	Action    string      `json:"action"`
	Traits    interface{} `json:"trailts"`
	Timestamp string      `json:"timestamp"`
}

//
// Alias message
//

type alias struct {
	Action     string `json:"action"`
	PreviousId string `json:"previousId"`
	Timestamp  string `json:"timestamp"`
}

//
// Track message
//

type track struct {
	Action     string      `json:"action"`
	Event      string      `json:"event"`
	Properties interface{} `json:"properties"`
	Timestamp  string      `json:"timestamp"`
}

//
// Group message
//

type group struct {
	Action    string      `json:"action"`
	GroupId   string      `json:"groupId"`
	Traits    interface{} `json:"trailts"`
	Timestamp string      `json:"timestamp"`
}

//
// Page message
//

type page struct {
	Action     string      `json:"action"`
	Category   string      `json:"category"`
	Name       string      `json:"name"`
	Properties interface{} `json:"properties"`
	Timestamp  string      `json:"timestamp"`
}

//
// Batch message
//

type batch struct {
	Context   context        `json:"context"`
	RequestId string         `json:"requestId"`
	Messages  []*interface{} `json:"batch"`
}

//
// Return a new Segment.io client
// with the given write key.
//

func Client(key string) *client {
	c := &client{
		key:    key,
		url:    api,
		buffer: make([]*interface{}, 0),
	}

	c.FlushAt(500)
	c.FlushAfter(10 * time.Second)

	return c
}

//
// Set buffer max.
//

func (c *client) FlushAt(n int) {
	c.flushCount = n
}

//
// Set buffer flush interal.
//

func (c *client) FlushAfter(interval time.Duration) {
	c.flushInterval = interval

	if c.running {
		return
	}

	c.running = true

	go func() {
		for {
			time.Sleep(c.flushInterval)
			c.log("interval %v reached", c.flushInterval)
			c.flush()
		}
	}()
}

//
// Enable debug mode.
//

func (c *client) Debug() {
	c.debug = true
}

//
// Set target url
//

func (c *client) URL(url string) {
	c.url = url
}

//
// Return formatted timestamp.
//

func timestamp() string {
	return strftime.Format("%Y-%m-%dT%H:%M:%S%z", time.Now())
}

// Return a batch message primed
// with context properties
//

func createBatch(msgs []*interface{}) (*batch, error) {
	uid, err := uuid.NewV4()

	if err != nil {
		return nil, err
	}

	batch := &batch{
		RequestId: uid.String(),
		Messages:  msgs,
		Context: context{
			Library: contextLibrary{
				Name:    "analytics-go",
				Version: Version,
			},
		},
	}

	return batch, nil
}

//
// Flush the buffered messages.
//

func (c *client) flush() error {
	if len(c.buffer) == 0 {
		c.log("no messages to flush")
		return nil
	}

	c.log("flushing %d messages", len(c.buffer))
	batch, err := createBatch(c.buffer)

	if err != nil {
		return err
	}

	json, err := Marshal(batch)

	if err != nil {
		return err
	}

	c.buffer = nil

	client := &http.Client{}
	req, err := http.NewRequest("POST", c.url+"/v1/batch", bytes.NewBuffer(json))

	if err != nil {
		return err
	}

	req.Header.Add("User-Agent", "analytics-go (version: "+Version+")")
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Content-Length", string(len(json)))
	req.SetBasicAuth(c.key, "")

	_, err = client.Do(req)

	return err
}

//
// Buffer the given message and flush
// when the buffer exceeds .flushCount.
//

func (c *client) bufferMessage(msg interface{}) error {
	c.buffer = append(c.buffer, &msg)

	c.log("buffer (%d/%d) %v", len(c.buffer), c.flushCount, msg)

	if len(c.buffer) >= c.flushCount {
		return c.flush()
	}

	return nil
}

//
// Log in debug mode.
//

func (c *client) log(format string, v ...interface{}) {
	if c.debug {
		log.Printf(format, v...)
	}
}

//
// Buffer an alias message
//

func (c *client) Alias(previousId string) error {
	return c.bufferMessage(&alias{"Alias", previousId, timestamp()})
}

//
// Buffer a page message
//

func (c *client) Page(name string, category string, properties interface{}) error {
	return c.bufferMessage(&page{"Page", name, category, properties, timestamp()})
}

//
// Buffer a screen message
//

func (c *client) Screen(name string, category string, properties interface{}) error {
	return c.bufferMessage(&page{"Screen", name, category, properties, timestamp()})
}

//
// Buffer a group message
//

func (c *client) Group(id string, traits interface{}) error {
	return c.bufferMessage(&group{"Group", id, traits, timestamp()})
}

//
// Buffer an identify message
//

func (c *client) Identify(traits interface{}) error {
	return c.bufferMessage(&identify{"Identify", traits, timestamp()})
}

//
// Buffer a track message
//

func (c *client) Track(event string, properties interface{}) error {
	return c.bufferMessage(&track{"Track", event, properties, timestamp()})
}
<|endoftext|>