<commit_before>/*
Copyright 2017 Google Inc.

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

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

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
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 srvtopo

import (
	"bytes"
	"html/template"
	"sync"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/youtube/vitess/go/vt/status"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/topo/memorytopo"
	"golang.org/x/net/context"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
	vschemapb "github.com/youtube/vitess/go/vt/proto/vschema"
)

// TestGetSrvKeyspace will test we properly return updated SrvKeyspace.
func TestGetSrvKeyspace(t *testing.T) {
	ts := memorytopo.NewServer("test_cell")
	rs := NewResilientServer(ts, "TestGetSrvKeyspace")

	// Ask for a not-yet-created keyspace
	_, err := rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
	if err != topo.ErrNoNode {
		t.Fatalf("GetSrvKeyspace(not created) got unexpected error: %v", err)
	}

	// Set SrvKeyspace with value
	want := &topodatapb.SrvKeyspace{
		ShardingColumnName: "id",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
	}
	ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks", want)

	// wait until we get the right value
	var got *topodatapb.SrvKeyspace
	expiry := time.Now().Add(5 * time.Second)
	for {
		got, err = rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
		if err != nil {
			t.Fatalf("GetSrvKeyspace got unexpected error: %v", err)
		}
		if proto.Equal(want, got) {
			break
		}
		if time.Now().After(expiry) {
			t.Fatalf("GetSrvKeyspace() timeout = %+v, want %+v", got, want)
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Now delete the SrvKeyspace, wait until we get the error.
	if err := ts.DeleteSrvKeyspace(context.Background(), "test_cell", "test_ks"); err != nil {
		t.Fatalf("DeleteSrvKeyspace() failed: %v", err)
	}
	expiry = time.Now().Add(5 * time.Second)
	for {
		got, err = rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
		if err == topo.ErrNoNode {
			break
		}
		if time.Now().After(expiry) {
			t.Fatalf("timeout waiting for no keyspace error")
		}
		time.Sleep(time.Millisecond)
	}

	// Now send an updated real value, see it come through.
	want = &topodatapb.SrvKeyspace{
		ShardingColumnName: "id2",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
	}
	ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks", want)
	expiry = time.Now().Add(5 * time.Second)
	for {
		got, err = rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
		if err == nil && proto.Equal(want, got) {
			break
		}
		if time.Now().After(expiry) {
			t.Fatalf("timeout waiting for new keyspace value")
		}
		time.Sleep(time.Millisecond)
	}

	// make sure the HTML template works
	templ := template.New("").Funcs(status.StatusFuncs)
	templ, err = templ.Parse(TopoTemplate)
	if err != nil {
		t.Fatalf("error parsing template: %v", err)
	}
	wr := &bytes.Buffer{}
	if err := templ.Execute(wr, rs.CacheStatus()); err != nil {
		t.Fatalf("error executing template: %v", err)
	}
}

// TestSrvKeyspaceCachedError will test we properly re-try to query
// the topo server upon failure.
func TestSrvKeyspaceCachedError(t *testing.T) {
	ts := memorytopo.NewServer("test_cell")
	rs := NewResilientServer(ts, "TestSrvKeyspaceCachedErrors")

	// Ask for an unknown keyspace, should get an error.
	ctx := context.Background()
	_, err := rs.GetSrvKeyspace(ctx, "test_cell", "unknown_ks")
	if err == nil {
		t.Fatalf("First GetSrvKeyspace didn't return an error")
	}
	entry := rs.getSrvKeyspaceEntry("test_cell", "unknown_ks")
	if err != entry.lastError {
		t.Errorf("Error wasn't saved properly")
	}
	if ctx != entry.lastErrorCtx {
		t.Errorf("Context wasn't saved properly")
	}

	// Ask again with a different context, should get an error and
	// save that context.
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	_, err2 := rs.GetSrvKeyspace(ctx, "test_cell", "unknown_ks")
	if err2 == nil {
		t.Fatalf("Second GetSrvKeyspace didn't return an error")
	}
	if err2 != entry.lastError {
		t.Errorf("Error wasn't saved properly")
	}
	if ctx != entry.lastErrorCtx {
		t.Errorf("Context wasn't saved properly")
	}
}

// TestGetSrvKeyspaceCreated will test we properly get the initial
// value if the SrvKeyspace already exists.
func TestGetSrvKeyspaceCreated(t *testing.T) {
	ts := memorytopo.NewServer("test_cell")
	rs := NewResilientServer(ts, "TestGetSrvKeyspaceCreated")

	// Set SrvKeyspace with value.
	want := &topodatapb.SrvKeyspace{
		ShardingColumnName: "id",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
	}
	ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks", want)

	// Wait until we get the right value.
	expiry := time.Now().Add(5 * time.Second)
	for {
		got, err := rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
		switch err {
		case topo.ErrNoNode:
			// keep trying
		case nil:
			// we got a value, see if it's good
			if proto.Equal(want, got) {
				return
			}
		default:
			t.Fatalf("GetSrvKeyspace got unexpected error: %v", err)
		}
		if time.Now().After(expiry) {
			t.Fatalf("GetSrvKeyspace() timeout = %+v, want %+v", got, want)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestWatchSrvVSchema(t *testing.T) {
	ctx := context.Background()
	ts := memorytopo.NewServer("test_cell")
	rs := NewResilientServer(ts, "TestWatchSrvVSchema")
	watchSrvVSchemaSleepTime = 10 * time.Millisecond

	// mu protects watchValue and watchErr.
	mu := sync.Mutex{}
	var watchValue *vschemapb.SrvVSchema
	var watchErr error
	rs.WatchSrvVSchema(ctx, "test_cell", func(v *vschemapb.SrvVSchema, e error) {
		mu.Lock()
		defer mu.Unlock()
		watchValue = v
		watchErr = e
	})
	get := func() (*vschemapb.SrvVSchema, error) {
		mu.Lock()
		defer mu.Unlock()
		return watchValue, watchErr
	}

	// WatchSrvVSchema won't return until it gets the initial value,
	// which is not there, so we should get watchErr=topo.ErrNoNode.
	if _, err := get(); err != topo.ErrNoNode {
		t.Fatalf("WatchSrvVSchema didn't return topo.ErrNoNode at first, but got: %v", err)
	}

	// Save a value, wait for it.
	newValue := &vschemapb.SrvVSchema{
		Keyspaces: map[string]*vschemapb.Keyspace{
			"ks1": {},
		},
	}
	if err := ts.UpdateSrvVSchema(ctx, "test_cell", newValue); err != nil {
		t.Fatalf("UpdateSrvVSchema failed: %v", err)
	}
	start := time.Now()
	for {
		if v, err := get(); err == nil && proto.Equal(newValue, v) {
			break
		}
		if time.Since(start) > 5*time.Second {
			t.Fatalf("timed out waiting for new SrvVschema")
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Update value, wait for it.
	updatedValue := &vschemapb.SrvVSchema{
		Keyspaces: map[string]*vschemapb.Keyspace{
			"ks1": {},
		},
	}
	if err := ts.UpdateSrvVSchema(ctx, "test_cell", updatedValue); err != nil {
		t.Fatalf("UpdateSrvVSchema failed: %v", err)
	}
	start = time.Now()
	for {
		if v, err := get(); err == nil && proto.Equal(updatedValue, v) {
			break
		}
		if time.Since(start) > 5*time.Second {
			t.Fatalf("timed out waiting for updated SrvVschema")
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Delete the value, wait for topo.ErrNoNode
	if err := ts.DeleteSrvVSchema(ctx, "test_cell"); err != nil {
		t.Fatalf("DeleteSrvVSchema failed: %v", err)
	}
	start = time.Now()
	for {
		if _, err := get(); err == topo.ErrNoNode {
			break
		}
		if time.Since(start) > 5*time.Second {
			t.Fatalf("timed out waiting for deleted SrvVschema")
		}
		time.Sleep(10 * time.Millisecond)
	}
}
<commit_msg>Fixes test that asserts that changes to vschema are updated correctly by topo watcher<commit_after>/*
Copyright 2017 Google Inc.

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

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

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
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 srvtopo

import (
	"bytes"
	"html/template"
	"sync"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/youtube/vitess/go/vt/status"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/topo/memorytopo"
	"golang.org/x/net/context"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
	vschemapb "github.com/youtube/vitess/go/vt/proto/vschema"
)

// TestGetSrvKeyspace will test we properly return updated SrvKeyspace.
func TestGetSrvKeyspace(t *testing.T) {
	ts := memorytopo.NewServer("test_cell")
	rs := NewResilientServer(ts, "TestGetSrvKeyspace")

	// Ask for a not-yet-created keyspace
	_, err := rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
	if err != topo.ErrNoNode {
		t.Fatalf("GetSrvKeyspace(not created) got unexpected error: %v", err)
	}

	// Set SrvKeyspace with value
	want := &topodatapb.SrvKeyspace{
		ShardingColumnName: "id",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
	}
	ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks", want)

	// wait until we get the right value
	var got *topodatapb.SrvKeyspace
	expiry := time.Now().Add(5 * time.Second)
	for {
		got, err = rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
		if err != nil {
			t.Fatalf("GetSrvKeyspace got unexpected error: %v", err)
		}
		if proto.Equal(want, got) {
			break
		}
		if time.Now().After(expiry) {
			t.Fatalf("GetSrvKeyspace() timeout = %+v, want %+v", got, want)
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Now delete the SrvKeyspace, wait until we get the error.
	if err := ts.DeleteSrvKeyspace(context.Background(), "test_cell", "test_ks"); err != nil {
		t.Fatalf("DeleteSrvKeyspace() failed: %v", err)
	}
	expiry = time.Now().Add(5 * time.Second)
	for {
		got, err = rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
		if err == topo.ErrNoNode {
			break
		}
		if time.Now().After(expiry) {
			t.Fatalf("timeout waiting for no keyspace error")
		}
		time.Sleep(time.Millisecond)
	}

	// Now send an updated real value, see it come through.
	want = &topodatapb.SrvKeyspace{
		ShardingColumnName: "id2",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
	}
	ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks", want)
	expiry = time.Now().Add(5 * time.Second)
	for {
		got, err = rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
		if err == nil && proto.Equal(want, got) {
			break
		}
		if time.Now().After(expiry) {
			t.Fatalf("timeout waiting for new keyspace value")
		}
		time.Sleep(time.Millisecond)
	}

	// make sure the HTML template works
	templ := template.New("").Funcs(status.StatusFuncs)
	templ, err = templ.Parse(TopoTemplate)
	if err != nil {
		t.Fatalf("error parsing template: %v", err)
	}
	wr := &bytes.Buffer{}
	if err := templ.Execute(wr, rs.CacheStatus()); err != nil {
		t.Fatalf("error executing template: %v", err)
	}
}

// TestSrvKeyspaceCachedError will test we properly re-try to query
// the topo server upon failure.
func TestSrvKeyspaceCachedError(t *testing.T) {
	ts := memorytopo.NewServer("test_cell")
	rs := NewResilientServer(ts, "TestSrvKeyspaceCachedErrors")

	// Ask for an unknown keyspace, should get an error.
	ctx := context.Background()
	_, err := rs.GetSrvKeyspace(ctx, "test_cell", "unknown_ks")
	if err == nil {
		t.Fatalf("First GetSrvKeyspace didn't return an error")
	}
	entry := rs.getSrvKeyspaceEntry("test_cell", "unknown_ks")
	if err != entry.lastError {
		t.Errorf("Error wasn't saved properly")
	}
	if ctx != entry.lastErrorCtx {
		t.Errorf("Context wasn't saved properly")
	}

	// Ask again with a different context, should get an error and
	// save that context.
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	_, err2 := rs.GetSrvKeyspace(ctx, "test_cell", "unknown_ks")
	if err2 == nil {
		t.Fatalf("Second GetSrvKeyspace didn't return an error")
	}
	if err2 != entry.lastError {
		t.Errorf("Error wasn't saved properly")
	}
	if ctx != entry.lastErrorCtx {
		t.Errorf("Context wasn't saved properly")
	}
}

// TestGetSrvKeyspaceCreated will test we properly get the initial
// value if the SrvKeyspace already exists.
func TestGetSrvKeyspaceCreated(t *testing.T) {
	ts := memorytopo.NewServer("test_cell")
	rs := NewResilientServer(ts, "TestGetSrvKeyspaceCreated")

	// Set SrvKeyspace with value.
	want := &topodatapb.SrvKeyspace{
		ShardingColumnName: "id",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
	}
	ts.UpdateSrvKeyspace(context.Background(), "test_cell", "test_ks", want)

	// Wait until we get the right value.
	expiry := time.Now().Add(5 * time.Second)
	for {
		got, err := rs.GetSrvKeyspace(context.Background(), "test_cell", "test_ks")
		switch err {
		case topo.ErrNoNode:
			// keep trying
		case nil:
			// we got a value, see if it's good
			if proto.Equal(want, got) {
				return
			}
		default:
			t.Fatalf("GetSrvKeyspace got unexpected error: %v", err)
		}
		if time.Now().After(expiry) {
			t.Fatalf("GetSrvKeyspace() timeout = %+v, want %+v", got, want)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestWatchSrvVSchema(t *testing.T) {
	ctx := context.Background()
	ts := memorytopo.NewServer("test_cell")
	rs := NewResilientServer(ts, "TestWatchSrvVSchema")
	watchSrvVSchemaSleepTime = 10 * time.Millisecond

	// mu protects watchValue and watchErr.
	mu := sync.Mutex{}
	var watchValue *vschemapb.SrvVSchema
	var watchErr error
	rs.WatchSrvVSchema(ctx, "test_cell", func(v *vschemapb.SrvVSchema, e error) {
		mu.Lock()
		defer mu.Unlock()
		watchValue = v
		watchErr = e
	})
	get := func() (*vschemapb.SrvVSchema, error) {
		mu.Lock()
		defer mu.Unlock()
		return watchValue, watchErr
	}

	// WatchSrvVSchema won't return until it gets the initial value,
	// which is not there, so we should get watchErr=topo.ErrNoNode.
	if _, err := get(); err != topo.ErrNoNode {
		t.Fatalf("WatchSrvVSchema didn't return topo.ErrNoNode at first, but got: %v", err)
	}

	// Save a value, wait for it.
	newValue := &vschemapb.SrvVSchema{
		Keyspaces: map[string]*vschemapb.Keyspace{
			"ks1": {},
		},
	}
	if err := ts.UpdateSrvVSchema(ctx, "test_cell", newValue); err != nil {
		t.Fatalf("UpdateSrvVSchema failed: %v", err)
	}
	start := time.Now()
	for {
		if v, err := get(); err == nil && proto.Equal(newValue, v) {
			break
		}
		if time.Since(start) > 5*time.Second {
			t.Fatalf("timed out waiting for new SrvVschema")
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Update value, wait for it.
	updatedValue := &vschemapb.SrvVSchema{
		Keyspaces: map[string]*vschemapb.Keyspace{
			"ks2": {},
		},
	}
	if err := ts.UpdateSrvVSchema(ctx, "test_cell", updatedValue); err != nil {
		t.Fatalf("UpdateSrvVSchema failed: %v", err)
	}
	start = time.Now()
	for {
		if v, err := get(); err == nil && proto.Equal(updatedValue, v) {
			break
		}
		if time.Since(start) > 5*time.Second {
			t.Fatalf("timed out waiting for updated SrvVschema")
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Delete the value, wait for topo.ErrNoNode
	if err := ts.DeleteSrvVSchema(ctx, "test_cell"); err != nil {
		t.Fatalf("DeleteSrvVSchema failed: %v", err)
	}
	start = time.Now()
	for {
		if _, err := get(); err == topo.ErrNoNode {
			break
		}
		if time.Since(start) > 5*time.Second {
			t.Fatalf("timed out waiting for deleted SrvVschema")
		}
		time.Sleep(10 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>package lfs

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

	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

func ClearTempObjects() {
	d, err := os.Open(LocalObjectTempDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error opening %q to clear old temp files: %s", LocalObjectTempDir, err)
		return
	}

	filenames, _ := d.Readdirnames(-1)
	for _, filename := range filenames {
		path := filepath.Join(LocalObjectTempDir, filename)
		if shouldDeleteTempObject(path) {
			os.RemoveAll(path)
		}
	}
}

func shouldDeleteTempObject(path string) bool {
	info, err := os.Stat(path)
	if err != nil {
		return false
	}

	if info.IsDir() {
		return false
	}

	base := filepath.Base(path)
	parts := strings.SplitN(base, "-", 2)
	oid := parts[0]
	if len(parts) < 2 || len(oid) != 64 {
		tracerx.Printf("Removing invalid tmp object file: %s", path)
		return true
	}

	if FileExists(localMediaPathNoCreate(oid)) {
		tracerx.Printf("Removing existing tmp object file: %s", path)
		return true
	}

	if time.Since(info.ModTime()) > time.Hour {
		tracerx.Printf("Removing old tmp object file: %s", path)
		return true
	}

	return false
}
<commit_msg>ンンン ンンン ンン<commit_after>package lfs

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

	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

func ClearTempObjects() {
	if len(LocalObjectTempDir) == 0 {
		return
	}

	d, err := os.Open(LocalObjectTempDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error opening %q to clear old temp files: %s\n", LocalObjectTempDir, err)
		return
	}

	filenames, _ := d.Readdirnames(-1)
	for _, filename := range filenames {
		path := filepath.Join(LocalObjectTempDir, filename)
		if shouldDeleteTempObject(path) {
			os.RemoveAll(path)
		}
	}
}

func shouldDeleteTempObject(path string) bool {
	info, err := os.Stat(path)
	if err != nil {
		return false
	}

	if info.IsDir() {
		return false
	}

	base := filepath.Base(path)
	parts := strings.SplitN(base, "-", 2)
	oid := parts[0]
	if len(parts) < 2 || len(oid) != 64 {
		tracerx.Printf("Removing invalid tmp object file: %s", path)
		return true
	}

	if FileExists(localMediaPathNoCreate(oid)) {
		tracerx.Printf("Removing existing tmp object file: %s", path)
		return true
	}

	if time.Since(info.ModTime()) > time.Hour {
		tracerx.Printf("Removing old tmp object file: %s", path)
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package lfs

import (
	"fmt"
	"os"
	"path/filepath"
)

// Uploadable describes a file that can be uploaded.
type Uploadable struct {
	oid      string
	OidPath  string
	Filename string
	size     int64
	object   *objectResource
}

// NewUploadable builds the Uploadable from the given information.
func NewUploadable(oid, filename string, index, totalFiles int) (*Uploadable, *WrappedError) {
	path, err := LocalMediaPath(oid)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	if err := ensureFile(filename, path); err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	fi, err := os.Stat(filename)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	return &Uploadable{oid: oid, OidPath: path, Filename: filename, size: fi.Size()}, nil
}

func (u *Uploadable) Check() (*objectResource, *WrappedError) {
	return UploadCheck(u.OidPath)
}

func (u *Uploadable) Transfer(cb CopyCallback) *WrappedError {
	wcb := func(total, read int64, current int) error {
		cb(total, read, current)
		return nil
	}

	return UploadObject(u.object, wcb)
}

func (u *Uploadable) Object() *objectResource {
	return u.object
}

func (u *Uploadable) Oid() string {
	return u.oid
}

func (u *Uploadable) Size() int64 {
	return u.size
}

func (u *Uploadable) Name() string {
	return u.Filename
}

func (u *Uploadable) SetObject(o *objectResource) {
	u.object = o
}

// NewUploadQueue builds an UploadQueue, allowing `workers` concurrent uploads.
func NewUploadQueue(workers, files int) *TransferQueue {
	q := newTransferQueue(workers, files)
	q.transferKind = "upload"
	return q
}

// ensureFile makes sure that the cleanPath exists before pushing it.  If it
// does not exist, it attempts to clean it by reading the file at smudgePath.
func ensureFile(smudgePath, cleanPath string) error {
	if _, err := os.Stat(cleanPath); err == nil {
		return nil
	}

	expectedOid := filepath.Base(cleanPath)
	localPath := filepath.Join(LocalWorkingDir, smudgePath)
	file, err := os.Open(localPath)
	if err != nil {
		return err
	}

	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return err
	}

	cleaned, err := PointerClean(file, stat.Size(), nil)
	if err != nil {
		return err
	}

	cleaned.Close()

	if expectedOid != cleaned.Oid {
		return fmt.Errorf("Expected %s to have an OID of %s, got %s", smudgePath, expectedOid, cleaned.Oid)
	}

	return nil
}
<commit_msg>we just need the file size. the file doesn't need to be in the working dir<commit_after>package lfs

import (
	"fmt"
	"os"
	"path/filepath"
)

// Uploadable describes a file that can be uploaded.
type Uploadable struct {
	oid      string
	OidPath  string
	Filename string
	size     int64
	object   *objectResource
}

// NewUploadable builds the Uploadable from the given information.
func NewUploadable(oid, filename string, index, totalFiles int) (*Uploadable, *WrappedError) {
	path, err := LocalMediaPath(oid)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	if err := ensureFile(filename, path); err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	fi, err := os.Stat(path)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	return &Uploadable{oid: oid, OidPath: path, Filename: filename, size: fi.Size()}, nil
}

func (u *Uploadable) Check() (*objectResource, *WrappedError) {
	return UploadCheck(u.OidPath)
}

func (u *Uploadable) Transfer(cb CopyCallback) *WrappedError {
	wcb := func(total, read int64, current int) error {
		cb(total, read, current)
		return nil
	}

	return UploadObject(u.object, wcb)
}

func (u *Uploadable) Object() *objectResource {
	return u.object
}

func (u *Uploadable) Oid() string {
	return u.oid
}

func (u *Uploadable) Size() int64 {
	return u.size
}

func (u *Uploadable) Name() string {
	return u.Filename
}

func (u *Uploadable) SetObject(o *objectResource) {
	u.object = o
}

// NewUploadQueue builds an UploadQueue, allowing `workers` concurrent uploads.
func NewUploadQueue(workers, files int) *TransferQueue {
	q := newTransferQueue(workers, files)
	q.transferKind = "upload"
	return q
}

// ensureFile makes sure that the cleanPath exists before pushing it.  If it
// does not exist, it attempts to clean it by reading the file at smudgePath.
func ensureFile(smudgePath, cleanPath string) error {
	if _, err := os.Stat(cleanPath); err == nil {
		return nil
	}

	expectedOid := filepath.Base(cleanPath)
	localPath := filepath.Join(LocalWorkingDir, smudgePath)
	file, err := os.Open(localPath)
	if err != nil {
		return err
	}

	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return err
	}

	cleaned, err := PointerClean(file, stat.Size(), nil)
	if err != nil {
		return err
	}

	cleaned.Close()

	if expectedOid != cleaned.Oid {
		return fmt.Errorf("Expected %s to have an OID of %s, got %s", smudgePath, expectedOid, cleaned.Oid)
	}

	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017 Red Hat, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy ofthe 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 specificlanguage governing permissions and
 * limitations under the License.
 *
 */

package common

import (
	"fmt"
	"net/http"
	"strings"
	"sync"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/graffiti/graph"
	"github.com/skydive-project/skydive/graffiti/graph/traversal"
	gws "github.com/skydive-project/skydive/graffiti/websocket"
	"github.com/skydive-project/skydive/logging"
	ws "github.com/skydive-project/skydive/websocket"
)

type subscriber struct {
	graph         *graph.Graph
	gremlinFilter string
	ts            *traversal.GremlinTraversalSequence
}

// SubscriberEndpoint sends all the modifications to its subscribers.
type SubscriberEndpoint struct {
	common.RWMutex
	ws.DefaultSpeakerEventHandler
	pool          ws.StructSpeakerPool
	Graph         *graph.Graph
	wg            sync.WaitGroup
	gremlinParser *traversal.GremlinTraversalParser
	subscribers   map[ws.Speaker]*subscriber
}

func (t *SubscriberEndpoint) getGraph(gremlinQuery string, ts *traversal.GremlinTraversalSequence, lockGraph bool) (*graph.Graph, error) {
	res, err := ts.Exec(t.Graph, lockGraph)
	if err != nil {
		return nil, err
	}

	tv, ok := res.(*traversal.GraphTraversal)
	if !ok {
		return nil, fmt.Errorf("Gremlin query '%s' did not return a graph", gremlinQuery)
	}

	return tv.Graph, nil
}

func (t *SubscriberEndpoint) newSubscriber(host string, gremlinFilter string, lockGraph bool) (*subscriber, error) {
	ts, err := t.gremlinParser.Parse(strings.NewReader(gremlinFilter))
	if err != nil {
		return nil, fmt.Errorf("Invalid Gremlin filter '%s' for client %s", gremlinFilter, host)
	}

	g, err := t.getGraph(gremlinFilter, ts, lockGraph)
	if err != nil {
		return nil, err
	}

	return &subscriber{graph: g, ts: ts, gremlinFilter: gremlinFilter}, nil
}

// OnConnected called when a subscriber got connected.
func (t *SubscriberEndpoint) OnConnected(c ws.Speaker) {
	gremlinFilter := c.GetHeaders().Get("X-Gremlin-Filter")
	if gremlinFilter == "" {
		gremlinFilter = c.GetURL().Query().Get("x-gremlin-filter")
	}

	if gremlinFilter != "" {
		host := c.GetRemoteHost()

		subscriber, err := t.newSubscriber(host, gremlinFilter, false)
		if err != nil {
			logging.GetLogger().Error(err)
			return
		}

		logging.GetLogger().Infof("Client %s subscribed with filter %s during the connection", host, gremlinFilter)
		t.subscribers[c] = subscriber
	}
}

// OnDisconnected called when a subscriber got disconnected.
func (t *SubscriberEndpoint) OnDisconnected(c ws.Speaker) {
	t.Lock()
	delete(t.subscribers, c)
	t.Unlock()
}

// OnStructMessage is triggered when receiving a message from a subscriber.
// It only responds to SyncRequestMsgType messages
func (t *SubscriberEndpoint) OnStructMessage(c ws.Speaker, msg *ws.StructMessage) {
	msgType, obj, err := gws.UnmarshalMessage(msg)
	if err != nil {
		logging.GetLogger().Errorf("Graph: Unable to parse the event %v: %s", msg, err)
		return
	}

	// this kind of message usually comes from external clients like the WebUI
	if msgType == gws.SyncRequestMsgType {
		t.Graph.RLock()
		defer t.Graph.RUnlock()

		syncMsg, status := obj.(*gws.SyncRequestMsg), http.StatusOK
		result, err := t.Graph.CloneWithContext(syncMsg.Context)
		if err != nil {
			logging.GetLogger().Errorf("unable to get a graph with context %+v: %s", syncMsg, err)
			reply := msg.Reply(nil, gws.SyncReplyMsgType, http.StatusBadRequest)
			c.SendMessage(reply)
			return
		}

		host := c.GetRemoteHost()

		if syncMsg.GremlinFilter != nil {
			// filter reset
			if *syncMsg.GremlinFilter == "" {
				t.Lock()
				delete(t.subscribers, c)
				t.Unlock()
			} else {
				subscriber, err := t.newSubscriber(host, *syncMsg.GremlinFilter, false)
				if err != nil {
					logging.GetLogger().Error(err)

					reply := msg.Reply(err.Error(), gws.SyncReplyMsgType, http.StatusBadRequest)
					c.SendMessage(reply)

					t.Lock()
					t.subscribers[c] = nil
					t.Unlock()

					return
				}

				logging.GetLogger().Infof("Client %s requested subscription with filter %s", host, *syncMsg.GremlinFilter)
				result = subscriber.graph

				t.Lock()
				t.subscribers[c] = subscriber
				t.Unlock()
			}
		} else {
			t.RLock()
			subscriber := t.subscribers[c]
			t.RUnlock()

			if subscriber != nil {
				result = subscriber.graph
			}
		}

		reply := msg.Reply(result, gws.SyncReplyMsgType, status)
		c.SendMessage(reply)

		return
	}
}

// notifyClients forwards local graph modification to subscribers. If a subscriber
// specified a Gremlin filter, a 'Diff' is applied between the previous graph state
// for this subscriber and the current graph state.
func (t *SubscriberEndpoint) notifyClients(msg *ws.StructMessage) {
	for _, c := range t.pool.GetSpeakers() {
		t.RLock()
		subscriber, found := t.subscribers[c]
		t.RUnlock()

		if found {
			// in the case of a error during the subscription we got a nil subscriber
			if subscriber == nil {
				return
			}

			g, err := t.getGraph(subscriber.gremlinFilter, subscriber.ts, false)
			if err != nil {
				logging.GetLogger().Error(err)
				continue
			}

			addedNodes, removedNodes, addedEdges, removedEdges := subscriber.graph.Diff(g)

			for _, n := range addedNodes {
				c.SendMessage(gws.NewStructMessage(gws.NodeAddedMsgType, n))
			}

			for _, n := range removedNodes {
				c.SendMessage(gws.NewStructMessage(gws.NodeDeletedMsgType, n))
			}

			for _, e := range addedEdges {
				c.SendMessage(gws.NewStructMessage(gws.EdgeAddedMsgType, e))
			}

			for _, e := range removedEdges {
				c.SendMessage(gws.NewStructMessage(gws.EdgeDeletedMsgType, e))
			}

			subscriber.graph = g
		} else {
			c.SendMessage(msg)
		}
	}
}

// OnNodeUpdated graph node updated event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnNodeUpdated(n *graph.Node) {
	t.notifyClients(gws.NewStructMessage(gws.NodeUpdatedMsgType, n))
}

// OnNodeAdded graph node added event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnNodeAdded(n *graph.Node) {
	t.notifyClients(gws.NewStructMessage(gws.NodeAddedMsgType, n))
}

// OnNodeDeleted graph node deleted event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnNodeDeleted(n *graph.Node) {
	t.notifyClients(gws.NewStructMessage(gws.NodeDeletedMsgType, n))
}

// OnEdgeUpdated graph edge updated event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnEdgeUpdated(e *graph.Edge) {
	t.notifyClients(gws.NewStructMessage(gws.EdgeUpdatedMsgType, e))
}

// OnEdgeAdded graph edge added event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnEdgeAdded(e *graph.Edge) {
	t.notifyClients(gws.NewStructMessage(gws.EdgeAddedMsgType, e))
}

// OnEdgeDeleted graph edge deleted event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnEdgeDeleted(e *graph.Edge) {
	t.notifyClients(gws.NewStructMessage(gws.EdgeDeletedMsgType, e))
}

// NewSubscriberEndpoint returns a new server to be used by external subscribers,
// for instance the WebUI.
func NewSubscriberEndpoint(pool ws.StructSpeakerPool, g *graph.Graph, tr *traversal.GremlinTraversalParser) *SubscriberEndpoint {
	t := &SubscriberEndpoint{
		Graph:         g,
		pool:          pool,
		subscribers:   make(map[ws.Speaker]*subscriber),
		gremlinParser: tr,
	}

	pool.AddEventHandler(t)

	// subscribe to the graph messages
	pool.AddStructMessageHandler(t, []string{gws.Namespace})

	// subscribe to the local graph event
	g.AddEventListener(t)
	return t
}
<commit_msg>graffiti: handle updates in gremlin filters<commit_after>/*
 * Copyright (C) 2017 Red Hat, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy ofthe 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 specificlanguage governing permissions and
 * limitations under the License.
 *
 */

package common

import (
	"fmt"
	"net/http"
	"strings"
	"sync"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/graffiti/graph"
	"github.com/skydive-project/skydive/graffiti/graph/traversal"
	gws "github.com/skydive-project/skydive/graffiti/websocket"
	"github.com/skydive-project/skydive/logging"
	ws "github.com/skydive-project/skydive/websocket"
)

type subscriber struct {
	graph         *graph.Graph
	gremlinFilter string
	ts            *traversal.GremlinTraversalSequence
}

// SubscriberEndpoint sends all the modifications to its subscribers.
type SubscriberEndpoint struct {
	common.RWMutex
	ws.DefaultSpeakerEventHandler
	pool          ws.StructSpeakerPool
	Graph         *graph.Graph
	wg            sync.WaitGroup
	gremlinParser *traversal.GremlinTraversalParser
	subscribers   map[ws.Speaker]*subscriber
}

func (t *SubscriberEndpoint) getGraph(gremlinQuery string, ts *traversal.GremlinTraversalSequence, lockGraph bool) (*graph.Graph, error) {
	res, err := ts.Exec(t.Graph, lockGraph)
	if err != nil {
		return nil, err
	}

	tv, ok := res.(*traversal.GraphTraversal)
	if !ok {
		return nil, fmt.Errorf("Gremlin query '%s' did not return a graph", gremlinQuery)
	}

	return tv.Graph, nil
}

func (t *SubscriberEndpoint) newSubscriber(host string, gremlinFilter string, lockGraph bool) (*subscriber, error) {
	ts, err := t.gremlinParser.Parse(strings.NewReader(gremlinFilter))
	if err != nil {
		return nil, fmt.Errorf("Invalid Gremlin filter '%s' for client %s", gremlinFilter, host)
	}

	g, err := t.getGraph(gremlinFilter, ts, lockGraph)
	if err != nil {
		return nil, err
	}

	return &subscriber{graph: g, ts: ts, gremlinFilter: gremlinFilter}, nil
}

// OnConnected called when a subscriber got connected.
func (t *SubscriberEndpoint) OnConnected(c ws.Speaker) {
	gremlinFilter := c.GetHeaders().Get("X-Gremlin-Filter")
	if gremlinFilter == "" {
		gremlinFilter = c.GetURL().Query().Get("x-gremlin-filter")
	}

	if gremlinFilter != "" {
		host := c.GetRemoteHost()

		subscriber, err := t.newSubscriber(host, gremlinFilter, false)
		if err != nil {
			logging.GetLogger().Error(err)
			return
		}

		logging.GetLogger().Infof("Client %s subscribed with filter %s during the connection", host, gremlinFilter)
		t.subscribers[c] = subscriber
	}
}

// OnDisconnected called when a subscriber got disconnected.
func (t *SubscriberEndpoint) OnDisconnected(c ws.Speaker) {
	t.Lock()
	delete(t.subscribers, c)
	t.Unlock()
}

// OnStructMessage is triggered when receiving a message from a subscriber.
// It only responds to SyncRequestMsgType messages
func (t *SubscriberEndpoint) OnStructMessage(c ws.Speaker, msg *ws.StructMessage) {
	msgType, obj, err := gws.UnmarshalMessage(msg)
	if err != nil {
		logging.GetLogger().Errorf("Graph: Unable to parse the event %v: %s", msg, err)
		return
	}

	// this kind of message usually comes from external clients like the WebUI
	if msgType == gws.SyncRequestMsgType {
		t.Graph.RLock()
		defer t.Graph.RUnlock()

		syncMsg, status := obj.(*gws.SyncRequestMsg), http.StatusOK
		result, err := t.Graph.CloneWithContext(syncMsg.Context)
		if err != nil {
			logging.GetLogger().Errorf("unable to get a graph with context %+v: %s", syncMsg, err)
			reply := msg.Reply(nil, gws.SyncReplyMsgType, http.StatusBadRequest)
			c.SendMessage(reply)
			return
		}

		host := c.GetRemoteHost()

		if syncMsg.GremlinFilter != nil {
			// filter reset
			if *syncMsg.GremlinFilter == "" {
				t.Lock()
				delete(t.subscribers, c)
				t.Unlock()
			} else {
				subscriber, err := t.newSubscriber(host, *syncMsg.GremlinFilter, false)
				if err != nil {
					logging.GetLogger().Error(err)

					reply := msg.Reply(err.Error(), gws.SyncReplyMsgType, http.StatusBadRequest)
					c.SendMessage(reply)

					t.Lock()
					t.subscribers[c] = nil
					t.Unlock()

					return
				}

				logging.GetLogger().Infof("Client %s requested subscription with filter %s", host, *syncMsg.GremlinFilter)
				result = subscriber.graph

				t.Lock()
				t.subscribers[c] = subscriber
				t.Unlock()
			}
		} else {
			t.RLock()
			subscriber := t.subscribers[c]
			t.RUnlock()

			if subscriber != nil {
				result = subscriber.graph
			}
		}

		reply := msg.Reply(result, gws.SyncReplyMsgType, status)
		c.SendMessage(reply)

		return
	}
}

// notifyClients forwards local graph modification to subscribers. If a subscriber
// specified a Gremlin filter, a 'Diff' is applied between the previous graph state
// for this subscriber and the current graph state.
func (t *SubscriberEndpoint) notifyClients(typ string, i interface{}) {
	for _, c := range t.pool.GetSpeakers() {
		t.RLock()
		subscriber, found := t.subscribers[c]
		t.RUnlock()

		if found {
			// in the case of a error during the subscription we got a nil subscriber
			if subscriber == nil {
				return
			}

			g, err := t.getGraph(subscriber.gremlinFilter, subscriber.ts, false)
			if err != nil {
				logging.GetLogger().Error(err)
				continue
			}

			addedNodes, removedNodes, addedEdges, removedEdges := subscriber.graph.Diff(g)

			for _, n := range addedNodes {
				c.SendMessage(gws.NewStructMessage(gws.NodeAddedMsgType, n))
			}

			for _, n := range removedNodes {
				c.SendMessage(gws.NewStructMessage(gws.NodeDeletedMsgType, n))
			}

			for _, e := range addedEdges {
				c.SendMessage(gws.NewStructMessage(gws.EdgeAddedMsgType, e))
			}

			for _, e := range removedEdges {
				c.SendMessage(gws.NewStructMessage(gws.EdgeDeletedMsgType, e))
			}

			// handle updates
			switch typ {
			case gws.NodeUpdatedMsgType:
				if g.GetNode(i.(*graph.Node).ID) != nil {
					c.SendMessage(gws.NewStructMessage(gws.NodeUpdatedMsgType, i))
				}
			case gws.EdgeUpdatedMsgType:
				if g.GetEdge(i.(*graph.Edge).ID) != nil {
					c.SendMessage(gws.NewStructMessage(gws.EdgeUpdatedMsgType, i))
				}
			}

			subscriber.graph = g
		} else {
			c.SendMessage(gws.NewStructMessage(typ, i))
		}
	}
}

// OnNodeUpdated graph node updated event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnNodeUpdated(n *graph.Node) {
	t.notifyClients(gws.NodeUpdatedMsgType, n)
}

// OnNodeAdded graph node added event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnNodeAdded(n *graph.Node) {
	t.notifyClients(gws.NodeAddedMsgType, n)
}

// OnNodeDeleted graph node deleted event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnNodeDeleted(n *graph.Node) {
	t.notifyClients(gws.NodeDeletedMsgType, n)
}

// OnEdgeUpdated graph edge updated event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnEdgeUpdated(e *graph.Edge) {
	t.notifyClients(gws.EdgeUpdatedMsgType, e)
}

// OnEdgeAdded graph edge added event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnEdgeAdded(e *graph.Edge) {
	t.notifyClients(gws.EdgeAddedMsgType, e)
}

// OnEdgeDeleted graph edge deleted event. Implements the GraphEventListener interface.
func (t *SubscriberEndpoint) OnEdgeDeleted(e *graph.Edge) {
	t.notifyClients(gws.EdgeDeletedMsgType, e)
}

// NewSubscriberEndpoint returns a new server to be used by external subscribers,
// for instance the WebUI.
func NewSubscriberEndpoint(pool ws.StructSpeakerPool, g *graph.Graph, tr *traversal.GremlinTraversalParser) *SubscriberEndpoint {
	t := &SubscriberEndpoint{
		Graph:         g,
		pool:          pool,
		subscribers:   make(map[ws.Speaker]*subscriber),
		gremlinParser: tr,
	}

	pool.AddEventHandler(t)

	// subscribe to the graph messages
	pool.AddStructMessageHandler(t, []string{gws.Namespace})

	// subscribe to the local graph event
	g.AddEventListener(t)
	return t
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 probe

import (
	"net/http"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/unrolled/render"
)

const (
	defaultPluginName string = "PROMETHEUS"

	// DefaultMetricsPath default Prometheus metrics URL
	DefaultMetricsPath string = "/metrics"

	// Namespace namespace to use for Prometheus metrics
	Namespace string = ""
	// Subsystem subsystem to use for Prometheus metrics
	Subsystem string = ""
	// ServiceLabel label for service field
	ServiceLabel string = "service"
	// DependencyLabel label for dependency field
	DependencyLabel string = "dependency"
	// BuildVersionLabel label for build version field
	BuildVersionLabel string = "build_version"
	// BuildDateLabel label for build date field
	BuildDateLabel string = "build_date"

	// ServiceHealthName name of service health metric
	ServiceHealthName string = "service_health"

	// ServiceHealthHelp help text for service health metric
	// Adapt Ligato status code for now.
	// TODO: Consolidate with that from the "Common Container Telemetry" proposal.
	// ServiceHealthHelp    string = "The health of the ServiceLabel 0 = INIT, 1 = UP, 2 = DOWN, 3 = OUTAGE"
	ServiceHealthHelp string = "The health of the ServiceLabel 0 = INIT, 1 = OK, 2 = ERROR"

	// DependencyHealthName name of dependency health metric
	DependencyHealthName string = "service_dependency_health"

	// DependencyHealthHelp help text for dependency health metric
	// Adapt Ligato status code for now.
	// TODO: Consolidate with that from the "Common Container Telemetry" proposal.
	// DependencyHealthHelp string = "The health of the DependencyLabel 0 = INIT, 1 = UP, 2 = DOWN, 3 = OUTAGE"
	DependencyHealthHelp string = "The health of the DependencyLabel 0 = INIT, 1 = OK, 2 = ERROR"

	// ServiceInfoName name of service info metric
	ServiceInfoName string = "service_info"
	// ServiceInfoHelp help text for service info metric
	ServiceInfoHelp string = "Build info for the service.  Value is always 1, build info is in the tags."
)

// PrometheusPlugin struct holds all plugin-related data.
type PrometheusPlugin struct {
	Deps
}

// Init may create a new (custom) instance of HTTP if the injected instance uses
// different HTTP port than requested.
func (p *PrometheusPlugin) Init() (err error) {
	serviceLabel := p.String()
	if p.Deps.ServiceLabel != nil {
		serviceLabel = p.Deps.ServiceLabel.GetAgentLabel()
	}

	p.registerGauge(
		Namespace,
		Subsystem,
		ServiceHealthName,
		ServiceHealthHelp,
		prometheus.Labels{ServiceLabel: serviceLabel},
		p.getServiceHealth,
	)

	agentStatus := p.StatusCheck.GetAgentStatus()
	p.registerGauge(
		Namespace,
		Subsystem,
		ServiceInfoName,
		ServiceInfoHelp,
		prometheus.Labels{
			ServiceLabel:      serviceLabel,
			BuildVersionLabel: agentStatus.BuildVersion,
			BuildDateLabel:    agentStatus.BuildDate},
		func() float64 { return 1 },
	)

	return nil
}

// AfterInit registers HTTP handlers.
func (p *PrometheusPlugin) AfterInit() error {
	if p.HTTP != nil {
		if p.StatusCheck != nil {
			p.Log.Info("Starting Prometheus metrics handlers")
			p.HTTP.RegisterHTTPHandler(DefaultMetricsPath, p.metricsHandler, "GET")
		} else {
			p.Log.Info("Unable to register Prometheus metrics handlers, StatusCheck is nil")
		}
	} else {
		p.Log.Info("Unable to register Prometheus metrics handlers, HTTP is nil")
	}

	return nil
}

// Close shutdowns HTTP if a custom instance was created in Init().
func (p *PrometheusPlugin) Close() error {
	return nil
}

// metricsHandler handles Prometheus metrics collection.
func (p *PrometheusPlugin) metricsHandler(formatter *render.Render) http.HandlerFunc {
	return promhttp.Handler().ServeHTTP
}

func (p *PrometheusPlugin) getServiceHealth() float64 {
	agentStatus := p.StatusCheck.GetAgentStatus()
	// Adapt Ligato status code for now.
	// TODO: Consolidate with that from the "Common Container Telemetry" proposal.
	health := float64(agentStatus.State)
	p.Log.Infof("getServiceHealth(): %f", health)
	return health
}

// RegisterGauge registers custom gauge with specific valueFunc to report status when invoked.
func (p *PrometheusPlugin) registerGauge(namespace string, subsystem string, name string, help string,
	labels prometheus.Labels, valueFunc func() float64) {
	gaugeName := name
	if subsystem != "" {
		gaugeName = subsystem + "_" + gaugeName
	}
	if namespace != "" {
		gaugeName = namespace + "_" + gaugeName
	}
	if err := prometheus.DefaultRegisterer.Register(prometheus.NewGaugeFunc(
		prometheus.GaugeOpts{
			// Namespace, Subsystem, and Name are components of the fully-qualified
			// name of the Metric (created by joining these components with
			// "_"). Only Name is mandatory, the others merely help structuring the
			// name. Note that the fully-qualified name of the metric must be a
			// valid Prometheus metric name.
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      name,

			// Help provides information about this metric. Mandatory!
			//
			// Metrics with the same fully-qualified name must have the same Help
			// string.
			Help: help,

			// ConstLabels are used to attach fixed labels to this metric. Metrics
			// with the same fully-qualified name must have the same label names in
			// their ConstLabels.
			//
			// Note that in most cases, labels have a value that varies during the
			// lifetime of a process. Those labels are usually managed with a metric
			// vector collector (like CounterVec, GaugeVec, UntypedVec). ConstLabels
			// serve only special purposes. One is for the special case where the
			// value of a label does not change during the lifetime of a process,
			// e.g. if the revision of the running binary is put into a
			// label. Another, more advanced purpose is if more than one Collector
			// needs to collect Metrics with the same fully-qualified name. In that
			// case, those Metrics must differ in the values of their
			// ConstLabels. See the Collector examples.
			//
			// If the value of a label never changes (not even between binaries),
			// that label most likely should not be a label at all (but part of the
			// metric name).
			ConstLabels: labels,
		},
		valueFunc,
	)); err == nil {
		p.Log.Infof("GaugeFunc('%s') registered.", gaugeName)
	} else {
		p.Log.Errorf("GaugeFunc('%s') registration failed: %s", gaugeName, err)
	}
}

// String returns plugin name if it was injected, defaultPluginName otherwise.
func (p *PrometheusPlugin) String() string {
	if len(string(p.PluginName)) > 0 {
		return string(p.PluginName)
	}
	return defaultPluginName
}
<commit_msg>Isolated health metrics by /health endpoint<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 probe

import (
	"net/http"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/unrolled/render"
)

const (
	defaultPluginName string = "PROMETHEUS"

	// DefaultMetricsPath default Prometheus metrics URL
	DefaultMetricsPath string = "/metrics"
	// DefaultHealthPath default Prometheus health metrics URL
	DefaultHealthPath string = "/health"

	// Namespace namespace to use for Prometheus health metrics
	Namespace string = ""
	// Subsystem subsystem to use for Prometheus health metrics
	Subsystem string = ""
	// ServiceLabel label for service field
	ServiceLabel string = "service"
	// DependencyLabel label for dependency field
	DependencyLabel string = "dependency"
	// BuildVersionLabel label for build version field
	BuildVersionLabel string = "build_version"
	// BuildDateLabel label for build date field
	BuildDateLabel string = "build_date"

	// ServiceHealthName name of service health metric
	ServiceHealthName string = "service_health"

	// ServiceHealthHelp help text for service health metric
	// Adapt Ligato status code for now.
	// TODO: Consolidate with that from the "Common Container Telemetry" proposal.
	// ServiceHealthHelp    string = "The health of the ServiceLabel 0 = INIT, 1 = UP, 2 = DOWN, 3 = OUTAGE"
	ServiceHealthHelp string = "The health of the ServiceLabel 0 = INIT, 1 = OK, 2 = ERROR"

	// DependencyHealthName name of dependency health metric
	DependencyHealthName string = "service_dependency_health"

	// DependencyHealthHelp help text for dependency health metric
	// Adapt Ligato status code for now.
	// TODO: Consolidate with that from the "Common Container Telemetry" proposal.
	// DependencyHealthHelp string = "The health of the DependencyLabel 0 = INIT, 1 = UP, 2 = DOWN, 3 = OUTAGE"
	DependencyHealthHelp string = "The health of the DependencyLabel 0 = INIT, 1 = OK, 2 = ERROR"

	// ServiceInfoName name of service info metric
	ServiceInfoName string = "service_info"
	// ServiceInfoHelp help text for service info metric
	ServiceInfoHelp string = "Build info for the service.  Value is always 1, build info is in the tags."
)

// PrometheusPlugin struct holds all plugin-related data.
type PrometheusPlugin struct {
	Deps
	healthRegistry *prometheus.Registry
}

// Init may create a new (custom) instance of HTTP if the injected instance uses
// different HTTP port than requested.
func (p *PrometheusPlugin) Init() (err error) {
	serviceLabel := p.String()
	if p.Deps.ServiceLabel != nil {
		serviceLabel = p.Deps.ServiceLabel.GetAgentLabel()
	}

	p.healthRegistry = prometheus.NewRegistry()

	p.registerGauge(
		Namespace,
		Subsystem,
		ServiceHealthName,
		ServiceHealthHelp,
		prometheus.Labels{ServiceLabel: serviceLabel},
		p.getServiceHealth,
	)

	agentStatus := p.StatusCheck.GetAgentStatus()
	p.registerGauge(
		Namespace,
		Subsystem,
		ServiceInfoName,
		ServiceInfoHelp,
		prometheus.Labels{
			ServiceLabel:      serviceLabel,
			BuildVersionLabel: agentStatus.BuildVersion,
			BuildDateLabel:    agentStatus.BuildDate},
		func() float64 { return 1 },
	)

	return nil
}

// AfterInit registers HTTP handlers.
func (p *PrometheusPlugin) AfterInit() error {
	if p.HTTP != nil {
		if p.StatusCheck != nil {
			p.Log.Info("Starting Prometheus metrics handlers")
			p.HTTP.RegisterHTTPHandler(DefaultMetricsPath, p.metricsHandler, "GET")
			p.HTTP.RegisterHTTPHandler(DefaultHealthPath, p.healthMetricsHandler, "GET")
			p.Log.Infof("Serving %s on port %d", DefaultMetricsPath, p.HTTP.GetPort())
			p.Log.Infof("Serving %s on port %d", DefaultHealthPath, p.HTTP.GetPort())
		} else {
			p.Log.Info("Unable to register Prometheus metrics handlers, StatusCheck is nil")
		}
	} else {
		p.Log.Info("Unable to register Prometheus metrics handlers, HTTP is nil")
	}

	return nil
}

// Close shutdowns HTTP if a custom instance was created in Init().
func (p *PrometheusPlugin) Close() error {
	return nil
}

// metricsHandler handles Prometheus metrics collection.
func (p *PrometheusPlugin) metricsHandler(formatter *render.Render) http.HandlerFunc {
	return promhttp.Handler().ServeHTTP
}

// healthMetricsHandler handles custom health metrics for Prometheus.
func (p *PrometheusPlugin) healthMetricsHandler(formatter *render.Render) http.HandlerFunc {
	return promhttp.HandlerFor(p.healthRegistry, promhttp.HandlerOpts{}).ServeHTTP
}

func (p *PrometheusPlugin) getServiceHealth() float64 {
	agentStatus := p.StatusCheck.GetAgentStatus()
	// Adapt Ligato status code for now.
	// TODO: Consolidate with that from the "Common Container Telemetry" proposal.
	health := float64(agentStatus.State)
	p.Log.Infof("getServiceHealth(): %f", health)
	return health
}

// registerGauge registers custom gauge with specific valueFunc to report status when invoked.
func (p *PrometheusPlugin) registerGauge(namespace string, subsystem string, name string, help string,
	labels prometheus.Labels, valueFunc func() float64) {
	gaugeName := name
	if subsystem != "" {
		gaugeName = subsystem + "_" + gaugeName
	}
	if namespace != "" {
		gaugeName = namespace + "_" + gaugeName
	}
	if err := p.healthRegistry.Register(prometheus.NewGaugeFunc(
		prometheus.GaugeOpts{
			// Namespace, Subsystem, and Name are components of the fully-qualified
			// name of the Metric (created by joining these components with
			// "_"). Only Name is mandatory, the others merely help structuring the
			// name. Note that the fully-qualified name of the metric must be a
			// valid Prometheus metric name.
			Namespace: namespace,
			Subsystem: subsystem,
			Name:      name,

			// Help provides information about this metric. Mandatory!
			//
			// Metrics with the same fully-qualified name must have the same Help
			// string.
			Help: help,

			// ConstLabels are used to attach fixed labels to this metric. Metrics
			// with the same fully-qualified name must have the same label names in
			// their ConstLabels.
			//
			// Note that in most cases, labels have a value that varies during the
			// lifetime of a process. Those labels are usually managed with a metric
			// vector collector (like CounterVec, GaugeVec, UntypedVec). ConstLabels
			// serve only special purposes. One is for the special case where the
			// value of a label does not change during the lifetime of a process,
			// e.g. if the revision of the running binary is put into a
			// label. Another, more advanced purpose is if more than one Collector
			// needs to collect Metrics with the same fully-qualified name. In that
			// case, those Metrics must differ in the values of their
			// ConstLabels. See the Collector examples.
			//
			// If the value of a label never changes (not even between binaries),
			// that label most likely should not be a label at all (but part of the
			// metric name).
			ConstLabels: labels,
		},
		valueFunc,
	)); err == nil {
		p.Log.Infof("GaugeFunc('%s') registered.", gaugeName)
	} else {
		p.Log.Errorf("GaugeFunc('%s') registration failed: %s", gaugeName, err)
	}
}

// String returns plugin name if it was injected, defaultPluginName otherwise.
func (p *PrometheusPlugin) String() string {
	if len(string(p.PluginName)) > 0 {
		return string(p.PluginName)
	}
	return defaultPluginName
}
<|endoftext|>


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

package related

import (
	"bytes"
	"mime"
	"net/textproto"
	"reflect"
	"testing"
)

func TestWriter(t *testing.T) {
	fileContents := []byte(`Life? Don't talk to me about life!`)

	var b bytes.Buffer
	w := NewWriter(&b)
	{
		part, err := w.CreateRoot("", "", nil)
		if err != nil {
			t.Fatalf("CreateRoot: %v", err)
		}
		part.Write(fileContents)

		if err := w.Close(); err != nil {
			t.Fatalf("Close: %v", err)
		}

		s := b.String()
		if len(s) == 0 {
			t.Fatal("String: unexpected empty result")
		}
	}

	// NewReader(&b, w.Boundary())
	// TODO Check output
}

func TestCreateRootFail(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)

	// Error handling
	testsError := []struct {
		id    string
		media string
	}{
		{"dont", "dont/panic"},
		{"", "dont;panic"},
		{"dont", ""},
	}

	for i, tt := range testsError {
		if _, err := w.CreateRoot(tt.id, tt.media, nil); err == nil {
			t.Errorf("%d. Content-Id: %s, Media-Type: %s", i, tt.id, tt.media)
		}
	}

	for i := 2; i > 0; i-- {
		_, err := w.CreateRoot("", "a/b", nil)
		if i == 1 && err != ErrRootExists {
			t.Errorf("%d. Multiple CreateRoot: Expected error", i)
		}
	}

	w.Close()
}

func TestCreatePartFail(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)

	// Error handling
	id := "&&;&;&"
	if _, err := w.CreatePart(id, nil); err == nil {
		t.Errorf("Content-Id: %s", id)
	}

	w.Close()
}

func TestCreatePartFirst(t *testing.T) {
	h := textproto.MIMEHeader{}
	h.Add("Content-Type", "a/b")

	tests := []struct {
		id        string
		header    textproto.MIMEHeader
		mediaType string
	}{
		{"a@b.c", nil, DefaultMediaType},
		{"a@b.c", h, "a/b"},
	}

	for i, tt := range tests {
		var b bytes.Buffer
		w := NewWriter(&b)

		if w.firstPart != false {
			t.Errorf("Before:\n%d. firstPart = %t, want %t", i, w.firstPart, false)
		}

		if _, err := w.CreatePart(tt.id, tt.header); err != nil {
			t.Fatalf("%d. CreatePart: %v", i, err)
		}
		if w.mediaType != tt.mediaType {
			t.Errorf("%d. type = %s, want %s", i, w.mediaType, tt.mediaType)
		}
		if w.rootMediaType != tt.mediaType {
			t.Errorf("%d. type = %s, want %s", i, w.rootMediaType, tt.mediaType)
		}

		if w.firstPart != true {
			t.Errorf("After:\n%d. firstPart = %t, want %t", i, w.firstPart, true)
		}

		w.Close()
	}
}

func TestSetStart(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)
	tests := []struct {
		id string
		w  string
		ok bool
	}{
		{"a@b.c", "<a@b.c>", true},
		{"aa", "", false},
	}

	for i, tt := range tests {
		err := w.SetStart(tt.id)
		got := err == nil
		if got != tt.ok {
			t.Errorf("%d. start %q = %v (%v), want %v", i, tt.id, got, err, tt.ok)
		} else if tt.ok {
			got := w.start
			if got != tt.w {
				t.Errorf("start = %q; want %q", got, tt.w)
			}
		}
	}
	w.Close()
}

func TestSetType(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)
	tests := []struct {
		t  string
		ok bool
	}{
		{"application/json", true},
		{";", false},
	}

	for i, tt := range tests {
		err := w.SetType(tt.t)
		got := err == nil
		if got != tt.ok {
			t.Errorf("%d. start %q = %v (%v), want %v", i, tt.t, got, err, tt.ok)
		} else if tt.ok {
			got := w.mediaType
			if got != tt.t {
				t.Errorf("start = %q; want %q", got, tt.t)
			}
		}
	}
	w.Close()
}

func TestSetBoundary(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)
	tests := []struct {
		b  string
		ok bool
	}{
		{"abc", true},
		{"ungültig", false},
	}

	for i, tt := range tests {
		err := w.SetBoundary(tt.b)
		got := err == nil
		if got != tt.ok {
			t.Errorf("%d. start %q = %v (%v), want %v", i, tt.b, got, err, tt.ok)
		} else if tt.ok {
			got := w.Boundary()
			if got != tt.b {
				t.Errorf("start = %q; want %q", got, tt.b)
			}
		}
	}

	w.Close()
}

func TestClose(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)

	if _, err := w.CreateRoot("a@b.c", "text/plain", nil); err != nil {
		t.Fatalf("CreateRoot: %v", err)
	}
	if err := w.SetType("text/html"); err != nil {
		t.Fatalf("SetType: %v", err)
	}
	if err := w.Close(); err != ErrTypeMatch {
		t.Errorf("NoMediaType = %v; want %q", err, ErrTypeMatch)
	}
	w.Close()
}

func TestFormDataContentType(t *testing.T) {
	var b bytes.Buffer

	in := map[string]string{
		"boundary":   "abc",
		"type":       "text/plain",
		"start":      "a@b.c",
		"start-info": `-o p"s`,
	}
	want := map[string]string{
		"boundary":   "abc",
		"type":       "text/plain",
		"start":      "<a@b.c>",
		"start-info": `-o p\"s`,
	}

	w := NewWriter(&b)

	if err := w.SetBoundary(in["boundary"]); err != nil {
		t.Fatalf("SetBoundary: %v", err)
	}
	if err := w.SetType(in["type"]); err != nil {
		t.Fatalf("SetType: %v", err)
	}
	if err := w.SetStart(in["start"]); err != nil {
		t.Fatalf("SetStart: %v", err)
	}
	w.SetStartInfo(in["start-info"])

	got := w.FormDataContentType()
	mediatype, params, err := mime.ParseMediaType(got)
	if err != nil {
		t.Fatalf("ParseMediaType: %v", err)
	}
	if mediatype != "multipart/related" {
		t.Errorf("mediatype = %s, want multipart/related", mediatype)
	}
	if !reflect.DeepEqual(params, want) {
		t.Errorf("params = %v, want %v", params, want)
	}

	w.Close()
}

func TestFormatContentId(t *testing.T) {
	tests := []struct {
		id string
		w  string
		ok bool
	}{
		{"a@b.c", "<a@b.c>", true},
		{"<aa>", "", false},
		{"", "", false},
	}

	for i, tt := range tests {
		got, err := formatContentId(tt.id)
		if err == nil != tt.ok {
			t.Errorf("%d. start %q = %v (%v), want %v", i, tt.id, got, err, tt.ok)
		} else if tt.ok {
			if got != tt.w {
				t.Errorf("start = %q; want %q", got, tt.w)
			}
		}
	}
}
<commit_msg>Complete Write test<commit_after>// Copyright 2015 The multipart-related Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package related

import (
	"bytes"
	"io/ioutil"
	"mime"
	"net/textproto"
	"reflect"
	"testing"
)

func TestWriter(t *testing.T) {
	fileContents1 := []byte(`Life? Don't talk to me about life!`)
	fileContents2 := []byte(`Marvin`)

	var b bytes.Buffer
	w := NewWriter(&b)
	{
		part, err := w.CreateRoot("", "a/b", nil)
		if err != nil {
			t.Fatalf("CreateRoot: %v", err)
		}
		part.Write(fileContents1)

		nextPart, err := w.CreatePart("", nil)
		if err != nil {
			t.Fatalf("CreatePart 2: %v", err)
		}

		nextPart.Write(fileContents2)

		if err := w.Close(); err != nil {
			t.Fatalf("Close: %v", err)
		}

		s := b.String()
		if len(s) == 0 {
			t.Fatal("String: unexpected empty result")
		}
	}

	r := NewReader(&b, map[string]string{
		"boundary": w.Boundary(),
	})
	part, err := r.NextPart()
	if err != nil {
		t.Fatalf("part root: %v", err)
	}
	if g, w := part.Header.Get("Content-Type"), "a/b"; g != w {
		t.Errorf("part root: got content-type: %s, want %s", g, w)
	}
	slurp, err := ioutil.ReadAll(part)
	if err != nil {
		t.Fatalf("part root: ReadAll: %v", err)
	}
	if g, w := string(slurp), string(fileContents1); w != g {
		t.Errorf("part root: got contents %q, want %q", g, w)
	}
	part, err = r.NextPart()
	if err != nil {
		t.Fatalf("part 2: %v", err)
	}
	if g, w := part.Header.Get("Content-Type"), "text/plain; charset=utf-8"; g != w {
		t.Errorf("part 2: got content-type: %s, want %s", g, w)
	}
	slurp, err = ioutil.ReadAll(part)
	if err != nil {
		t.Fatalf("part 2: ReadAll: %v", err)
	}
	if g, w := string(slurp), string(fileContents2); w != g {
		t.Errorf("part 2: got contents %q, want %q", g, w)
	}
	part, err = r.NextPart()
	if part != nil || err == nil {
		t.Fatalf("expected end of parts; got %v, %v", part, err)
	}
}

func TestCreateRootFail(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)

	// Error handling
	testsError := []struct {
		id    string
		media string
	}{
		{"dont", "dont/panic"},
		{"", "dont;panic"},
		{"dont", ""},
	}

	for i, tt := range testsError {
		if _, err := w.CreateRoot(tt.id, tt.media, nil); err == nil {
			t.Errorf("%d. Content-Id: %s, Media-Type: %s", i, tt.id, tt.media)
		}
	}

	for i := 2; i > 0; i-- {
		_, err := w.CreateRoot("", "a/b", nil)
		if i == 1 && err != ErrRootExists {
			t.Errorf("%d. Multiple CreateRoot: Expected error", i)
		}
	}

	w.Close()
}

func TestCreatePartFail(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)

	// Error handling
	id := "&&;&;&"
	if _, err := w.CreatePart(id, nil); err == nil {
		t.Errorf("Content-Id: %s", id)
	}

	w.Close()
}

func TestCreatePartFirst(t *testing.T) {
	h := textproto.MIMEHeader{}
	h.Add("Content-Type", "a/b")

	tests := []struct {
		id        string
		header    textproto.MIMEHeader
		mediaType string
	}{
		{"a@b.c", nil, DefaultMediaType},
		{"a@b.c", h, "a/b"},
	}

	for i, tt := range tests {
		var b bytes.Buffer
		w := NewWriter(&b)

		if w.firstPart != false {
			t.Errorf("Before:\n%d. firstPart = %t, want %t", i, w.firstPart, false)
		}

		if _, err := w.CreatePart(tt.id, tt.header); err != nil {
			t.Fatalf("%d. CreatePart: %v", i, err)
		}
		if w.mediaType != tt.mediaType {
			t.Errorf("%d. type = %s, want %s", i, w.mediaType, tt.mediaType)
		}
		if w.rootMediaType != tt.mediaType {
			t.Errorf("%d. type = %s, want %s", i, w.rootMediaType, tt.mediaType)
		}

		if w.firstPart != true {
			t.Errorf("After:\n%d. firstPart = %t, want %t", i, w.firstPart, true)
		}

		w.Close()
	}
}

func TestSetStart(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)
	tests := []struct {
		id string
		w  string
		ok bool
	}{
		{"a@b.c", "<a@b.c>", true},
		{"aa", "", false},
	}

	for i, tt := range tests {
		err := w.SetStart(tt.id)
		got := err == nil
		if got != tt.ok {
			t.Errorf("%d. start %q = %v (%v), want %v", i, tt.id, got, err, tt.ok)
		} else if tt.ok {
			got := w.start
			if got != tt.w {
				t.Errorf("start = %q; want %q", got, tt.w)
			}
		}
	}
	w.Close()
}

func TestSetType(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)
	tests := []struct {
		t  string
		ok bool
	}{
		{"application/json", true},
		{";", false},
	}

	for i, tt := range tests {
		err := w.SetType(tt.t)
		got := err == nil
		if got != tt.ok {
			t.Errorf("%d. start %q = %v (%v), want %v", i, tt.t, got, err, tt.ok)
		} else if tt.ok {
			got := w.mediaType
			if got != tt.t {
				t.Errorf("start = %q; want %q", got, tt.t)
			}
		}
	}
	w.Close()
}

func TestSetBoundary(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)
	tests := []struct {
		b  string
		ok bool
	}{
		{"abc", true},
		{"ungültig", false},
	}

	for i, tt := range tests {
		err := w.SetBoundary(tt.b)
		got := err == nil
		if got != tt.ok {
			t.Errorf("%d. start %q = %v (%v), want %v", i, tt.b, got, err, tt.ok)
		} else if tt.ok {
			got := w.Boundary()
			if got != tt.b {
				t.Errorf("start = %q; want %q", got, tt.b)
			}
		}
	}

	w.Close()
}

func TestClose(t *testing.T) {
	var b bytes.Buffer
	w := NewWriter(&b)

	if _, err := w.CreateRoot("a@b.c", "text/plain", nil); err != nil {
		t.Fatalf("CreateRoot: %v", err)
	}
	if err := w.SetType("text/html"); err != nil {
		t.Fatalf("SetType: %v", err)
	}
	if err := w.Close(); err != ErrTypeMatch {
		t.Errorf("NoMediaType = %v; want %q", err, ErrTypeMatch)
	}
	w.Close()
}

func TestFormDataContentType(t *testing.T) {
	var b bytes.Buffer

	in := map[string]string{
		"boundary":   "abc",
		"type":       "text/plain",
		"start":      "a@b.c",
		"start-info": `-o p"s`,
	}
	want := map[string]string{
		"boundary":   "abc",
		"type":       "text/plain",
		"start":      "<a@b.c>",
		"start-info": `-o p\"s`,
	}

	w := NewWriter(&b)

	if err := w.SetBoundary(in["boundary"]); err != nil {
		t.Fatalf("SetBoundary: %v", err)
	}
	if err := w.SetType(in["type"]); err != nil {
		t.Fatalf("SetType: %v", err)
	}
	if err := w.SetStart(in["start"]); err != nil {
		t.Fatalf("SetStart: %v", err)
	}
	w.SetStartInfo(in["start-info"])

	got := w.FormDataContentType()
	mediatype, params, err := mime.ParseMediaType(got)
	if err != nil {
		t.Fatalf("ParseMediaType: %v", err)
	}
	if mediatype != "multipart/related" {
		t.Errorf("mediatype = %s, want multipart/related", mediatype)
	}
	if !reflect.DeepEqual(params, want) {
		t.Errorf("params = %v, want %v", params, want)
	}

	w.Close()
}

func TestFormatContentId(t *testing.T) {
	tests := []struct {
		id string
		w  string
		ok bool
	}{
		{"a@b.c", "<a@b.c>", true},
		{"<aa>", "", false},
		{"", "", false},
	}

	for i, tt := range tests {
		got, err := formatContentId(tt.id)
		if err == nil != tt.ok {
			t.Errorf("%d. start %q = %v (%v), want %v", i, tt.id, got, err, tt.ok)
		} else if tt.ok {
			if got != tt.w {
				t.Errorf("start = %q; want %q", got, tt.w)
			}
		}
	}
}
<|endoftext|>


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

package git

import (
	"strconv"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestRefShortName(t *testing.T) {
	tests := []struct {
		ref    string
		expVal string
	}{
		{
			ref:    "refs/heads/master",
			expVal: "master",
		},
		{
			ref:    "refs/tags/v1.0.0",
			expVal: "v1.0.0",
		},
		{
			ref:    "refs/pull/98",
			expVal: "refs/pull/98",
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			assert.Equal(t, test.expVal, RefShortName(test.ref))
		})
	}
}

func TestRepository_ShowRefVerify(t *testing.T) {
	t.Run("reference does not exsit", func(t *testing.T) {
		rev, err := testrepo.ShowRefVerify("bad_reference")
		assert.NotNil(t, err)
		assert.Empty(t, rev)
	})

	rev, err := testrepo.ShowRefVerify("refs/heads/release-1.0")
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, "0eedd79eba4394bbef888c804e899731644367fe", rev)
}

func TestRepository_HasReference(t *testing.T) {
	tests := []struct {
		ref    string
		opt    ShowRefVerifyOptions
		expVal bool
	}{
		{
			ref:    RefsHeads + "master",
			expVal: true,
		},
		{
			ref:    "master",
			expVal: false,
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			assert.Equal(t, test.expVal, testrepo.HasReference(test.ref, test.opt))
		})
	}
}

func TestRepository_SymbolicRef(t *testing.T) {
	r, cleanup, err := setupTempRepo()
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()

	// Get HEAD
	ref, err := r.SymbolicRef()
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, RefsHeads+"master", ref)

	// Set a symbolic reference
	_, err = r.SymbolicRef(SymbolicRefOptions{
		Name: "TEST-REF",
		Ref:  RefsHeads + "develop",
	})
	if err != nil {
		t.Fatal(err)
	}

	// Get the symbolic reference we just set
	ref, err = r.SymbolicRef(SymbolicRefOptions{
		Name: "TEST-REF",
	})
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, RefsHeads+"develop", ref)
}

func TestRepository_ShowRef(t *testing.T) {
	tests := []struct {
		opt     ShowRefOptions
		expRefs []*Reference
	}{
		{
			opt: ShowRefOptions{
				Heads:    true,
				Patterns: []string{"release-1.0"},
			},
			expRefs: []*Reference{
				{
					ID:      "0eedd79eba4394bbef888c804e899731644367fe",
					Refspec: "refs/heads/release-1.0",
				},
			},
		}, {
			opt: ShowRefOptions{
				Tags:     true,
				Patterns: []string{"v1.0.0"},
			},
			expRefs: []*Reference{
				{
					ID:      "0eedd79eba4394bbef888c804e899731644367fe",
					Refspec: "refs/tags/v1.0.0",
				},
			},
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			refs, err := testrepo.ShowRef(test.opt)
			if err != nil {
				t.Fatal(err)
			}

			assert.Equal(t, test.expRefs, refs)
		})
	}
}

func TestRepository_DeleteBranch(t *testing.T) {
	r, cleanup, err := setupTempRepo()
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()

	tests := []struct {
		opt DeleteBranchOptions
	}{
		{
			opt: DeleteBranchOptions{
				Force: false,
			},
		},
		{
			opt: DeleteBranchOptions{
				Force: true,
			},
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			branch := strconv.Itoa(int(time.Now().UnixNano()))
			err := r.Checkout(branch, CheckoutOptions{
				BaseBranch: "master",
			})
			if err != nil {
				t.Fatal(err)
			}

			assert.True(t, r.HasReference(RefsHeads+branch))

			err = r.Checkout("master")
			if err != nil {
				t.Fatal(err)
			}

			err = r.DeleteBranch(branch, test.opt)
			if err != nil {
				t.Fatal(err)
			}

			assert.False(t, r.HasReference(RefsHeads+branch))
		})
	}
}
<commit_msg>repo_reference: adds more tests<commit_after>// Copyright 2020 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package git

import (
	"strconv"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestRefShortName(t *testing.T) {
	tests := []struct {
		ref    string
		expVal string
	}{
		{
			ref:    "refs/heads/master",
			expVal: "master",
		},
		{
			ref:    "refs/tags/v1.0.0",
			expVal: "v1.0.0",
		},
		{
			ref:    "refs/pull/98",
			expVal: "refs/pull/98",
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			assert.Equal(t, test.expVal, RefShortName(test.ref))
		})
	}
}

func TestRepository_ShowRefVerify(t *testing.T) {
	t.Run("reference does not exsit", func(t *testing.T) {
		rev, err := testrepo.ShowRefVerify("bad_reference")
		assert.NotNil(t, err)
		assert.Empty(t, rev)
	})

	rev, err := testrepo.ShowRefVerify("refs/heads/release-1.0")
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, "0eedd79eba4394bbef888c804e899731644367fe", rev)
}

func TestRepository_HasReference(t *testing.T) {
	tests := []struct {
		ref    string
		opt    ShowRefVerifyOptions
		expVal bool
	}{
		{
			ref:    RefsHeads + "master",
			expVal: true,
		},
		{
			ref:    RefsTags + "v1.0.0",
			expVal: true,
		},
		{
			ref:    "master",
			expVal: false,
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			assert.Equal(t, test.expVal, testrepo.HasReference(test.ref, test.opt))
		})
	}
}

func TestRepository_HasBranch(t *testing.T) {
	tests := []struct {
		ref    string
		opt    ShowRefVerifyOptions
		expVal bool
	}{
		{
			ref:    "master",
			expVal: true,
		},
		{
			ref:    RefsHeads + "master",
			expVal: false,
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			assert.Equal(t, test.expVal, testrepo.HasBranch(test.ref, test.opt))
		})
	}
}

func TestRepository_HasTag(t *testing.T) {
	tests := []struct {
		ref    string
		opt    ShowRefVerifyOptions
		expVal bool
	}{
		{
			ref:    "v1.0.0",
			expVal: true,
		},
		{
			ref:    RefsTags + "v1.0.0",
			expVal: false,
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			assert.Equal(t, test.expVal, testrepo.HasTag(test.ref, test.opt))
		})
	}
}

func TestRepository_SymbolicRef(t *testing.T) {
	r, cleanup, err := setupTempRepo()
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()

	// Get HEAD
	ref, err := r.SymbolicRef()
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, RefsHeads+"master", ref)

	// Set a symbolic reference
	_, err = r.SymbolicRef(SymbolicRefOptions{
		Name: "TEST-REF",
		Ref:  RefsHeads + "develop",
	})
	if err != nil {
		t.Fatal(err)
	}

	// Get the symbolic reference we just set
	ref, err = r.SymbolicRef(SymbolicRefOptions{
		Name: "TEST-REF",
	})
	if err != nil {
		t.Fatal(err)
	}
	assert.Equal(t, RefsHeads+"develop", ref)
}

func TestRepository_ShowRef(t *testing.T) {
	tests := []struct {
		opt     ShowRefOptions
		expRefs []*Reference
	}{
		{
			opt: ShowRefOptions{
				Heads:    true,
				Patterns: []string{"release-1.0"},
			},
			expRefs: []*Reference{
				{
					ID:      "0eedd79eba4394bbef888c804e899731644367fe",
					Refspec: "refs/heads/release-1.0",
				},
			},
		}, {
			opt: ShowRefOptions{
				Tags:     true,
				Patterns: []string{"v1.0.0"},
			},
			expRefs: []*Reference{
				{
					ID:      "0eedd79eba4394bbef888c804e899731644367fe",
					Refspec: "refs/tags/v1.0.0",
				},
			},
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			refs, err := testrepo.ShowRef(test.opt)
			if err != nil {
				t.Fatal(err)
			}

			assert.Equal(t, test.expRefs, refs)
		})
	}
}

func TestRepository_DeleteBranch(t *testing.T) {
	r, cleanup, err := setupTempRepo()
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()

	tests := []struct {
		opt DeleteBranchOptions
	}{
		{
			opt: DeleteBranchOptions{
				Force: false,
			},
		},
		{
			opt: DeleteBranchOptions{
				Force: true,
			},
		},
	}
	for _, test := range tests {
		t.Run("", func(t *testing.T) {
			branch := strconv.Itoa(int(time.Now().UnixNano()))
			err := r.Checkout(branch, CheckoutOptions{
				BaseBranch: "master",
			})
			if err != nil {
				t.Fatal(err)
			}

			assert.True(t, r.HasReference(RefsHeads+branch))

			err = r.Checkout("master")
			if err != nil {
				t.Fatal(err)
			}

			err = r.DeleteBranch(branch, test.opt)
			if err != nil {
				t.Fatal(err)
			}

			assert.False(t, r.HasReference(RefsHeads+branch))
		})
	}
}
<|endoftext|>


<commit_before>package baa

import (
	"errors"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
)

const (
	// DEV mode
	DEV = "development"
	// PROD mode
	PROD = "production"
	// TEST mode
	TEST = "test"
)

// Env default application runtime environment
var Env string

// Baa provlider an application
type Baa struct {
	debug           bool
	name            string
	di              DIer
	router          Router
	pool            sync.Pool
	errorHandler    ErrorHandleFunc
	notFoundHandler HandlerFunc
	middleware      []HandlerFunc
}

// Middleware middleware handler
type Middleware interface{}

// HandlerFunc context handler func
type HandlerFunc func(*Context)

// ErrorHandleFunc HTTP error handleFunc
type ErrorHandleFunc func(error, *Context)

// appInstances storage application instances
var appInstances map[string]*Baa

// defaultAppName default application name
const defaultAppName = "_default_"

// New create a baa application without any config.
func New() *Baa {
	b := new(Baa)
	b.middleware = make([]HandlerFunc, 0)
	b.pool = sync.Pool{
		New: func() interface{} {
			return NewContext(nil, nil, b)
		},
	}
	if Env != PROD {
		b.debug = true
	}
	b.SetDIer(NewDI())
	b.SetDI("router", NewTree(b))
	b.SetDI("logger", log.New(os.Stderr, "[Baa] ", log.LstdFlags))
	b.SetDI("render", newRender())
	b.SetNotFound(b.DefaultNotFoundHandler)
	return b
}

// Instance register or returns named application
func Instance(name string) *Baa {
	if name == "" {
		name = defaultAppName
	}
	if appInstances[name] == nil {
		appInstances[name] = New()
		appInstances[name].name = name
	}
	return appInstances[name]
}

// Default initial a default app then returns
func Default() *Baa {
	return Instance(defaultAppName)
}

// Server returns the internal *http.Server.
func (b *Baa) Server(addr string) *http.Server {
	s := &http.Server{Addr: addr}
	return s
}

// Run runs a server.
func (b *Baa) Run(addr string) {
	b.run(b.Server(addr))
}

// RunTLS runs a server with TLS configuration.
func (b *Baa) RunTLS(addr, certfile, keyfile string) {
	b.run(b.Server(addr), certfile, keyfile)
}

// RunServer runs a custom server.
func (b *Baa) RunServer(s *http.Server) {
	b.run(s)
}

// RunTLSServer runs a custom server with TLS configuration.
func (b *Baa) RunTLSServer(s *http.Server, crtFile, keyFile string) {
	b.run(s, crtFile, keyFile)
}

func (b *Baa) run(s *http.Server, files ...string) {
	s.Handler = b
	b.Logger().Printf("Run mode: %s", Env)
	if len(files) == 0 {
		b.Logger().Printf("Listen %s", s.Addr)
		b.Logger().Fatal(s.ListenAndServe())
	} else if len(files) == 2 {
		b.Logger().Printf("Listen %s with TLS", s.Addr)
		b.Logger().Fatal(s.ListenAndServeTLS(files[0], files[1]))
	} else {
		panic("invalid TLS configuration")
	}
}

func (b *Baa) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	c := b.pool.Get().(*Context)
	c.Reset(w, r)

	// build handler chain
	h, name := b.Router().Match(r.Method, r.URL.Path, c)
	c.routeName = name

	// notFound
	if h == nil {
		c.handlers = append(c.handlers, b.notFoundHandler)
	} else {
		c.handlers = append(c.handlers, h...)
	}

	c.Next()

	b.pool.Put(c)
}

// SetDIer set baa di
func (b *Baa) SetDIer(v DIer) {
	b.di = v
}

// SetDebug set baa debug
func (b *Baa) SetDebug(v bool) {
	b.debug = v
}

// Debug returns baa debug state
func (b *Baa) Debug() bool {
	return b.debug
}

// Logger return baa logger
func (b *Baa) Logger() Logger {
	return b.GetDI("logger").(Logger)
}

// Render return baa render
func (b *Baa) Render() Renderer {
	return b.GetDI("render").(Renderer)
}

// Router return baa router
func (b *Baa) Router() Router {
	if b.router == nil {
		b.router = b.GetDI("router").(Router)
	}
	return b.router
}

// Use registers a middleware
func (b *Baa) Use(m ...Middleware) {
	for i := range m {
		if m[i] != nil {
			b.middleware = append(b.middleware, wrapMiddleware(m[i]))
		}
	}
}

// SetDI registers a dependency injection
func (b *Baa) SetDI(name string, h interface{}) {
	switch name {
	case "logger":
		if _, ok := h.(Logger); !ok {
			panic("DI logger must be implement interface baa.Logger")
		}
	case "render":
		if _, ok := h.(Renderer); !ok {
			panic("DI render must be implement interface baa.Renderer")
		}
	case "router":
		if _, ok := h.(Router); !ok {
			panic("DI router must be implement interface baa.Router")
		}
	}
	b.di.Set(name, h)
}

// GetDI fetch a registered dependency injection
func (b *Baa) GetDI(name string) interface{} {
	return b.di.Get(name)
}

// Static set static file route
// h used for set Expries ...
func (b *Baa) Static(prefix string, dir string, index bool, h HandlerFunc) {
	if prefix == "" {
		panic("baa.Static prefix can not be empty")
	}
	if dir == "" {
		panic("baa.Static dir can not be empty")
	}
	b.Get(prefix+"*", newStatic(prefix, dir, index, h))
}

// StaticFile shortcut for serve file
func (b *Baa) StaticFile(pattern string, path string) RouteNode {
	return b.Get(pattern, func(c *Context) {
		if err := serveFile(path, c); err != nil {
			c.Error(err)
		}
	})
}

// SetAutoHead sets the value who determines whether add HEAD method automatically
// when GET method is added. Combo router will not be affected by this value.
func (b *Baa) SetAutoHead(v bool) {
	b.Router().SetAutoHead(v)
}

// SetAutoTrailingSlash optional trailing slash.
func (b *Baa) SetAutoTrailingSlash(v bool) {
	b.Router().SetAutoTrailingSlash(v)
}

// Route is a shortcut for same handlers but different HTTP methods.
//
// Example:
// 		baa.Route("/", "GET,POST", h)
func (b *Baa) Route(pattern, methods string, h ...HandlerFunc) RouteNode {
	var ru RouteNode
	var ms []string
	if methods == "*" {
		for m := range RouterMethods {
			ms = append(ms, m)
		}
	} else {
		ms = strings.Split(methods, ",")
	}
	for _, m := range ms {
		ru = b.Router().Add(strings.TrimSpace(m), pattern, h)
	}
	return ru
}

// Group registers a list of same prefix route
func (b *Baa) Group(pattern string, f func(), h ...HandlerFunc) {
	b.Router().GroupAdd(pattern, f, h)
}

// Any is a shortcut for b.Router().handle("*", pattern, handlers)
func (b *Baa) Any(pattern string, h ...HandlerFunc) RouteNode {
	var ru RouteNode
	for m := range RouterMethods {
		ru = b.Router().Add(m, pattern, h)
	}
	return ru
}

// Delete is a shortcut for b.Route(pattern, "DELETE", handlers)
func (b *Baa) Delete(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("DELETE", pattern, h)
}

// Get is a shortcut for b.Route(pattern, "GET", handlers)
func (b *Baa) Get(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("GET", pattern, h)
}

// Head is a shortcut forb.Route(pattern, "Head", handlers)
func (b *Baa) Head(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("HEAD", pattern, h)
}

// Options is a shortcut for b.Route(pattern, "Options", handlers)
func (b *Baa) Options(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("OPTIONS", pattern, h)
}

// Patch is a shortcut for b.Route(pattern, "PATCH", handlers)
func (b *Baa) Patch(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("PATCH", pattern, h)
}

// Post is a shortcut for b.Route(pattern, "POST", handlers)
func (b *Baa) Post(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("POST", pattern, h)
}

// Put is a shortcut for b.Route(pattern, "Put", handlers)
func (b *Baa) Put(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("PUT", pattern, h)
}

// Websocket register a websocket router handler
func (b *Baa) Websocket(pattern string, h func(*websocket.Conn)) RouteNode {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:    4096,
		WriteBufferSize:   4096,
		EnableCompression: true,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	return b.Route(pattern, "GET,POST", func(c *Context) {
		conn, err := upgrader.Upgrade(c.Resp, c.Req, nil)
		if err != nil {
			b.Logger().Panicf("websocket upgrade connection error: %v", err)
			return
		}
		h(conn)
	})
}

// SetNotFound set not found route handler
func (b *Baa) SetNotFound(h HandlerFunc) {
	b.notFoundHandler = h
}

// NotFound execute not found handler
func (b *Baa) NotFound(c *Context) {
	if b.notFoundHandler != nil {
		b.notFoundHandler(c)
		return
	}
	http.NotFound(c.Resp, c.Req)
}

// SetError set error handler
func (b *Baa) SetError(h ErrorHandleFunc) {
	b.errorHandler = h
}

// Error execute internal error handler
func (b *Baa) Error(err error, c *Context) {
	if err == nil {
		err = errors.New("Internal Server Error")
	}
	if b.errorHandler != nil {
		b.errorHandler(err, c)
		return
	}
	code := http.StatusInternalServerError
	msg := http.StatusText(code)
	if b.debug {
		msg = err.Error()
	}
	b.Logger().Println(err)
	http.Error(c.Resp, msg, code)
}

// DefaultNotFoundHandler invokes the default HTTP error handler.
func (b *Baa) DefaultNotFoundHandler(c *Context) {
	code := http.StatusNotFound
	msg := http.StatusText(code)
	http.Error(c.Resp, msg, code)
}

// URLFor use named route return format url
func (b *Baa) URLFor(name string, args ...interface{}) string {
	return b.Router().URLFor(name, args...)
}

// wrapMiddleware wraps middleware.
func wrapMiddleware(m Middleware) HandlerFunc {
	switch m := m.(type) {
	case HandlerFunc:
		return m
	case func(*Context):
		return m
	case http.Handler, http.HandlerFunc:
		return WrapHandlerFunc(func(c *Context) {
			m.(http.Handler).ServeHTTP(c.Resp, c.Req)
		})
	case func(http.ResponseWriter, *http.Request):
		return WrapHandlerFunc(func(c *Context) {
			m(c.Resp, c.Req)
		})
	default:
		panic("unknown middleware")
	}
}

// WrapHandlerFunc wrap for context handler chain
func WrapHandlerFunc(h HandlerFunc) HandlerFunc {
	return func(c *Context) {
		h(c)
		c.Next()
	}
}

func init() {
	appInstances = make(map[string]*Baa)
	Env = os.Getenv("BAA_ENV")
	if Env == "" {
		Env = DEV
	}
}
<commit_msg>fix url path with two slash<commit_after>package baa

import (
	"errors"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
)

const (
	// DEV mode
	DEV = "development"
	// PROD mode
	PROD = "production"
	// TEST mode
	TEST = "test"
)

// Env default application runtime environment
var Env string

// Baa provlider an application
type Baa struct {
	debug           bool
	name            string
	di              DIer
	router          Router
	pool            sync.Pool
	errorHandler    ErrorHandleFunc
	notFoundHandler HandlerFunc
	middleware      []HandlerFunc
}

// Middleware middleware handler
type Middleware interface{}

// HandlerFunc context handler func
type HandlerFunc func(*Context)

// ErrorHandleFunc HTTP error handleFunc
type ErrorHandleFunc func(error, *Context)

// appInstances storage application instances
var appInstances map[string]*Baa

// defaultAppName default application name
const defaultAppName = "_default_"

// New create a baa application without any config.
func New() *Baa {
	b := new(Baa)
	b.middleware = make([]HandlerFunc, 0)
	b.pool = sync.Pool{
		New: func() interface{} {
			return NewContext(nil, nil, b)
		},
	}
	if Env != PROD {
		b.debug = true
	}
	b.SetDIer(NewDI())
	b.SetDI("router", NewTree(b))
	b.SetDI("logger", log.New(os.Stderr, "[Baa] ", log.LstdFlags))
	b.SetDI("render", newRender())
	b.SetNotFound(b.DefaultNotFoundHandler)
	return b
}

// Instance register or returns named application
func Instance(name string) *Baa {
	if name == "" {
		name = defaultAppName
	}
	if appInstances[name] == nil {
		appInstances[name] = New()
		appInstances[name].name = name
	}
	return appInstances[name]
}

// Default initial a default app then returns
func Default() *Baa {
	return Instance(defaultAppName)
}

// Server returns the internal *http.Server.
func (b *Baa) Server(addr string) *http.Server {
	s := &http.Server{Addr: addr}
	return s
}

// Run runs a server.
func (b *Baa) Run(addr string) {
	b.run(b.Server(addr))
}

// RunTLS runs a server with TLS configuration.
func (b *Baa) RunTLS(addr, certfile, keyfile string) {
	b.run(b.Server(addr), certfile, keyfile)
}

// RunServer runs a custom server.
func (b *Baa) RunServer(s *http.Server) {
	b.run(s)
}

// RunTLSServer runs a custom server with TLS configuration.
func (b *Baa) RunTLSServer(s *http.Server, crtFile, keyFile string) {
	b.run(s, crtFile, keyFile)
}

func (b *Baa) run(s *http.Server, files ...string) {
	s.Handler = b
	b.Logger().Printf("Run mode: %s", Env)
	if len(files) == 0 {
		b.Logger().Printf("Listen %s", s.Addr)
		b.Logger().Fatal(s.ListenAndServe())
	} else if len(files) == 2 {
		b.Logger().Printf("Listen %s with TLS", s.Addr)
		b.Logger().Fatal(s.ListenAndServeTLS(files[0], files[1]))
	} else {
		panic("invalid TLS configuration")
	}
}

func (b *Baa) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	c := b.pool.Get().(*Context)
	c.Reset(w, r)

	// build handler chain
	path := strings.Replace(r.URL.Path, "//", "/", -1)
	h, name := b.Router().Match(r.Method, path, c)
	c.routeName = name

	// notFound
	if h == nil {
		c.handlers = append(c.handlers, b.notFoundHandler)
	} else {
		c.handlers = append(c.handlers, h...)
	}

	c.Next()

	b.pool.Put(c)
}

// SetDIer set baa di
func (b *Baa) SetDIer(v DIer) {
	b.di = v
}

// SetDebug set baa debug
func (b *Baa) SetDebug(v bool) {
	b.debug = v
}

// Debug returns baa debug state
func (b *Baa) Debug() bool {
	return b.debug
}

// Logger return baa logger
func (b *Baa) Logger() Logger {
	return b.GetDI("logger").(Logger)
}

// Render return baa render
func (b *Baa) Render() Renderer {
	return b.GetDI("render").(Renderer)
}

// Router return baa router
func (b *Baa) Router() Router {
	if b.router == nil {
		b.router = b.GetDI("router").(Router)
	}
	return b.router
}

// Use registers a middleware
func (b *Baa) Use(m ...Middleware) {
	for i := range m {
		if m[i] != nil {
			b.middleware = append(b.middleware, wrapMiddleware(m[i]))
		}
	}
}

// SetDI registers a dependency injection
func (b *Baa) SetDI(name string, h interface{}) {
	switch name {
	case "logger":
		if _, ok := h.(Logger); !ok {
			panic("DI logger must be implement interface baa.Logger")
		}
	case "render":
		if _, ok := h.(Renderer); !ok {
			panic("DI render must be implement interface baa.Renderer")
		}
	case "router":
		if _, ok := h.(Router); !ok {
			panic("DI router must be implement interface baa.Router")
		}
	}
	b.di.Set(name, h)
}

// GetDI fetch a registered dependency injection
func (b *Baa) GetDI(name string) interface{} {
	return b.di.Get(name)
}

// Static set static file route
// h used for set Expries ...
func (b *Baa) Static(prefix string, dir string, index bool, h HandlerFunc) {
	if prefix == "" {
		panic("baa.Static prefix can not be empty")
	}
	if dir == "" {
		panic("baa.Static dir can not be empty")
	}
	b.Get(prefix+"*", newStatic(prefix, dir, index, h))
}

// StaticFile shortcut for serve file
func (b *Baa) StaticFile(pattern string, path string) RouteNode {
	return b.Get(pattern, func(c *Context) {
		if err := serveFile(path, c); err != nil {
			c.Error(err)
		}
	})
}

// SetAutoHead sets the value who determines whether add HEAD method automatically
// when GET method is added. Combo router will not be affected by this value.
func (b *Baa) SetAutoHead(v bool) {
	b.Router().SetAutoHead(v)
}

// SetAutoTrailingSlash optional trailing slash.
func (b *Baa) SetAutoTrailingSlash(v bool) {
	b.Router().SetAutoTrailingSlash(v)
}

// Route is a shortcut for same handlers but different HTTP methods.
//
// Example:
// 		baa.Route("/", "GET,POST", h)
func (b *Baa) Route(pattern, methods string, h ...HandlerFunc) RouteNode {
	var ru RouteNode
	var ms []string
	if methods == "*" {
		for m := range RouterMethods {
			ms = append(ms, m)
		}
	} else {
		ms = strings.Split(methods, ",")
	}
	for _, m := range ms {
		ru = b.Router().Add(strings.TrimSpace(m), pattern, h)
	}
	return ru
}

// Group registers a list of same prefix route
func (b *Baa) Group(pattern string, f func(), h ...HandlerFunc) {
	b.Router().GroupAdd(pattern, f, h)
}

// Any is a shortcut for b.Router().handle("*", pattern, handlers)
func (b *Baa) Any(pattern string, h ...HandlerFunc) RouteNode {
	var ru RouteNode
	for m := range RouterMethods {
		ru = b.Router().Add(m, pattern, h)
	}
	return ru
}

// Delete is a shortcut for b.Route(pattern, "DELETE", handlers)
func (b *Baa) Delete(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("DELETE", pattern, h)
}

// Get is a shortcut for b.Route(pattern, "GET", handlers)
func (b *Baa) Get(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("GET", pattern, h)
}

// Head is a shortcut forb.Route(pattern, "Head", handlers)
func (b *Baa) Head(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("HEAD", pattern, h)
}

// Options is a shortcut for b.Route(pattern, "Options", handlers)
func (b *Baa) Options(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("OPTIONS", pattern, h)
}

// Patch is a shortcut for b.Route(pattern, "PATCH", handlers)
func (b *Baa) Patch(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("PATCH", pattern, h)
}

// Post is a shortcut for b.Route(pattern, "POST", handlers)
func (b *Baa) Post(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("POST", pattern, h)
}

// Put is a shortcut for b.Route(pattern, "Put", handlers)
func (b *Baa) Put(pattern string, h ...HandlerFunc) RouteNode {
	return b.Router().Add("PUT", pattern, h)
}

// Websocket register a websocket router handler
func (b *Baa) Websocket(pattern string, h func(*websocket.Conn)) RouteNode {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:    4096,
		WriteBufferSize:   4096,
		EnableCompression: true,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	return b.Route(pattern, "GET,POST", func(c *Context) {
		conn, err := upgrader.Upgrade(c.Resp, c.Req, nil)
		if err != nil {
			b.Logger().Panicf("websocket upgrade connection error: %v", err)
			return
		}
		h(conn)
	})
}

// SetNotFound set not found route handler
func (b *Baa) SetNotFound(h HandlerFunc) {
	b.notFoundHandler = h
}

// NotFound execute not found handler
func (b *Baa) NotFound(c *Context) {
	if b.notFoundHandler != nil {
		b.notFoundHandler(c)
		return
	}
	http.NotFound(c.Resp, c.Req)
}

// SetError set error handler
func (b *Baa) SetError(h ErrorHandleFunc) {
	b.errorHandler = h
}

// Error execute internal error handler
func (b *Baa) Error(err error, c *Context) {
	if err == nil {
		err = errors.New("Internal Server Error")
	}
	if b.errorHandler != nil {
		b.errorHandler(err, c)
		return
	}
	code := http.StatusInternalServerError
	msg := http.StatusText(code)
	if b.debug {
		msg = err.Error()
	}
	b.Logger().Println(err)
	http.Error(c.Resp, msg, code)
}

// DefaultNotFoundHandler invokes the default HTTP error handler.
func (b *Baa) DefaultNotFoundHandler(c *Context) {
	code := http.StatusNotFound
	msg := http.StatusText(code)
	http.Error(c.Resp, msg, code)
}

// URLFor use named route return format url
func (b *Baa) URLFor(name string, args ...interface{}) string {
	return b.Router().URLFor(name, args...)
}

// wrapMiddleware wraps middleware.
func wrapMiddleware(m Middleware) HandlerFunc {
	switch m := m.(type) {
	case HandlerFunc:
		return m
	case func(*Context):
		return m
	case http.Handler, http.HandlerFunc:
		return WrapHandlerFunc(func(c *Context) {
			m.(http.Handler).ServeHTTP(c.Resp, c.Req)
		})
	case func(http.ResponseWriter, *http.Request):
		return WrapHandlerFunc(func(c *Context) {
			m(c.Resp, c.Req)
		})
	default:
		panic("unknown middleware")
	}
}

// WrapHandlerFunc wrap for context handler chain
func WrapHandlerFunc(h HandlerFunc) HandlerFunc {
	return func(c *Context) {
		h(c)
		c.Next()
	}
}

func init() {
	appInstances = make(map[string]*Baa)
	Env = os.Getenv("BAA_ENV")
	if Env == "" {
		Env = DEV
	}
}
<|endoftext|>