<commit_before>package scipipe

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
)

// Task represents a single static shell command, or go function, to be
// executed, and are scheduled and managed by a corresponding Process
type Task struct {
	Name          string
	Command       string
	CustomExecute func(*Task)
	InIPs         map[string]*FileIP
	OutIPs        map[string]*FileIP
	Params        map[string]string
	Tags          map[string]string
	Done          chan int
	cores         int
	workflow      *Workflow
	process       *Process
}

// ------------------------------------------------------------------------
// Factory method(s)
// ------------------------------------------------------------------------

// NewTask instantiates and initializes a new Task
func NewTask(workflow *Workflow, process *Process, name string, cmdPat string, inIPs map[string]*FileIP, outPathFuncs map[string]func(*Task) string, outPortsDoStream map[string]bool, params map[string]string, tags map[string]string, prepend string, customExecute func(*Task), cores int) *Task {
	t := &Task{
		Name:          name,
		InIPs:         inIPs,
		OutIPs:        make(map[string]*FileIP),
		Params:        params,
		Tags:          tags,
		Command:       "",
		CustomExecute: customExecute,
		Done:          make(chan int),
		cores:         cores,
		workflow:      workflow,
		process:       process,
	}

	// Create Out-IPs
	for oname, outPathFunc := range outPathFuncs {
		oip := NewFileIP(outPathFunc(t))
		if outPortsDoStream[oname] {
			oip.doStream = true
		}
		t.OutIPs[oname] = oip
	}
	t.Command = formatCommand(cmdPat, inIPs, t.OutIPs, params, tags, prepend)
	return t
}

// formatCommand is a helper function for NewTask, that formats a shell command
// based on concrete file paths and parameter values
func formatCommand(cmd string, inIPs map[string]*FileIP, outIPs map[string]*FileIP, params map[string]string, tags map[string]string, prepend string) string {
	r := getShellCommandPlaceHolderRegex()
	placeHolderMatches := r.FindAllStringSubmatch(cmd, -1)
	for _, placeHolderMatch := range placeHolderMatches {
		var filePath string

		placeHolderStr := placeHolderMatch[0]
		portType := placeHolderMatch[1]
		portName := placeHolderMatch[2]
		sep := " " // Default

		switch portType {
		case "o":
			if outIPs[portName] == nil {
				Fail("Missing outpath for outport '", portName, "' for command '", cmd, "'")
			}
			filePath = outIPs[portName].TempPath()
		case "os":
			if outIPs[portName] == nil {
				Fail("Missing outpath for outport '", portName, "' for command '", cmd, "'")
			}
			filePath = outIPs[portName].FifoPath()
		case "i":
			if inIPs[portName] == nil {
				Fail("Missing in-IP for inport '", portName, "' for command '", cmd, "'")
			}
			// Identify if the place holder represents a reduce-type in-port
			reduceInputs := false
			if len(placeHolderMatch) > 5 {
				sep = placeHolderMatch[7]
				reduceInputs = true
			}
			if reduceInputs && inIPs[portName].Path() == "" {
				// Merge multiple input paths from a substream on the IP, into one string
				ips := []*FileIP{}
				for ip := range inIPs[portName].SubStream.Chan {
					ips = append(ips, ip)
				}
				paths := []string{}
				for _, ip := range ips {
					paths = append(paths, ip.Path())
				}
				filePath = strings.Join(paths, sep)
			} else {
				if inIPs[portName].Path() == "" {
					Fail("Missing inpath for inport '", portName, "', and no substream, for command '", cmd, "'")
				}
				if inIPs[portName].doStream {
					filePath = inIPs[portName].FifoPath()
				} else {
					filePath = inIPs[portName].Path()
				}
			}
		case "p":
			if params[portName] == "" {
				msg := fmt.Sprint("Missing param value for param '", portName, "' for command '", cmd, "'")
				Fail(msg)
			} else {
				filePath = params[portName]
			}
		case "t":
			if tags[portName] == "" {
				msg := fmt.Sprint("Missing tag value for tag '", portName, "' for command '", cmd, "'")
				Fail(msg)
			} else {
				filePath = tags[portName]
			}
		default:
			Fail("Replace failed for port ", portName, " for command '", cmd, "'")
		}
		cmd = strings.Replace(cmd, placeHolderStr, filePath, -1)
	}

	// Add prepend string to the command
	if prepend != "" {
		cmd = fmt.Sprintf("%s %s", prepend, cmd)
	}

	return cmd
}

// ------------------------------------------------------------------------
// Main API methods: Accessor methods
// ------------------------------------------------------------------------

// InIP returns an IP for the in-port with name portName
func (t *Task) InIP(portName string) *FileIP {
	if t.InIPs[portName] == nil {
		Failf("No such in-portname (%s) in task (%s)\n", portName, t.Name)
	}
	return t.InIPs[portName]
}

// InPath returns the path name of an input file for the task
func (t *Task) InPath(portName string) string {
	return t.InIP(portName).Path()
}

// OutIP returns an IP for the in-port with name portName
func (t *Task) OutIP(portName string) *FileIP {
	if t.OutIPs[portName] == nil {
		Failf("No such out-portname (%s) in task (%s)\n", portName, t.Name)
	}
	return t.OutIPs[portName]
}

// OutPath returns the path name of an input file for the task
func (t *Task) OutPath(portName string) string {
	return t.OutIP(portName).Path()
}

// Param returns the value of a param, for the task
func (t *Task) Param(portName string) string {
	if param, ok := t.Params[portName]; ok {
		return param
	}
	Failf("No such param port '%s' for task '%s'\n", portName, t.Name)
	return "invalid"
}

// Tag returns the value of a param, for the task
func (t *Task) Tag(tagName string) string {
	if tag, ok := t.Tags[tagName]; ok {
		return tag
	}
	Failf("No such tag '%s' for task '%s'\n", tagName, t.Name)
	return "invalid"
}

// ------------------------------------------------------------------------
// Execute the task
// ------------------------------------------------------------------------

// Execute executes the task (the shell command or go function in CustomExecute)
func (t *Task) Execute() {
	defer close(t.Done)

	// Do some sanity checks
	if t.anyTempfileExists() {
		Failf("| %-32s | Existing temp files found so existing. Clean up .tmp files before restarting the workflow!", t.Name)
	}

	if t.anyOutputsExist() {
		t.Done <- 1
		return
	}

	// Execute task
	t.workflow.IncConcurrentTasks(t.cores) // Will block if max concurrent tasks is reached
	t.createDirs()                         // Create output directories needed for any outputs
	startTime := time.Now()
	if t.CustomExecute != nil {
		outputsStr := ""
		for oipName, oip := range t.OutIPs {
			outputsStr += " " + oipName + ": " + oip.Path()
		}
		Audit.Printf("| %-32s | Executing: Custom Go function with outputs: %s\n", t.Name, outputsStr)
		t.CustomExecute(t)
		Audit.Printf("| %-32s | Finished:  Custom Go function with outputs: %s\n", t.Name, outputsStr)
	} else {
		Audit.Printf("| %-32s | Executing: %s\n", t.Name, t.Command)
		t.executeCommand(t.Command)
		Audit.Printf("| %-32s | Finished:  %s\n", t.Name, t.Command)
	}
	finishTime := time.Now()
	t.writeAuditLogs(startTime, finishTime)
	t.atomizeIPs()
	t.workflow.DecConcurrentTasks(t.cores)

	t.Done <- 1
}

// ------------------------------------------------------------------------
// Helper methods for the Execute method
// ------------------------------------------------------------------------

// anyTempFileExists checks if any temporary workflow files exist and if so, returns true
func (t *Task) anyTempfileExists() (anyTempfileExists bool) {
	anyTempfileExists = false
	for _, oip := range t.OutIPs {
		if !oip.doStream {
			otmpPath := oip.TempPath()
			if _, err := os.Stat(otmpPath); err == nil {
				Warning.Printf("| %-32s | Temp file already exists: %s (Note: If resuming from a failed run, clean up .tmp files first. Also, make sure that two processes don't produce the same output files!).\n", t.Name, otmpPath)
				anyTempfileExists = true
			}
		}
	}
	return
}

// anyOutputsExist if any output file IP, or temporary file IPs, exist
func (t *Task) anyOutputsExist() (anyFileExists bool) {
	anyFileExists = false
	for _, oip := range t.OutIPs {
		if !oip.doStream {
			opath := oip.Path()
			if _, err := os.Stat(opath); err == nil {
				Audit.Printf("| %-32s | Output file already exists, so skipping: %s\n", t.Name, opath)
				anyFileExists = true
			}
		}
	}
	return
}

// createDirs creates directories for out-IPs of the task
func (t *Task) createDirs() {
	for _, oip := range t.OutIPs {
		oipDir := oip.TempDir() // This will create all out dirs, including the temp dir
		err := os.MkdirAll(oipDir, 0777)
		CheckWithMsg(err, "Could not create directory: "+oipDir)
	}

}

// executeCommand executes the shell command cmd via bash
func (t *Task) executeCommand(cmd string) {
	out, err := exec.Command("bash", "-c", cmd).CombinedOutput()
	if err != nil {
		Failf("Command failed!\nCommand:\n%s\n\nOutput:\n%s\nOriginal error:%s\n", cmd, string(out), err.Error())
	}
}

func (t *Task) writeAuditLogs(startTime time.Time, finishTime time.Time) {
	// Append audit info for the task to all its output IPs
	auditInfo := NewAuditInfo()
	auditInfo.Command = t.Command
	auditInfo.ProcessName = t.process.Name()
	auditInfo.Params = t.Params
	auditInfo.StartTime = startTime
	auditInfo.FinishTime = finishTime
	auditInfo.ExecTimeMS = finishTime.Sub(startTime) / time.Millisecond
	// Set the audit infos from incoming IPs into the "Upstream" map
	for _, iip := range t.InIPs {
		iipPath := iip.Path()
		iipAuditInfo := iip.AuditInfo()
		auditInfo.Upstream[iipPath] = iipAuditInfo
	}
	// Add the current audit info to output ips and write them to file
	for _, oip := range t.OutIPs {
		oip.SetAuditInfo(auditInfo)
		for _, iip := range t.InIPs {
			oip.AddTags(iip.Tags())
		}
		oip.WriteAuditLogToFile()
	}
}

// Rename temporary output files to their proper file names
func (t *Task) atomizeIPs() {
	for _, oip := range t.OutIPs {
		if !oip.doStream {
			oip.Atomize()
		}
	}
}
<commit_msg>Don't create temp dir for FIFO files<commit_after>package scipipe

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

// Task represents a single static shell command, or go function, to be
// executed, and are scheduled and managed by a corresponding Process
type Task struct {
	Name          string
	Command       string
	CustomExecute func(*Task)
	InIPs         map[string]*FileIP
	OutIPs        map[string]*FileIP
	Params        map[string]string
	Tags          map[string]string
	Done          chan int
	cores         int
	workflow      *Workflow
	process       *Process
}

// ------------------------------------------------------------------------
// Factory method(s)
// ------------------------------------------------------------------------

// NewTask instantiates and initializes a new Task
func NewTask(workflow *Workflow, process *Process, name string, cmdPat string, inIPs map[string]*FileIP, outPathFuncs map[string]func(*Task) string, outPortsDoStream map[string]bool, params map[string]string, tags map[string]string, prepend string, customExecute func(*Task), cores int) *Task {
	t := &Task{
		Name:          name,
		InIPs:         inIPs,
		OutIPs:        make(map[string]*FileIP),
		Params:        params,
		Tags:          tags,
		Command:       "",
		CustomExecute: customExecute,
		Done:          make(chan int),
		cores:         cores,
		workflow:      workflow,
		process:       process,
	}

	// Create Out-IPs
	for oname, outPathFunc := range outPathFuncs {
		oip := NewFileIP(outPathFunc(t))
		if outPortsDoStream[oname] {
			oip.doStream = true
		}
		t.OutIPs[oname] = oip
	}
	t.Command = formatCommand(cmdPat, inIPs, t.OutIPs, params, tags, prepend)
	return t
}

// formatCommand is a helper function for NewTask, that formats a shell command
// based on concrete file paths and parameter values
func formatCommand(cmd string, inIPs map[string]*FileIP, outIPs map[string]*FileIP, params map[string]string, tags map[string]string, prepend string) string {
	r := getShellCommandPlaceHolderRegex()
	placeHolderMatches := r.FindAllStringSubmatch(cmd, -1)
	for _, placeHolderMatch := range placeHolderMatches {
		var filePath string

		placeHolderStr := placeHolderMatch[0]
		portType := placeHolderMatch[1]
		portName := placeHolderMatch[2]
		sep := " " // Default

		switch portType {
		case "o":
			if outIPs[portName] == nil {
				Fail("Missing outpath for outport '", portName, "' for command '", cmd, "'")
			}
			filePath = outIPs[portName].TempPath()
		case "os":
			if outIPs[portName] == nil {
				Fail("Missing outpath for outport '", portName, "' for command '", cmd, "'")
			}
			filePath = outIPs[portName].FifoPath()
		case "i":
			if inIPs[portName] == nil {
				Fail("Missing in-IP for inport '", portName, "' for command '", cmd, "'")
			}
			// Identify if the place holder represents a reduce-type in-port
			reduceInputs := false
			if len(placeHolderMatch) > 5 {
				sep = placeHolderMatch[7]
				reduceInputs = true
			}
			if reduceInputs && inIPs[portName].Path() == "" {
				// Merge multiple input paths from a substream on the IP, into one string
				ips := []*FileIP{}
				for ip := range inIPs[portName].SubStream.Chan {
					ips = append(ips, ip)
				}
				paths := []string{}
				for _, ip := range ips {
					paths = append(paths, ip.Path())
				}
				filePath = strings.Join(paths, sep)
			} else {
				if inIPs[portName].Path() == "" {
					Fail("Missing inpath for inport '", portName, "', and no substream, for command '", cmd, "'")
				}
				if inIPs[portName].doStream {
					filePath = inIPs[portName].FifoPath()
				} else {
					filePath = inIPs[portName].Path()
				}
			}
		case "p":
			if params[portName] == "" {
				msg := fmt.Sprint("Missing param value for param '", portName, "' for command '", cmd, "'")
				Fail(msg)
			} else {
				filePath = params[portName]
			}
		case "t":
			if tags[portName] == "" {
				msg := fmt.Sprint("Missing tag value for tag '", portName, "' for command '", cmd, "'")
				Fail(msg)
			} else {
				filePath = tags[portName]
			}
		default:
			Fail("Replace failed for port ", portName, " for command '", cmd, "'")
		}
		cmd = strings.Replace(cmd, placeHolderStr, filePath, -1)
	}

	// Add prepend string to the command
	if prepend != "" {
		cmd = fmt.Sprintf("%s %s", prepend, cmd)
	}

	return cmd
}

// ------------------------------------------------------------------------
// Main API methods: Accessor methods
// ------------------------------------------------------------------------

// InIP returns an IP for the in-port with name portName
func (t *Task) InIP(portName string) *FileIP {
	if t.InIPs[portName] == nil {
		Failf("No such in-portname (%s) in task (%s)\n", portName, t.Name)
	}
	return t.InIPs[portName]
}

// InPath returns the path name of an input file for the task
func (t *Task) InPath(portName string) string {
	return t.InIP(portName).Path()
}

// OutIP returns an IP for the in-port with name portName
func (t *Task) OutIP(portName string) *FileIP {
	if t.OutIPs[portName] == nil {
		Failf("No such out-portname (%s) in task (%s)\n", portName, t.Name)
	}
	return t.OutIPs[portName]
}

// OutPath returns the path name of an input file for the task
func (t *Task) OutPath(portName string) string {
	return t.OutIP(portName).Path()
}

// Param returns the value of a param, for the task
func (t *Task) Param(portName string) string {
	if param, ok := t.Params[portName]; ok {
		return param
	}
	Failf("No such param port '%s' for task '%s'\n", portName, t.Name)
	return "invalid"
}

// Tag returns the value of a param, for the task
func (t *Task) Tag(tagName string) string {
	if tag, ok := t.Tags[tagName]; ok {
		return tag
	}
	Failf("No such tag '%s' for task '%s'\n", tagName, t.Name)
	return "invalid"
}

// ------------------------------------------------------------------------
// Execute the task
// ------------------------------------------------------------------------

// Execute executes the task (the shell command or go function in CustomExecute)
func (t *Task) Execute() {
	defer close(t.Done)

	// Do some sanity checks
	if t.anyTempfileExists() {
		Failf("| %-32s | Existing temp files found so existing. Clean up .tmp files before restarting the workflow!", t.Name)
	}

	if t.anyOutputsExist() {
		t.Done <- 1
		return
	}

	// Execute task
	t.workflow.IncConcurrentTasks(t.cores) // Will block if max concurrent tasks is reached
	t.createDirs()                         // Create output directories needed for any outputs
	startTime := time.Now()
	if t.CustomExecute != nil {
		outputsStr := ""
		for oipName, oip := range t.OutIPs {
			outputsStr += " " + oipName + ": " + oip.Path()
		}
		Audit.Printf("| %-32s | Executing: Custom Go function with outputs: %s\n", t.Name, outputsStr)
		t.CustomExecute(t)
		Audit.Printf("| %-32s | Finished:  Custom Go function with outputs: %s\n", t.Name, outputsStr)
	} else {
		Audit.Printf("| %-32s | Executing: %s\n", t.Name, t.Command)
		t.executeCommand(t.Command)
		Audit.Printf("| %-32s | Finished:  %s\n", t.Name, t.Command)
	}
	finishTime := time.Now()
	t.writeAuditLogs(startTime, finishTime)
	t.atomizeIPs()
	t.workflow.DecConcurrentTasks(t.cores)

	t.Done <- 1
}

// ------------------------------------------------------------------------
// Helper methods for the Execute method
// ------------------------------------------------------------------------

// anyTempFileExists checks if any temporary workflow files exist and if so, returns true
func (t *Task) anyTempfileExists() (anyTempfileExists bool) {
	anyTempfileExists = false
	for _, oip := range t.OutIPs {
		if !oip.doStream {
			otmpPath := oip.TempPath()
			if _, err := os.Stat(otmpPath); err == nil {
				Warning.Printf("| %-32s | Temp file already exists: %s (Note: If resuming from a failed run, clean up .tmp files first. Also, make sure that two processes don't produce the same output files!).\n", t.Name, otmpPath)
				anyTempfileExists = true
			}
		}
	}
	return
}

// anyOutputsExist if any output file IP, or temporary file IPs, exist
func (t *Task) anyOutputsExist() (anyFileExists bool) {
	anyFileExists = false
	for _, oip := range t.OutIPs {
		if !oip.doStream {
			opath := oip.Path()
			if _, err := os.Stat(opath); err == nil {
				Audit.Printf("| %-32s | Output file already exists, so skipping: %s\n", t.Name, opath)
				anyFileExists = true
			}
		}
	}
	return
}

// createDirs creates directories for out-IPs of the task
func (t *Task) createDirs() {
	for _, oip := range t.OutIPs {
		oipDir := oip.TempDir() // This will create all out dirs, including the temp dir
		if oip.doStream {       // Temp dirs are not created for fifo files
			oipDir = filepath.Dir(oip.FifoPath())
		}
		err := os.MkdirAll(oipDir, 0777)
		CheckWithMsg(err, "Could not create directory: "+oipDir)
	}

}

// executeCommand executes the shell command cmd via bash
func (t *Task) executeCommand(cmd string) {
	out, err := exec.Command("bash", "-c", cmd).CombinedOutput()
	if err != nil {
		Failf("Command failed!\nCommand:\n%s\n\nOutput:\n%s\nOriginal error:%s\n", cmd, string(out), err.Error())
	}
}

func (t *Task) writeAuditLogs(startTime time.Time, finishTime time.Time) {
	// Append audit info for the task to all its output IPs
	auditInfo := NewAuditInfo()
	auditInfo.Command = t.Command
	auditInfo.ProcessName = t.process.Name()
	auditInfo.Params = t.Params
	auditInfo.StartTime = startTime
	auditInfo.FinishTime = finishTime
	auditInfo.ExecTimeMS = finishTime.Sub(startTime) / time.Millisecond
	// Set the audit infos from incoming IPs into the "Upstream" map
	for _, iip := range t.InIPs {
		iipPath := iip.Path()
		iipAuditInfo := iip.AuditInfo()
		auditInfo.Upstream[iipPath] = iipAuditInfo
	}
	// Add the current audit info to output ips and write them to file
	for _, oip := range t.OutIPs {
		oip.SetAuditInfo(auditInfo)
		for _, iip := range t.InIPs {
			oip.AddTags(iip.Tags())
		}
		oip.WriteAuditLogToFile()
	}
}

// Rename temporary output files to their proper file names
func (t *Task) atomizeIPs() {
	for _, oip := range t.OutIPs {
		if !oip.doStream {
			oip.Atomize()
		}
	}
}
<|endoftext|>


<commit_before>package test

import (
	"errors"
	"fmt"
	"runtime"
	"strings"
	"testing"
)

// ErrPlain is the default error that is returned for functions in this package.
var ErrPlain = errors.New("error")

func fileline(i int) string {
	_, file, line, ok := runtime.Caller(i)
	if !ok {
		return ""
	}
	parts := strings.Split(file, "/")
	file = parts[len(parts)-1]
	return fmt.Sprintf("%s:%d", file, line)
}

func trace() string {
	trace2 := fileline(2)
	trace3 := fileline(3)
	return "\r\t" + strings.Repeat(" ", len(fmt.Sprintf("%s:", trace2))) + "\r\t" + trace3
}

func printable(s string) string {
	s = strings.Replace(s, "\n", `\n`, -1)
	s = strings.Replace(s, "\r", `\r`, -1)
	s = strings.Replace(s, "\t", `\t`, -1)
	return s
}

func Assert(t *testing.T, condition bool, msgs ...string) {
	if !condition {
		t.Errorf("%s: %s\n", trace(), strings.Join(msgs, " "))
	}
}

func Error(t *testing.T, err, expected error, msgs ...string) {
	if err != expected {
		t.Errorf("%s: %s\n   error: %v\nexpected: %v\n", trace(), strings.Join(msgs, " "), err, expected)
	}
}

func String(t *testing.T, output, expected string, msgs ...string) {
	if output != expected {
		t.Errorf("%s: %s\nminified: %s\nexpected: %s\n", trace(), strings.Join(msgs, " "), printable(output), printable(expected))
	}
}

func Minify(t *testing.T, input string, err error, output, expected string, msgs ...string) {
	if err != nil {
		t.Errorf("%s: %s\n   given: %s\n   error: %v\n", trace(), strings.Join(msgs, " "), printable(input), err)
	}
	if output != expected {
		t.Errorf("%s: %s\n   given: %s\nminified: %s\nexpected: %s\n", trace(), strings.Join(msgs, " "), printable(input), printable(output), printable(expected))
	}
}
<commit_msg>Add That<commit_after>package test

import (
	"bytes"
	"errors"
	"fmt"
	"runtime"
	"strings"
	"testing"
)

// ErrPlain is the default error that is returned for functions in this package.
var ErrPlain = errors.New("error")

func fileline(i int) string {
	_, file, line, ok := runtime.Caller(i)
	if !ok {
		return ""
	}
	parts := strings.Split(file, "/")
	file = parts[len(parts)-1]
	return fmt.Sprintf("%s:%d", file, line)
}

func trace() string {
	trace2 := fileline(2)
	trace3 := fileline(3)
	return "\r\t" + strings.Repeat(" ", len(fmt.Sprintf("%s:", trace2))) + "\r\t" + trace3
}

func printable(s string) string {
	s = strings.Replace(s, "\n", `\n`, -1)
	s = strings.Replace(s, "\r", `\r`, -1)
	s = strings.Replace(s, "\t", `\t`, -1)
	return s
}

func That(t *testing.T, condition bool, msg ...interface{}) {
	if !condition {
		t.Errorf("%s: %s\n", trace(), fmt.Sprint(msg...))
	}
}

func Error(t *testing.T, err, expected error, msg ...interface{}) {
	if err != expected {
		t.Errorf("%s: %s\n   error: %v\nexpected: %v\n", trace(), fmt.Sprint(msg...), err, expected)
	}
}

func String(t *testing.T, output, expected string, msg ...interface{}) {
	if output != expected {
		t.Errorf("%s: %s\n  output: %s\nexpected: %s\n", trace(), fmt.Sprint(msg...), printable(output), printable(expected))
	}
}

func Bytes(t *testing.T, output, expected []byte, msg ...interface{}) {
	if !bytes.Equal(output, expected) {
		t.Errorf("%s: %s\n  output: %s\nexpected: %s\n", trace(), fmt.Sprint(msg...), printable(string(output)), printable(string(expected)))
	}
}

func Minify(t *testing.T, input string, err error, output, expected string, msg ...interface{}) {
	if err != nil {
		t.Errorf("%s: %s\n   given: %s\n   error: %v\n", trace(), fmt.Sprint(msg...), printable(input), err)
	}
	if output != expected {
		t.Errorf("%s: %s\n   given: %s\nminified: %s\nexpected: %s\n", trace(), fmt.Sprint(msg...), printable(input), printable(output), printable(expected))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"strings"
	"sync"

	"github.com/godbus/dbus"
)

const (
	dest = `org.kde.kdeconnect`
	path = `/modules/kdeconnect`

	signalReachableStatusChanged = dest + `.device.reachableStatusChanged`
	signalStateChanged           = dest + `.device.stateChanged`
	signalTrustedChanged         = dest + `.device.trustedChanged`
	signalNameChanged            = dest + `.device.nameChanged`
	signalPluginsChanged         = dest + `.device.pluginsChanged`

	pluginShare = `kdeconnect_share`
)

type deviceList struct {
	devices map[string]*Device
	conn    *dbus.Conn
	sync.RWMutex
}

func (d *deviceList) get(id string) (*Device, bool) {
	d.RLock()
	defer d.RUnlock()
	dev, ok := d.devices[id]
	return dev, ok
}

func (d *deviceList) add(id string) error {
	if _, ok := d.get(id); ok {
		return fmt.Errorf("Device already exists: %s", id)
	}

	d.Lock()
	defer d.Unlock()
	dev, err := newDevice(id)
	if err != nil {
		return err
	}
	d.devices[id] = dev

	return nil
}

func (d *deviceList) delete(id string) {
	d.Lock()
	delete(d.devices, id)
	d.Unlock()
}

func (d *deviceList) all() map[string]*Device {
	d.RLock()
	defer d.RUnlock()
	return d.devices
}

func (d *deviceList) Close() error {
	err := d.conn.Close()
	for _, d := range d.devices {
		if e := d.Close(); err != nil {
			log(e)
			if err == nil {
				err = e
			}
		}
	}
	return err
}

// Device maps to the DBUS device interface
type Device struct {
	ID               string              `json:"id"`
	Type             string              `json:"type"`
	Name             string              `json:"name"`
	IconName         string              `json:"iconName"`
	StatusIconName   string              `json:"statusIconName"`
	IsReachable      bool                `json:"isReachable"`
	IsTrusted        bool                `json:"isTrusted"`
	SupportedPlugins map[string]struct{} `json:"supportedPlugins"`
	conn             *dbus.Conn
	obj              dbus.BusObject
	signal           chan *dbus.Signal
	sync.RWMutex
}

/*
func init() {
	conn, err := dbus.SessionBus()
	if err != nil {
		panic(err)
	}
	// Debugging
	if err = introSpect(conn.Object(dest, dbus.ObjectPath(fmt.Sprintf("%s/devices", path)))); err != nil {
		log(err)
	}
}
*/

func (d *Device) watch() error {
	if err := d.addMatchSignal(`reachableStatusChanged`); err != nil {
		return err
	}
	if err := d.addMatchSignal(`stateChanged`); err != nil {
		return err
	}
	if err := d.addMatchSignal(`trustedChanged`); err != nil {
		return err
	}
	if err := d.addMatchSignal(`nameChanged`); err != nil {
		return err
	}
	if err := d.addMatchSignal(`pluginsChanged`); err != nil {
		return err
	}

	d.conn.Signal(d.signal)
	go func() {
		for s := range d.signal {
			var err error
			switch s.Name {
			case signalReachableStatusChanged:
				if err = d.getIsReachable(); err != nil {
					log(err)
				}
			case signalPluginsChanged:
				if err = d.getSupportedPlugins(); err != nil {
					log(err)
				}
			case signalNameChanged:
				if err = d.getName(); err != nil {
					log(err)
				}
			case signalTrustedChanged:
				if err = d.getIsTrusted(); err != nil {
					log(err)
				}
			case signalStateChanged:
				if err = d.update(); err != nil {
					log(err)
				}
			default:
				if err = d.update(); err != nil {
					log(err)
				}
			}
			update := &message{
				Type: typeDeviceUpdate,
			}
			if update.Data, err = json.Marshal(d); err != nil {
				log(err)
				continue
			}
			messageQueue <- update
		}
	}()

	return nil
}

func (d *Device) addMatchSignal(member string) error {
	call := d.conn.BusObject().Call(
		`org.freedesktop.DBus.AddMatch`,
		0,
		fmt.Sprintf("type='signal',path='%s',interface='%s.device',member='%s'", d.obj.Path(), dest, member),
	)
	return call.Err
}

func (d *Device) getType() error {
	v, err := d.obj.GetProperty(dest + `.device.type`)
	if err != nil {
		return err
	}
	d.Type = strings.Trim(v.String(), `"`)

	return nil
}

func (d *Device) getName() error {
	v, err := d.obj.GetProperty(dest + `.device.name`)
	if err != nil {
		return err
	}
	d.Name = strings.Trim(v.String(), `"`)

	return nil
}

func (d *Device) getIconName() error {
	v, err := d.obj.GetProperty(dest + `.device.iconName`)
	if err != nil {
		return err
	}
	d.IconName = strings.Trim(v.String(), `"`)

	return nil
}

func (d *Device) getStatusIconName() error {
	v, err := d.obj.GetProperty(dest + `.device.statusIconName`)
	if err != nil {
		return err
	}
	d.StatusIconName = strings.Trim(v.String(), `"`)

	return nil
}

func (d *Device) getIsReachable() error {
	v, err := d.obj.GetProperty(dest + `.device.isReachable`)
	if err != nil {
		return err
	}
	d.IsReachable = v.Value().(bool)

	return nil
}

func (d *Device) getIsTrusted() error {
	v, err := d.obj.GetProperty(dest + `.device.isTrusted`)
	if err != nil {
		return err
	}
	d.IsTrusted = v.Value().(bool)

	return nil
}

func (d *Device) getSupportedPlugins() error {
	v, err := d.obj.GetProperty(dest + `.device.supportedPlugins`)
	if err != nil {
		return err
	}
	plugins := make(map[string]struct{})
	for _, plugin := range v.Value().([]string) {
		plugins[plugin] = struct{}{}
	}
	d.Lock()
	d.SupportedPlugins = plugins
	d.Unlock()

	return nil
}

func (d *Device) update() error {
	if err := d.getName(); err != nil {
		logBadProp(d.ID, `name`, err)
	}
	if err := d.getType(); err != nil {
		logBadProp(d.ID, `type`, err)
	}
	if err := d.getIconName(); err != nil {
		logBadProp(d.ID, `iconName`, err)
	}
	if err := d.getStatusIconName(); err != nil {
		logBadProp(d.ID, `statusIconName`, err)
	}
	if err := d.getIsTrusted(); err != nil {
		logBadProp(d.ID, `isTrusted`, err)
	}
	if err := d.getIsReachable(); err != nil {
		logBadProp(d.ID, `isReachable`, err)
	}
	if err := d.getSupportedPlugins(); err != nil {
		logBadProp(d.ID, `supportedPlugins`, err)
	}

	return nil
}

func (d *Device) share(url string) error {
	if err := d.supported(pluginShare); err != nil {
		return err
	}
	return d.conn.Object(dest, d.obj.Path()+`/share`).Call(`shareUrl`, 0, url).Err
}

func (d *Device) supported(plugin string) error {
	d.RLock()
	defer d.RUnlock()
	if _, ok := d.SupportedPlugins[plugin]; !ok {
		return fmt.Errorf("Device does not currently support %s", plugin)
	}
	if !d.IsReachable {
		return fmt.Errorf("Device is not reachable")
	}
	if !d.IsTrusted {
		return fmt.Errorf("Device is not trusted")
	}

	return nil
}

// Close cleans up device signals, and removes it from the global list
func (d *Device) Close() error {
	devices.delete(d.ID)
	d.conn.RemoveSignal(d.signal)
	return d.conn.Close()
}

func (d *deviceList) getDevices() error {
	var ids []string

	obj := d.conn.Object(dest, path)
	// Find known devices, include unreachable, but exclude unpaired
	if err := obj.Call(`devices`, 0, false, true).Store(&ids); err != nil {
		return err
	}

	for _, id := range ids {
		var err error
		if _, ok := d.get(id); ok {
			continue
		}
		err = d.add(id)
		if err != nil {
			log(err)
			continue
		}
	}

	return nil
}

func logBadProp(id, prop string, err error) {
	log(fmt.Errorf("Device %s missing property (%s): %v\n", id, prop, err))
}

func newDevice(id string) (*Device, error) {
	conn, err := dbus.SessionBus()
	obj := conn.Object(dest, dbus.ObjectPath(fmt.Sprintf("%s/devices/%s", path, id)))
	if err != nil {
		return nil, err
	}
	d := &Device{
		ID:     id,
		conn:   conn,
		obj:    obj,
		signal: make(chan *dbus.Signal, 10),
	}
	if err = d.update(); err != nil {
		return nil, err
	}

	if err = d.watch(); err != nil {
		return nil, err
	}

	return d, nil
}

func newDeviceList() (*deviceList, error) {
	conn, err := dbus.SessionBus()
	if err != nil {
		return nil, err
	}
	return &deviceList{
		devices: make(map[string]*Device),
		conn:    conn,
	}, nil
}
<commit_msg>Add new reachableChanged signal for KDE Connect v1.2<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"strings"
	"sync"

	"github.com/godbus/dbus"
)

const (
	dest = `org.kde.kdeconnect`
	path = `/modules/kdeconnect`

	signalReachableStatusChanged = dest + `.device.reachableStatusChanged`
	signalStateChanged           = dest + `.device.stateChanged`
	signalTrustedChanged         = dest + `.device.trustedChanged`
	signalNameChanged            = dest + `.device.nameChanged`
	signalPluginsChanged         = dest + `.device.pluginsChanged`

	pluginShare = `kdeconnect_share`
)

type deviceList struct {
	devices map[string]*Device
	conn    *dbus.Conn
	sync.RWMutex
}

func (d *deviceList) get(id string) (*Device, bool) {
	d.RLock()
	defer d.RUnlock()
	dev, ok := d.devices[id]
	return dev, ok
}

func (d *deviceList) add(id string) error {
	if _, ok := d.get(id); ok {
		return fmt.Errorf("Device already exists: %s", id)
	}

	d.Lock()
	defer d.Unlock()
	dev, err := newDevice(id)
	if err != nil {
		return err
	}
	d.devices[id] = dev

	return nil
}

func (d *deviceList) delete(id string) {
	d.Lock()
	delete(d.devices, id)
	d.Unlock()
}

func (d *deviceList) all() map[string]*Device {
	d.RLock()
	defer d.RUnlock()
	return d.devices
}

func (d *deviceList) Close() error {
	err := d.conn.Close()
	for _, d := range d.devices {
		if e := d.Close(); err != nil {
			log(e)
			if err == nil {
				err = e
			}
		}
	}
	return err
}

// Device maps to the DBUS device interface
type Device struct {
	ID               string              `json:"id"`
	Type             string              `json:"type"`
	Name             string              `json:"name"`
	IconName         string              `json:"iconName"`
	StatusIconName   string              `json:"statusIconName"`
	IsReachable      bool                `json:"isReachable"`
	IsTrusted        bool                `json:"isTrusted"`
	SupportedPlugins map[string]struct{} `json:"supportedPlugins"`
	conn             *dbus.Conn
	obj              dbus.BusObject
	signal           chan *dbus.Signal
	sync.RWMutex
}

func (d *Device) watch() error {
	// kdeconnect < v1.2
	if err := d.addMatchSignal(`reachableStatusChanged`); err != nil {
		return err
	}
	// kdeconnect >= v1.2
	if err := d.addMatchSignal(`reachableChanged`); err != nil {
		return err
	}

	if err := d.addMatchSignal(`stateChanged`); err != nil {
		return err
	}
	if err := d.addMatchSignal(`trustedChanged`); err != nil {
		return err
	}
	if err := d.addMatchSignal(`nameChanged`); err != nil {
		return err
	}
	if err := d.addMatchSignal(`pluginsChanged`); err != nil {
		return err
	}

	d.conn.Signal(d.signal)
	go func() {
		for s := range d.signal {
			var err error
			switch s.Name {
			case signalReachableStatusChanged:
				if err = d.getIsReachable(); err != nil {
					log(err)
				}
			case signalPluginsChanged:
				if err = d.getSupportedPlugins(); err != nil {
					log(err)
				}
			case signalNameChanged:
				if err = d.getName(); err != nil {
					log(err)
				}
			case signalTrustedChanged:
				if err = d.getIsTrusted(); err != nil {
					log(err)
				}
			case signalStateChanged:
				if err = d.update(); err != nil {
					log(err)
				}
			default:
				if err = d.update(); err != nil {
					log(err)
				}
			}
			update := &message{
				Type: typeDeviceUpdate,
			}
			if update.Data, err = json.Marshal(d); err != nil {
				log(err)
				continue
			}
			messageQueue <- update
		}
	}()

	return nil
}

func (d *Device) addMatchSignal(member string) error {
	call := d.conn.BusObject().Call(
		`org.freedesktop.DBus.AddMatch`,
		0,
		fmt.Sprintf("type='signal',path='%s',interface='%s.device',member='%s'", d.obj.Path(), dest, member),
	)
	return call.Err
}

func (d *Device) getType() error {
	v, err := d.obj.GetProperty(dest + `.device.type`)
	if err != nil {
		return err
	}
	d.Type = strings.Trim(v.String(), `"`)

	return nil
}

func (d *Device) getName() error {
	v, err := d.obj.GetProperty(dest + `.device.name`)
	if err != nil {
		return err
	}
	d.Name = strings.Trim(v.String(), `"`)

	return nil
}

func (d *Device) getIconName() error {
	v, err := d.obj.GetProperty(dest + `.device.iconName`)
	if err != nil {
		return err
	}
	d.IconName = strings.Trim(v.String(), `"`)

	return nil
}

func (d *Device) getStatusIconName() error {
	v, err := d.obj.GetProperty(dest + `.device.statusIconName`)
	if err != nil {
		return err
	}
	d.StatusIconName = strings.Trim(v.String(), `"`)

	return nil
}

func (d *Device) getIsReachable() error {
	v, err := d.obj.GetProperty(dest + `.device.isReachable`)
	if err != nil {
		return err
	}
	d.IsReachable = v.Value().(bool)

	return nil
}

func (d *Device) getIsTrusted() error {
	v, err := d.obj.GetProperty(dest + `.device.isTrusted`)
	if err != nil {
		return err
	}
	d.IsTrusted = v.Value().(bool)

	return nil
}

func (d *Device) getSupportedPlugins() error {
	v, err := d.obj.GetProperty(dest + `.device.supportedPlugins`)
	if err != nil {
		return err
	}
	plugins := make(map[string]struct{})
	for _, plugin := range v.Value().([]string) {
		plugins[plugin] = struct{}{}
	}
	d.Lock()
	d.SupportedPlugins = plugins
	d.Unlock()

	return nil
}

func (d *Device) update() error {
	if err := d.getName(); err != nil {
		logBadProp(d.ID, `name`, err)
	}
	if err := d.getType(); err != nil {
		logBadProp(d.ID, `type`, err)
	}
	if err := d.getIconName(); err != nil {
		logBadProp(d.ID, `iconName`, err)
	}
	if err := d.getStatusIconName(); err != nil {
		logBadProp(d.ID, `statusIconName`, err)
	}
	if err := d.getIsTrusted(); err != nil {
		logBadProp(d.ID, `isTrusted`, err)
	}
	if err := d.getIsReachable(); err != nil {
		logBadProp(d.ID, `isReachable`, err)
	}
	if err := d.getSupportedPlugins(); err != nil {
		logBadProp(d.ID, `supportedPlugins`, err)
	}

	return nil
}

func (d *Device) share(url string) error {
	if err := d.supported(pluginShare); err != nil {
		return err
	}
	return d.conn.Object(dest, d.obj.Path()+`/share`).Call(`shareUrl`, 0, url).Err
}

func (d *Device) supported(plugin string) error {
	d.RLock()
	defer d.RUnlock()
	if _, ok := d.SupportedPlugins[plugin]; !ok {
		return fmt.Errorf("Device does not currently support %s", plugin)
	}
	if !d.IsReachable {
		return fmt.Errorf("Device is not reachable")
	}
	if !d.IsTrusted {
		return fmt.Errorf("Device is not trusted")
	}

	return nil
}

// Close cleans up device signals, and removes it from the global list
func (d *Device) Close() error {
	devices.delete(d.ID)
	d.conn.RemoveSignal(d.signal)
	return d.conn.Close()
}

func (d *deviceList) getDevices() error {
	var ids []string

	obj := d.conn.Object(dest, path)
	// Find known devices, include unreachable, but exclude unpaired
	if err := obj.Call(`devices`, 0, false, true).Store(&ids); err != nil {
		return err
	}

	for _, id := range ids {
		var err error
		if _, ok := d.get(id); ok {
			continue
		}
		err = d.add(id)
		if err != nil {
			log(err)
			continue
		}
	}

	return nil
}

func logBadProp(id, prop string, err error) {
	log(fmt.Errorf("Device %s missing property (%s): %v\n", id, prop, err))
}

func newDevice(id string) (*Device, error) {
	conn, err := dbus.SessionBus()
	obj := conn.Object(dest, dbus.ObjectPath(fmt.Sprintf("%s/devices/%s", path, id)))
	if err != nil {
		return nil, err
	}
	d := &Device{
		ID:     id,
		conn:   conn,
		obj:    obj,
		signal: make(chan *dbus.Signal, 10),
	}
	if err = d.update(); err != nil {
		return nil, err
	}

	if err = d.watch(); err != nil {
		return nil, err
	}

	return d, nil
}

func newDeviceList() (*deviceList, error) {
	conn, err := dbus.SessionBus()
	if err != nil {
		return nil, err
	}
	return &deviceList{
		devices: make(map[string]*Device),
		conn:    conn,
	}, nil
}
<|endoftext|>


<commit_before>package irest

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

type Test struct {
	Name   string `json:"name"`
	Error  error  `json:"err"`
	Status int    `json:"status"`

	Tests    []*Test
	Errors   []error
	Created  time.Time `json:"created"`
	Duration int64     `json:"duration"`

	Client   *http.Client
	Header   *http.Header
	Cookie   *http.Cookie
	Response *http.Response
}

func NewTest(name string) *Test {
	t := &Test{
		Name:    name,
		Error:   nil,
		Tests:   []*Test{},
		Errors:  []error{},
		Created: time.Now(),
		Client:  &http.Client{},
		Header:  &http.Header{},
	}

	return t
}

func (t *Test) NewTest(name string) *Test {
	testCase := &Test{
		Name:   name,
		Tests:  []*Test{},
		Client: t.Client,
		Header: &http.Header{},
	}

	t.Tests = append(t.Tests, testCase)

	return testCase
}

func (t *Test) AddHeader(name, value string) *Test {
	t.Header.Set(name, value)
	return t
}

func (t *Test) Post(baseURL, endpoint string, result interface{}) *Test {
	t.Header.Set("Content-Type", "application/json")

	addr, err := url.Parse(baseURL + endpoint)
	if err != nil {
		t.Error = err
		return t
	}

	req, err := http.NewRequest("POST", addr.String(), nil)
	if err != nil {
		t.Error = err
		return t
	}

	res, err := t.Client.Do(req)
	if err != nil {
		t.Error = err
		return t
	}

	t.Response = res
	t.Status = res.StatusCode

	body := res.Body
	defer body.Close()

	data, err := ioutil.ReadAll(res.Body)

	if err != nil {
		t.Error = err
		return t
	}

	json.Unmarshal(data, result)

	return t
}

func (t *Test) MustStatus(statusCode int) *Test {
	if t.Error != nil {
		return t
	}

	if t.Status != statusCode {
		t.Error = fmt.Errorf("expected status code response of %d, actual %d", statusCode, t.Status)
	}

	return t
}

func (t *Test) MustStringValue(expected, actual string) *Test {
	if t.Error != nil {
		return t
	}

	if expected != actual {
		t.Error = fmt.Errorf("expected %s, but got %s", expected, actual)
	}

	return t
}

func (t *Test) MustIntValue(expected, actual int) *Test {
	if t.Error != nil {
		return t
	}

	if expected != actual {
		t.Error = fmt.Errorf("expected %d, but got %d", expected, actual)
	}

	return t
}

// MustFunction adds the ability to create functions that can check something
// not covered by the current functions currently.
type MustFunction func() error

func (t *Test) Must(fn MustFunction) *Test {
	if t.Error != nil {
		return t
	}

	if err := fn(); err != nil {
		t.Error = err
	}

	return t
}

func (t *Test) SaveCookie(name string, cookie *http.Cookie) *Test {
	if t.Error != nil {
		return t
	}

	if t.Response == nil {
		t.Error = fmt.Errorf("http response not set, must have request before saving result")
		return t
	}

	for _, c := range t.Response.Cookies() {
		if c.Name == name {
			cookie.Name = c.Name
			cookie.Value = c.Value
			return t
		}
	}

	t.Error = fmt.Errorf("cookie name '%s' not found", name)

	return t
}
<commit_msg>Add initial documentation for functions and package.<commit_after>// Package iREST is an integration testing package for RESTful
// APIs. It simply makes HTTP requests and allows for checking of
// responses, status codes, etc.

package irest

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

// Test struct contains sub-tests that can be isolated test cases as well
// as the HTTP related objects and errors of current test and its sub-tests.
type Test struct {
	Name     string `json:"name"`
	Error    error  `json:"err"`
	Status   int    `json:"status"`
	Tests    []*Test
	Errors   []error
	Created  time.Time `json:"created"`
	Duration int64     `json:"duration"`

	// HTTP related fields for making requests and getting responses.
	Client   *http.Client
	Header   *http.Header
	Cookie   *http.Cookie
	Response *http.Response
}

// NewTest creates a new test with a given name.
func NewTest(name string) *Test {
	t := &Test{
		Name:    name,
		Error:   nil,
		Tests:   []*Test{},
		Errors:  []error{},
		Created: time.Now(),
		Client:  &http.Client{},
		Header:  &http.Header{},
	}

	return t
}

// NewTest adds a Test as a sub-test to the current one. Sub-tests act as
// individual test cases.
func (t *Test) NewTest(name string) *Test {
	testCase := &Test{
		Name:   name,
		Tests:  []*Test{},
		Client: t.Client,
		Header: &http.Header{},
	}

	t.Tests = append(t.Tests, testCase)

	return testCase
}

// AddHeader is a utility function to just wrap setting a header with a value
// by name.
func (t *Test) AddHeader(name, value string) *Test {
	t.Header.Set(name, value)
	return t
}

// Post sends a HTTP Post request with given URL from baseURL combined with
// endpoint and then saves the result.
func (t *Test) Post(baseURL, endpoint string, result interface{}) *Test {
	t.Header.Set("Content-Type", "application/json")

	addr, err := url.Parse(baseURL + endpoint)
	if err != nil {
		t.Error = err
		return t
	}

	req, err := http.NewRequest("POST", addr.String(), nil)
	if err != nil {
		t.Error = err
		return t
	}

	res, err := t.Client.Do(req)
	if err != nil {
		t.Error = err
		return t
	}

	t.Response = res
	t.Status = res.StatusCode

	body := res.Body
	defer body.Close()

	data, err := ioutil.ReadAll(res.Body)

	if err != nil {
		t.Error = err
		return t
	}

	json.Unmarshal(data, result)

	return t
}

// MustStatus sets the Test.Error if the status code is not the expected
// value. An HTTP request must have been made prior to this function call.
func (t *Test) MustStatus(statusCode int) *Test {
	if t.Error != nil {
		return t
	}

	if t.Status != statusCode {
		t.Error = fmt.Errorf("expected status code response of %d, actual %d", statusCode, t.Status)
	}

	return t
}

// MustStringValue compares two string values and sets the Test.Error if not
// equal.
func (t *Test) MustStringValue(expected, actual string) *Test {
	if t.Error != nil {
		return t
	}

	if expected != actual {
		t.Error = fmt.Errorf("expected %s, but got %s", expected, actual)
	}

	return t
}

// MustIntValue compares two int values and sets the Test.Error if not equal.
func (t *Test) MustIntValue(expected, actual int) *Test {
	if t.Error != nil {
		return t
	}

	if expected != actual {
		t.Error = fmt.Errorf("expected %d, but got %d", expected, actual)
	}

	return t
}

// MustFunction adds the ability to create functions that can check something
// not covered by the current functions currently.
type MustFunction func() error

// Must allows for passing in created functions matching the MustFunction
// pattern with no parameters returning an error.
func (t *Test) Must(fn MustFunction) *Test {
	if t.Error != nil {
		return t
	}

	if err := fn(); err != nil {
		t.Error = err
	}

	return t
}

// SaveCookie will store the cookie with the specified name if it exists in the
// response. An HTTP request must have been made prior to this function call.
func (t *Test) SaveCookie(name string, cookie *http.Cookie) *Test {
	if t.Error != nil {
		return t
	}

	if t.Response == nil {
		t.Error = fmt.Errorf("http response not set, must have request before saving result")
		return t
	}

	for _, c := range t.Response.Cookies() {
		if c.Name == name {
			cookie.Name = c.Name
			cookie.Value = c.Value
			return t
		}
	}

	t.Error = fmt.Errorf("cookie name '%s' not found", name)

	return t
}
<|endoftext|>


<commit_before>// Package iREST is an integration testing package for RESTful
// APIs. It simply makes HTTP requests and allows for checking of
// responses, status codes, etc.
package irest

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"
)

// Test struct contains sub-tests that can be isolated test cases as well
// as the HTTP related objects and errors of current test and its sub-tests.
type Test struct {
	Name     string `json:"name"`
	Error    error  `json:"err"`
	Status   int    `json:"status"`
	Tests    []*Test
	Errors   []error
	Created  time.Time `json:"created"`
	Duration int64     `json:"duration"`

	// HTTP related fields for making requests and getting responses.
	Client   *http.Client
	Header   *http.Header
	Cookie   *http.Cookie
	Response *http.Response
}

// NewTest creates a new test with a given name.
func NewTest(name string) *Test {
	t := &Test{
		Name:    name,
		Error:   nil,
		Tests:   []*Test{},
		Errors:  []error{},
		Created: time.Now(),
		Client:  &http.Client{},
		Header:  &http.Header{},
	}

	return t
}

// NewTest adds a Test as a sub-test to the current one. Sub-tests act as
// individual test cases.
func (t *Test) NewTest(name string) *Test {
	testCase := &Test{
		Name:   name,
		Tests:  []*Test{},
		Client: t.Client,
		Header: &http.Header{},
	}

	t.Tests = append(t.Tests, testCase)

	return testCase
}

// AddHeader is a utility function to just wrap setting a header with a value
// by name.
func (t *Test) AddHeader(name, value string) *Test {
	t.Header.Set(name, value)
	return t
}

// Get retrieves data from a specified endpoint.
func (t *Test) Get(baseURL, endpoint string, result interface{}) *Test {
	if t.Error != nil {
		return t
	}

	addr, err := url.Parse(baseURL + endpoint)
	if err != nil {
		t.Error = err
		return t
	}

	req, err := http.NewRequest("GET", addr.String(), nil)
	if err != nil {
		t.Error = err
		return t
	}

	res, err := t.Client.Do(req)
	if err != nil {
		t.Error = err
		return t
	}

	t.Response = res
	t.Status = res.StatusCode

	body := res.Body
	defer body.Close()

	resultBody, err := ioutil.ReadAll(res.Body)

	if err != nil {
		t.Error = err
		return t
	}

	if err := json.Unmarshal(resultBody, result); err != nil {
		t.Error = err
		return t
	}

	return t
}

// Post sends a HTTP Post request with given URL from baseURL combined with
// endpoint and then saves the result.
func (t *Test) Post(baseURL, endpoint string, data, result interface{}) *Test {
	t.Header.Set("Content-Type", "application/json")

	addr, err := url.Parse(baseURL + endpoint)
	if err != nil {
		t.Error = err
		return t
	}

	b := new(bytes.Buffer)
	if err := json.NewEncoder(b).Encode(data); err != nil {
		t.Error = err
		return t
	}

	req, err := http.NewRequest("POST", addr.String(), b)
	if err != nil {
		t.Error = err
		return t
	}

	res, err := t.Client.Do(req)
	if err != nil {
		t.Error = err
		return t
	}

	t.Response = res
	t.Status = res.StatusCode

	body := res.Body
	defer body.Close()

	resultBody, err := ioutil.ReadAll(res.Body)

	if err != nil {
		t.Error = err
		return t
	}

	if err := json.Unmarshal(resultBody, result); err != nil {
		t.Error = err
		return t
	}

	return t
}

// MustStatus sets the Test.Error if the status code is not the expected
// value. An HTTP request must have been made prior to this function call.
func (t *Test) MustStatus(statusCode int) *Test {
	if t.Error != nil {
		return t
	}

	if t.Status != statusCode {
		t.Error = fmt.Errorf("expected status code response of %d, actual %d", statusCode, t.Status)
	}

	return t
}

// MustStringValue compares two string values and sets the Test.Error if not
// equal.
func (t *Test) MustStringValue(expected, actual string) *Test {
	if t.Error != nil {
		return t
	}

	if expected != actual {
		t.Error = fmt.Errorf("expected %s, but got %s", expected, actual)
	}

	return t
}

// MustIntValue compares two int values and sets the Test.Error if not equal.
func (t *Test) MustIntValue(expected, actual int) *Test {
	if t.Error != nil {
		return t
	}

	if expected != actual {
		t.Error = fmt.Errorf("expected %d, but got %d", expected, actual)
	}

	return t
}

// MustFunction adds the ability to create functions that can check something
// not covered by the current functions currently.
type MustFunction func() error

// Must allows for passing in created functions matching the MustFunction
// pattern with no parameters returning an error.
func (t *Test) Must(fn MustFunction) *Test {
	if t.Error != nil {
		return t
	}

	if err := fn(); err != nil {
		t.Error = err
	}

	return t
}

// SaveCookie will store the cookie with the specified name if it exists in the
// response. An HTTP request must have been made prior to this function call.
func (t *Test) SaveCookie(name string, cookie *http.Cookie) *Test {
	if t.Error != nil {
		return t
	}

	if t.Response == nil {
		t.Error = fmt.Errorf("http response not set, must have request before saving result")
		return t
	}

	for _, c := range t.Response.Cookies() {
		if c.Name == name {
			cookie.Name = c.Name
			cookie.Value = c.Value
			return t
		}
	}

	t.Error = fmt.Errorf("cookie name '%s' not found", name)

	return t
}
<commit_msg>Refactor post function into a common function between put and post.<commit_after>// Package iREST is an integration testing package for RESTful
// APIs. It simply makes HTTP requests and allows for checking of
// responses, status codes, etc.
package irest

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"
)

// Test struct contains sub-tests that can be isolated test cases as well
// as the HTTP related objects and errors of current test and its sub-tests.
type Test struct {
	Name     string `json:"name"`
	Error    error  `json:"err"`
	Status   int    `json:"status"`
	Tests    []*Test
	Errors   []error
	Created  time.Time `json:"created"`
	Duration int64     `json:"duration"`

	// HTTP related fields for making requests and getting responses.
	Client   *http.Client
	Header   *http.Header
	Cookie   *http.Cookie
	Response *http.Response
}

// NewTest creates a new test with a given name.
func NewTest(name string) *Test {
	t := &Test{
		Name:    name,
		Error:   nil,
		Tests:   []*Test{},
		Errors:  []error{},
		Created: time.Now(),
		Client:  &http.Client{},
		Header:  &http.Header{},
	}

	return t
}

// NewTest adds a Test as a sub-test to the current one. Sub-tests act as
// individual test cases.
func (t *Test) NewTest(name string) *Test {
	testCase := &Test{
		Name:   name,
		Tests:  []*Test{},
		Client: t.Client,
		Header: &http.Header{},
	}

	t.Tests = append(t.Tests, testCase)

	return testCase
}

// AddHeader is a utility function to just wrap setting a header with a value
// by name.
func (t *Test) AddHeader(name, value string) *Test {
	t.Header.Set(name, value)
	return t
}

// Get retrieves data from a specified endpoint.
func (t *Test) Get(baseURL, endpoint string, result interface{}) *Test {
	if t.Error != nil {
		return t
	}

	addr, err := url.Parse(baseURL + endpoint)
	if err != nil {
		t.Error = err
		return t
	}

	req, err := http.NewRequest("GET", addr.String(), nil)
	if err != nil {
		t.Error = err
		return t
	}

	res, err := t.Client.Do(req)
	if err != nil {
		t.Error = err
		return t
	}

	t.Response = res
	t.Status = res.StatusCode

	body := res.Body
	defer body.Close()

	resultBody, err := ioutil.ReadAll(res.Body)

	if err != nil {
		t.Error = err
		return t
	}

	if err := json.Unmarshal(resultBody, result); err != nil {
		t.Error = err
		return t
	}

	return t
}

// Put sends a HTTP PUT request with given URL from baseURL combined with
// endpoint and then saves the result.
func (t *Test) Post(baseURL, endpoint string, data, result interface{}) *Test {
	if err := t.DoSaveRequest(baseURL+endpoint, "POST", data, result); err != nil {
		t.Error = err
	}

	return t
}

// Post sends a HTTP POST request with given URL from baseURL combined with
// endpoint and then saves the result.
func (t *Test) Put(baseURL, endpoint string, data, result interface{}) *Test {
	if err := t.DoSaveRequest(baseURL+endpoint, "PUT", data, result); err != nil {
		t.Error = err
	}

	return t
}

// DoSaveRequest extracts similar functionality of both PUT and POST into
// a single function.
func (t *Test) DoSaveRequest(path, method string, data, result interface{}) error {
	t.Header.Set("Content-Type", "application/json")

	addr, err := url.Parse(path)
	if err != nil {
		return err
	}

	b := new(bytes.Buffer)
	if err := json.NewEncoder(b).Encode(data); err != nil {
		return err
	}

	req, err := http.NewRequest(method, addr.String(), b)
	if err != nil {
		return err
	}

	res, err := t.Client.Do(req)
	if err != nil {
		return err
	}

	t.Response = res
	t.Status = res.StatusCode

	body := res.Body
	defer body.Close()

	resultBody, err := ioutil.ReadAll(res.Body)

	if err != nil {
		return err
	}

	return json.Unmarshal(resultBody, result)
}

// MustStatus sets the Test.Error if the status code is not the expected
// value. An HTTP request must have been made prior to this function call.
func (t *Test) MustStatus(statusCode int) *Test {
	if t.Error != nil {
		return t
	}

	if t.Status != statusCode {
		t.Error = fmt.Errorf("expected status code response of %d, actual %d", statusCode, t.Status)
	}

	return t
}

// MustStringValue compares two string values and sets the Test.Error if not
// equal.
func (t *Test) MustStringValue(expected, actual string) *Test {
	if t.Error != nil {
		return t
	}

	if expected != actual {
		t.Error = fmt.Errorf("expected %s, but got %s", expected, actual)
	}

	return t
}

// MustIntValue compares two int values and sets the Test.Error if not equal.
func (t *Test) MustIntValue(expected, actual int) *Test {
	if t.Error != nil {
		return t
	}

	if expected != actual {
		t.Error = fmt.Errorf("expected %d, but got %d", expected, actual)
	}

	return t
}

// MustFunction adds the ability to create functions that can check something
// not covered by the current functions currently.
type MustFunction func() error

// Must allows for passing in created functions matching the MustFunction
// pattern with no parameters returning an error.
func (t *Test) Must(fn MustFunction) *Test {
	if t.Error != nil {
		return t
	}

	if err := fn(); err != nil {
		t.Error = err
	}

	return t
}

// SaveCookie will store the cookie with the specified name if it exists in the
// response. An HTTP request must have been made prior to this function call.
func (t *Test) SaveCookie(name string, cookie *http.Cookie) *Test {
	if t.Error != nil {
		return t
	}

	if t.Response == nil {
		t.Error = fmt.Errorf("http response not set, must have request before saving result")
		return t
	}

	for _, c := range t.Response.Cookies() {
		if c.Name == name {
			cookie.Name = c.Name
			cookie.Value = c.Value
			return t
		}
	}

	t.Error = fmt.Errorf("cookie name '%s' not found", name)

	return t
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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
//
// https://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 fields contains tests related to fields accessors
package fields

import (
	"example.com/core"
)

func TestFieldAccessors(s core.Source, ptr *core.Source) {
	core.Sinkf("Data: %v", s.GetData()) // want "a source has reached a sink"
	core.Sinkf("ID: %v", s.GetID())

	core.Sinkf("Data: %v", ptr.GetData()) // want "a source has reached a sink"
	core.Sinkf("ID: %v", ptr.GetID())
}

func TestDirectFieldAccess(c *core.Source) {
	core.Sinkf("Data: %v", c.Data) // want "a source has reached a sink"
	core.Sinkf("ID: %v", c.ID)
}

func TestProtoStyleFieldAccessorSanitizedPII(c *core.Source) {
	core.Sinkf("Source data: %v", core.Sanitize(c.GetData()))
}

func TestProtoStyleFieldAccessorPIISecondLevel(wrapper struct{ *core.Source }) {
	core.Sinkf("Source data: %v", wrapper.Source.GetData()) // want "a source has reached a sink"
	core.Sinkf("Source id: %v", wrapper.Source.GetID())
}

func TestDirectFieldAccessorPIISecondLevel(wrapper struct{ *core.Source }) {
	core.Sinkf("Source data: %v", wrapper.Source.Data) // want "a source has reached a sink"
	core.Sinkf("Source id: %v", wrapper.Source.ID)
}

func TestTaggedStruct(s core.TaggedSource) {
	core.Sink(s) // want "a source has reached a sink"
}

func TestTaggedAndNonTaggedFields(s core.TaggedSource) {
	core.Sink(s.Data) // want "a source has reached a sink"
	core.Sink(s.ID)
}
<commit_msg>Add tests for field propagation (#195)<commit_after>// Copyright 2020 Google LLC
//
// 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
//
// https://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 fields contains tests related to fields accessors
package fields

import (
	"strconv"

	"example.com/core"
)

func TestFieldAccessors(s core.Source, ptr *core.Source) {
	core.Sinkf("Data: %v", s.GetData()) // want "a source has reached a sink"
	core.Sinkf("ID: %v", s.GetID())

	core.Sinkf("Data: %v", ptr.GetData()) // want "a source has reached a sink"
	core.Sinkf("ID: %v", ptr.GetID())
}

func TestDirectFieldAccess(c *core.Source) {
	core.Sinkf("Data: %v", c.Data) // want "a source has reached a sink"
	core.Sinkf("ID: %v", c.ID)
}

func TestProtoStyleFieldAccessorSanitizedPII(c *core.Source) {
	core.Sinkf("Source data: %v", core.Sanitize(c.GetData()))
}

func TestProtoStyleFieldAccessorPIISecondLevel(wrapper struct{ *core.Source }) {
	core.Sinkf("Source data: %v", wrapper.Source.GetData()) // want "a source has reached a sink"
	core.Sinkf("Source id: %v", wrapper.Source.GetID())
}

func TestDirectFieldAccessorPIISecondLevel(wrapper struct{ *core.Source }) {
	core.Sinkf("Source data: %v", wrapper.Source.Data) // want "a source has reached a sink"
	core.Sinkf("Source id: %v", wrapper.Source.ID)
}

func TestTaggedStruct(s core.TaggedSource) {
	core.Sink(s) // want "a source has reached a sink"
}

func TestTaggedAndNonTaggedFields(s core.TaggedSource) {
	core.Sink(s.Data) // want "a source has reached a sink"
	core.Sink(s.ID)
}

func TestTaintFieldOnNonSourceStruct(s core.Source, i *core.Innocuous) {
	i.Data = s.Data
	core.Sink(i)      // TODO(#228) want "a source has reached a sink"
	core.Sink(i.Data) // TODO(#228) want "a source has reached a sink"
}

func TestTaintNonSourceFieldOnSourceType(s core.Source, i *core.Innocuous) {
	s.ID, _ = strconv.Atoi(s.Data)
	core.Sink(s.ID) // TODO(#228) want "a source has reached a sink"
}

type Headers struct {
	Name  string
	Auth  map[string]string `levee:"source"`
	Other map[string]string
}

func fooByPtr(h *Headers) {}

func foo(h Headers) {}

func TestCallWithStructReferenceTaintsEveryField(h Headers) {
	fooByPtr(&h)       // without interprocedural assessment, foo can do anything, so this call should taint every field on h
	core.Sink(h.Name)  // TODO(#229) want "a source has reached a sink"
	core.Sink(h.Other) // TODO(#229) want "a source has reached a sink"
}

func TestCallWithStructValueDoesNotTaintNonReferenceFields(h Headers) {
	foo(h) // h is passed by value, so only its reference-like fields should be tainted
	core.Sink(h.Name)
	core.Sink(h.Other) // TODO(#229) want "a source has reached a sink"
}
<|endoftext|>


<commit_before>package boltdb

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"

	"github.com/Workiva/go-datastructures/augmentedtree"
	region "github.com/akhenakh/regionagogo"
	"github.com/akhenakh/regionagogo/geostore"
	"github.com/boltdb/bolt"
	"github.com/golang/geo/s2"
	"github.com/golang/protobuf/proto"
	lru "github.com/hashicorp/golang-lru"
)

const (
	loopBucket  = "loop"
	coverBucket = "cover"
)

// GeoFenceBoltDB provides an in memory index and boltdb query engine for fences lookup
type GeoFenceBoltDB struct {
	augmentedtree.Tree
	*bolt.DB
	cache *lru.Cache
	debug bool
}

// GeoSearchOption used to pass options to NewGeoSearch
type GeoFenceBoltDBOption func(*geoFenceBoltDBOptions)

type geoFenceBoltDBOptions struct {
	maxCachedEntries uint
	debug            bool
}

// WithCachedEntries enable an LRU cache default is disabled
func WithCachedEntries(maxCachedEntries uint) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.maxCachedEntries = maxCachedEntries
	}
}

// WithDebug enable debug
func WithDebug(debug bool) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.debug = debug
	}
}

// NewGeoFenceBoltDB creates a new geo database, needs a writable path for BoltDB
func NewGeoFenceBoltDB(dbpath string, opts ...GeoFenceBoltDBOption) (region.GeoFenceDB, error) {
	db, err := bolt.Open(dbpath, 0600, nil)
	if err != nil {
		return nil, err
	}

	if errdb := db.Update(func(tx *bolt.Tx) error {
		if _, errtx := tx.CreateBucketIfNotExists([]byte(loopBucket)); errtx != nil {
			return fmt.Errorf("create bucket: %s", errtx)
		}
		if _, errtx := tx.CreateBucketIfNotExists([]byte(coverBucket)); errtx != nil {
			return fmt.Errorf("create bucket: %s", errtx)
		}
		return nil
	}); errdb != nil {
		return nil, errdb
	}

	var geoOpts geoFenceBoltDBOptions

	for _, opt := range opts {
		opt(&geoOpts)
	}

	gs := &GeoFenceBoltDB{
		Tree:  augmentedtree.New(1),
		DB:    db,
		debug: geoOpts.debug,
	}

	if geoOpts.maxCachedEntries != 0 {
		cache, err := lru.New(int(geoOpts.maxCachedEntries))
		if err != nil {
			return nil, err
		}
		gs.cache = cache
	}

	if err := gs.importGeoData(); err != nil {
		return nil, err
	}

	return gs, nil
}

// index indexes each cells of the cover and set its loopID
func (gs *GeoFenceBoltDB) index(fc *geostore.FenceCover, loopID uint64) {
	for _, cell := range fc.Cellunion {
		s2interval := &region.S2Interval{CellID: s2.CellID(cell)}
		intervals := gs.Query(s2interval)
		found := false

		if len(intervals) != 0 {
			for _, existInterval := range intervals {
				if existInterval.LowAtDimension(1) == s2interval.LowAtDimension(1) &&
					existInterval.HighAtDimension(1) == s2interval.HighAtDimension(1) {
					if gs.debug {
						log.Println("added to existing interval", s2interval, loopID)
					}
					s2interval.LoopIDs = append(s2interval.LoopIDs, loopID)
					found = true
					break
				}
			}
		}

		if !found {
			// create new interval with current loop
			s2interval.LoopIDs = []uint64{loopID}
			gs.Add(s2interval)
			if gs.debug {
				log.Println("added new interval", s2interval, loopID)
			}
		}
	}
}

// importGeoData loads all existing cells into the segment tree
func (gs *GeoFenceBoltDB) importGeoData() error {
	var count int
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(coverBucket))
		cur := b.Cursor()

		// load the cell ranges into the tree
		var fc geostore.FenceCover
		for k, v := cur.First(); k != nil; k, v = cur.Next() {
			count++
			err := proto.Unmarshal(v, &fc)
			if err != nil {
				return err
			}
			if gs.debug {
				log.Println("read", fc.Cellunion)
			}

			// read back the loopID from the key
			var loopID uint64
			buf := bytes.NewBuffer(k)
			err = binary.Read(buf, binary.BigEndian, &loopID)
			if err != nil {
				return err
			}

			gs.index(&fc, loopID)
		}

		return nil
	})
	if err != nil {
		return err
	}

	log.Println("loaded", count, "fences")

	return nil
}

// FenceByID returns a region from DB by its id
func (gs *GeoFenceBoltDB) FenceByID(loopID uint64) *region.Fence {
	// TODO: refactor as Fence ?
	if gs.cache != nil {
		if val, ok := gs.cache.Get(loopID); ok {
			r, _ := val.(*region.Fence)
			return r
		}
	}
	var rs *geostore.FenceStorage
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(loopBucket))
		v := b.Get(itob(loopID))

		var frs geostore.FenceStorage
		err := proto.Unmarshal(v, &frs)
		if err != nil {
			return err
		}
		rs = &frs
		return nil
	})
	if err != nil {
		return nil
	}
	r := region.NewFenceFromStorage(rs)
	if gs.cache != nil && r != nil {
		gs.cache.Add(loopID, r)
	}
	return r
}

// StubbingQuery returns the fence for the corresponding lat, lng point
func (gs *GeoFenceBoltDB) StubbingQuery(lat, lng float64) *region.Fence {
	q := s2.CellIDFromLatLng(s2.LatLngFromDegrees(lat, lng))
	i := &region.S2Interval{CellID: q}

	if gs.debug {
		log.Println("lookup", lat, lng, q)
	}
	r := gs.Tree.Query(i)

	var foundRegion *region.Fence

	for _, itv := range r {
		sitv := itv.(*region.S2Interval)
		if gs.debug {
			log.Println("found", sitv, sitv.LoopIDs)
		}

		// a region can include a smaller region
		// return only the one that is contained in the other
		for _, loopID := range sitv.LoopIDs {
			region := gs.FenceByID(loopID)
			if region != nil && region.Loop.ContainsPoint(q.Point()) {
				if foundRegion == nil {
					foundRegion = region
					continue
				}
				// we take the 1st vertex of the region.Loop if it is contained in previousLoop
				// region loop  is more precise
				if foundRegion.Loop.ContainsPoint(region.Loop.Vertex(0)) {
					foundRegion = region
				}
			}
		}
	}

	return foundRegion
}

// RectQuery perform rectangular query ur upper right bl bottom left
func (gs *GeoFenceBoltDB) RectQuery(urlat, urlng, bllat, bllng float64, limit int) (region.Fences, error) {
	rect := s2.RectFromLatLng(s2.LatLngFromDegrees(bllat, bllng))
	rect = rect.AddPoint(s2.LatLngFromDegrees(urlat, urlng))

	rc := &s2.RegionCoverer{MaxLevel: 30, MaxCells: 1}
	covering := rc.Covering(rect)
	if len(covering) != 1 {
		return nil, errors.New("impossible covering")
	}
	i := &region.S2Interval{CellID: covering[0]}
	r := gs.Tree.Query(i)

	fences := make(map[uint64]*region.Fence)

	for _, itv := range r {
		sitv := itv.(*region.S2Interval)
		for _, loopID := range sitv.LoopIDs {
			var region *region.Fence
			if v, ok := fences[loopID]; ok {
				region = v
			} else {
				region = gs.FenceByID(loopID)
			}
			// testing the found loop is actually inside the rect
			// (since we are using only one large cover it may be outside)
			if rect.Contains(region.Loop.RectBound()) {
				fences[loopID] = region
			}
		}
	}

	var res []*region.Fence
	for _, v := range fences {
		res = append(res, v)
	}
	return region.Fences(res), nil
}

// StoreFence stores a fence into the database and load its index in memory
func (gs *GeoFenceBoltDB) StoreFence(fs *geostore.FenceStorage, cover []uint64) error {
	return gs.Update(func(tx *bolt.Tx) error {
		loopB := tx.Bucket([]byte(loopBucket))
		coverBucket := tx.Bucket([]byte(coverBucket))

		loopID, err := loopB.NextSequence()
		if err != nil {
			return err
		}

		buf, err := proto.Marshal(fs)
		if err != nil {
			return err
		}

		if gs.debug {
			log.Println("inserted", loopID, fs.Data, cover)
		}

		// convert our loopID to bigendian to be used as key
		k := itob(loopID)

		err = loopB.Put(k, buf)
		if err != nil {
			return err
		}

		// inserting into cover index using the same key
		fc := &geostore.FenceCover{Cellunion: cover}
		bufc, err := proto.Marshal(fc)
		if err != nil {
			return err
		}

		// also load into memory
		gs.index(fc, loopID)

		return coverBucket.Put(k, bufc)
	})
}

// itob returns an 8-byte big endian representation of v.
func itob(v uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	return b
}
<commit_msg>return a concrete type<commit_after>package boltdb

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"

	"github.com/Workiva/go-datastructures/augmentedtree"
	region "github.com/akhenakh/regionagogo"
	"github.com/akhenakh/regionagogo/geostore"
	"github.com/boltdb/bolt"
	"github.com/golang/geo/s2"
	"github.com/golang/protobuf/proto"
	lru "github.com/hashicorp/golang-lru"
)

const (
	loopBucket  = "loop"
	coverBucket = "cover"
)

// GeoFenceBoltDB provides an in memory index and boltdb query engine for fences lookup
type GeoFenceBoltDB struct {
	augmentedtree.Tree
	*bolt.DB
	cache *lru.Cache
	debug bool
}

// GeoSearchOption used to pass options to NewGeoSearch
type GeoFenceBoltDBOption func(*geoFenceBoltDBOptions)

type geoFenceBoltDBOptions struct {
	maxCachedEntries uint
	debug            bool
}

// WithCachedEntries enable an LRU cache default is disabled
func WithCachedEntries(maxCachedEntries uint) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.maxCachedEntries = maxCachedEntries
	}
}

// WithDebug enable debug
func WithDebug(debug bool) GeoFenceBoltDBOption {
	return func(o *geoFenceBoltDBOptions) {
		o.debug = debug
	}
}

// NewGeoFenceBoltDB creates a new geo database, needs a writable path for BoltDB
func NewGeoFenceBoltDB(dbpath string, opts ...GeoFenceBoltDBOption) (*GeoFenceBoltDB, error) {
	db, err := bolt.Open(dbpath, 0600, nil)
	if err != nil {
		return nil, err
	}

	if errdb := db.Update(func(tx *bolt.Tx) error {
		if _, errtx := tx.CreateBucketIfNotExists([]byte(loopBucket)); errtx != nil {
			return fmt.Errorf("create bucket: %s", errtx)
		}
		if _, errtx := tx.CreateBucketIfNotExists([]byte(coverBucket)); errtx != nil {
			return fmt.Errorf("create bucket: %s", errtx)
		}
		return nil
	}); errdb != nil {
		return nil, errdb
	}

	var geoOpts geoFenceBoltDBOptions

	for _, opt := range opts {
		opt(&geoOpts)
	}

	gs := &GeoFenceBoltDB{
		Tree:  augmentedtree.New(1),
		DB:    db,
		debug: geoOpts.debug,
	}

	if geoOpts.maxCachedEntries != 0 {
		cache, err := lru.New(int(geoOpts.maxCachedEntries))
		if err != nil {
			return nil, err
		}
		gs.cache = cache
	}

	if err := gs.importGeoData(); err != nil {
		return nil, err
	}

	return gs, nil
}

// index indexes each cells of the cover and set its loopID
func (gs *GeoFenceBoltDB) index(fc *geostore.FenceCover, loopID uint64) {
	for _, cell := range fc.Cellunion {
		s2interval := &region.S2Interval{CellID: s2.CellID(cell)}
		intervals := gs.Query(s2interval)
		found := false

		if len(intervals) != 0 {
			for _, existInterval := range intervals {
				if existInterval.LowAtDimension(1) == s2interval.LowAtDimension(1) &&
					existInterval.HighAtDimension(1) == s2interval.HighAtDimension(1) {
					if gs.debug {
						log.Println("added to existing interval", s2interval, loopID)
					}
					s2interval.LoopIDs = append(s2interval.LoopIDs, loopID)
					found = true
					break
				}
			}
		}

		if !found {
			// create new interval with current loop
			s2interval.LoopIDs = []uint64{loopID}
			gs.Add(s2interval)
			if gs.debug {
				log.Println("added new interval", s2interval, loopID)
			}
		}
	}
}

// importGeoData loads all existing cells into the segment tree
func (gs *GeoFenceBoltDB) importGeoData() error {
	var count int
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(coverBucket))
		cur := b.Cursor()

		// load the cell ranges into the tree
		var fc geostore.FenceCover
		for k, v := cur.First(); k != nil; k, v = cur.Next() {
			count++
			err := proto.Unmarshal(v, &fc)
			if err != nil {
				return err
			}
			if gs.debug {
				log.Println("read", fc.Cellunion)
			}

			// read back the loopID from the key
			var loopID uint64
			buf := bytes.NewBuffer(k)
			err = binary.Read(buf, binary.BigEndian, &loopID)
			if err != nil {
				return err
			}

			gs.index(&fc, loopID)
		}

		return nil
	})
	if err != nil {
		return err
	}

	log.Println("loaded", count, "fences")

	return nil
}

// FenceByID returns a region from DB by its id
func (gs *GeoFenceBoltDB) FenceByID(loopID uint64) *region.Fence {
	// TODO: refactor as Fence ?
	if gs.cache != nil {
		if val, ok := gs.cache.Get(loopID); ok {
			r, _ := val.(*region.Fence)
			return r
		}
	}
	var rs *geostore.FenceStorage
	err := gs.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(loopBucket))
		v := b.Get(itob(loopID))

		var frs geostore.FenceStorage
		err := proto.Unmarshal(v, &frs)
		if err != nil {
			return err
		}
		rs = &frs
		return nil
	})
	if err != nil {
		return nil
	}
	r := region.NewFenceFromStorage(rs)
	if gs.cache != nil && r != nil {
		gs.cache.Add(loopID, r)
	}
	return r
}

// StubbingQuery returns the fence for the corresponding lat, lng point
func (gs *GeoFenceBoltDB) StubbingQuery(lat, lng float64) *region.Fence {
	q := s2.CellIDFromLatLng(s2.LatLngFromDegrees(lat, lng))
	i := &region.S2Interval{CellID: q}

	if gs.debug {
		log.Println("lookup", lat, lng, q)
	}
	r := gs.Tree.Query(i)

	var foundRegion *region.Fence

	for _, itv := range r {
		sitv := itv.(*region.S2Interval)
		if gs.debug {
			log.Println("found", sitv, sitv.LoopIDs)
		}

		// a region can include a smaller region
		// return only the one that is contained in the other
		for _, loopID := range sitv.LoopIDs {
			region := gs.FenceByID(loopID)
			if region != nil && region.Loop.ContainsPoint(q.Point()) {
				if foundRegion == nil {
					foundRegion = region
					continue
				}
				// we take the 1st vertex of the region.Loop if it is contained in previousLoop
				// region loop  is more precise
				if foundRegion.Loop.ContainsPoint(region.Loop.Vertex(0)) {
					foundRegion = region
				}
			}
		}
	}

	return foundRegion
}

// RectQuery perform rectangular query ur upper right bl bottom left
func (gs *GeoFenceBoltDB) RectQuery(urlat, urlng, bllat, bllng float64, limit int) (region.Fences, error) {
	rect := s2.RectFromLatLng(s2.LatLngFromDegrees(bllat, bllng))
	rect = rect.AddPoint(s2.LatLngFromDegrees(urlat, urlng))

	rc := &s2.RegionCoverer{MaxLevel: 30, MaxCells: 1}
	covering := rc.Covering(rect)
	if len(covering) != 1 {
		return nil, errors.New("impossible covering")
	}
	i := &region.S2Interval{CellID: covering[0]}
	r := gs.Tree.Query(i)

	fences := make(map[uint64]*region.Fence)

	for _, itv := range r {
		sitv := itv.(*region.S2Interval)
		for _, loopID := range sitv.LoopIDs {
			var region *region.Fence
			if v, ok := fences[loopID]; ok {
				region = v
			} else {
				region = gs.FenceByID(loopID)
			}
			// testing the found loop is actually inside the rect
			// (since we are using only one large cover it may be outside)
			if rect.Contains(region.Loop.RectBound()) {
				fences[loopID] = region
			}
		}
	}

	var res []*region.Fence
	for _, v := range fences {
		res = append(res, v)
	}
	return region.Fences(res), nil
}

// StoreFence stores a fence into the database and load its index in memory
func (gs *GeoFenceBoltDB) StoreFence(fs *geostore.FenceStorage, cover []uint64) error {
	return gs.Update(func(tx *bolt.Tx) error {
		loopB := tx.Bucket([]byte(loopBucket))
		coverBucket := tx.Bucket([]byte(coverBucket))

		loopID, err := loopB.NextSequence()
		if err != nil {
			return err
		}

		buf, err := proto.Marshal(fs)
		if err != nil {
			return err
		}

		if gs.debug {
			log.Println("inserted", loopID, fs.Data, cover)
		}

		// convert our loopID to bigendian to be used as key
		k := itob(loopID)

		err = loopB.Put(k, buf)
		if err != nil {
			return err
		}

		// inserting into cover index using the same key
		fc := &geostore.FenceCover{Cellunion: cover}
		bufc, err := proto.Marshal(fc)
		if err != nil {
			return err
		}

		// also load into memory
		gs.index(fc, loopID)

		return coverBucket.Put(k, bufc)
	})
}

// itob returns an 8-byte big endian representation of v.
func itob(v uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	return b
}
<|endoftext|>