<commit_before>//  Copyright (c) 2019 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 vpp

import (
	"net"
	"testing"

	"github.com/ligato/cn-infra/logging/logrus"

	vpp_interfaces "github.com/ligato/vpp-agent/api/models/vpp/interfaces"
	vpp_l3 "github.com/ligato/vpp-agent/api/models/vpp/l3"
	_ "github.com/ligato/vpp-agent/plugins/vpp/ifplugin"
	"github.com/ligato/vpp-agent/plugins/vpp/ifplugin/ifaceidx"
	ifplugin_vppcalls "github.com/ligato/vpp-agent/plugins/vpp/ifplugin/vppcalls"
	_ "github.com/ligato/vpp-agent/plugins/vpp/l3plugin"
	l3plugin_vppcalls "github.com/ligato/vpp-agent/plugins/vpp/l3plugin/vppcalls"
	"github.com/ligato/vpp-agent/plugins/vpp/l3plugin/vrfidx"
)

func TestLoopbackInterface(t *testing.T) {
	ctx := setupVPP(t)
	defer ctx.teardownVPP()

	h := ifplugin_vppcalls.CompatibleInterfaceVppHandler(ctx.Chan, logrus.NewLogger("test"))

	index, err := h.AddLoopbackInterface("loop1")
	if err != nil {
		t.Fatalf("creating loopback interface failed: %v", err)
	}
	t.Logf("loopback index: %+v", index)

	ifaces, err := h.DumpInterfaces()
	if err != nil {
		t.Fatalf("dumping interfaces failed: %v", err)
	}
	iface, ok := ifaces[index]
	if !ok {
		t.Fatalf("loopback interface not found in dump")
	}
	t.Logf("interface: %+v", iface.Interface)
	if iface.Interface.Name != "loop1" {
		t.Fatalf("expected interface name to be loop1, got %v", iface.Interface.Name)
	}
	if iface.Interface.Type != vpp_interfaces.Interface_SOFTWARE_LOOPBACK {
		t.Fatalf("expected interface type to be loopback, got %v", iface.Interface.Type)
	}
}

func TestRoutes(t *testing.T) {
	ctx := setupVPP(t)
	defer ctx.teardownVPP()

	ifIndexes := ifaceidx.NewIfaceIndex(logrus.NewLogger("test-if"), "test-if")
	vrfIndexes := vrfidx.NewVRFIndex(logrus.NewLogger("test-vrf"), "test-vrf")
	vrfIndexes.Put("vrf1-ipv4", &vrfidx.VRFMetadata{Index: 0, Protocol: vpp_l3.VrfTable_IPV4})
	vrfIndexes.Put("vrf1-ipv6", &vrfidx.VRFMetadata{Index: 0, Protocol: vpp_l3.VrfTable_IPV6})

	h := l3plugin_vppcalls.CompatibleL3VppHandler(ctx.Chan, ifIndexes, vrfIndexes, logrus.NewLogger("test"))

	routes, err := h.DumpRoutes()
	if err != nil {
		t.Fatalf("dumping routes failed: %v", err)
	}
	t.Logf("%d routes dumped", len(routes))

	var hasIPv4, hasIPv6 bool
	for _, route := range routes {
		t.Logf(" - route: %+v", route.Route)

		ip, _, err := net.ParseCIDR(route.Route.DstNetwork)
		if err != nil {
			t.Fatalf("invalid dst network: %v", route.Route.DstNetwork)
		}
		if ip.To4() == nil {
			hasIPv4 = true
		} else {
			hasIPv6 = true
		}
	}

	if !hasIPv4 || !hasIPv6 {
		t.Fatalf("expected dump to contain both IPv4 and IPv6 routes")
	}
}
<commit_msg>added memif interface test to tests/integration (#1412)<commit_after>//  Copyright (c) 2019 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 vpp

import (
	"net"
	"testing"

	"github.com/ligato/cn-infra/logging/logrus"

	vpp_interfaces "github.com/ligato/vpp-agent/api/models/vpp/interfaces"
	vpp_l3 "github.com/ligato/vpp-agent/api/models/vpp/l3"
	_ "github.com/ligato/vpp-agent/plugins/vpp/ifplugin"
	"github.com/ligato/vpp-agent/plugins/vpp/ifplugin/ifaceidx"
	ifplugin_vppcalls "github.com/ligato/vpp-agent/plugins/vpp/ifplugin/vppcalls"
	_ "github.com/ligato/vpp-agent/plugins/vpp/l3plugin"
	l3plugin_vppcalls "github.com/ligato/vpp-agent/plugins/vpp/l3plugin/vppcalls"
	"github.com/ligato/vpp-agent/plugins/vpp/l3plugin/vrfidx"
)

func TestLoopbackInterface(t *testing.T) {
	ctx := setupVPP(t)
	defer ctx.teardownVPP()

	h := ifplugin_vppcalls.CompatibleInterfaceVppHandler(ctx.Chan, logrus.NewLogger("test"))

	index, err := h.AddLoopbackInterface("loop1")
	if err != nil {
		t.Fatalf("creating loopback interface failed: %v", err)
	}
	t.Logf("loopback index: %+v", index)

	ifaces, err := h.DumpInterfaces()
	if err != nil {
		t.Fatalf("dumping interfaces failed: %v", err)
	}
	iface, ok := ifaces[index]
	if !ok {
		t.Fatalf("loopback interface not found in dump")
	}
	t.Logf("interface: %+v", iface.Interface)
	if iface.Interface.Name != "loop1" {
		t.Fatalf("expected interface name to be loop1, got %v", iface.Interface.Name)
	}
	if iface.Interface.Type != vpp_interfaces.Interface_SOFTWARE_LOOPBACK {
		t.Fatalf("expected interface type to be loopback, got %v", iface.Interface.Type)
	}
}

func TestMemifInterface(t *testing.T) {
	ctx := setupVPP(t)
	defer ctx.teardownVPP()

	h := ifplugin_vppcalls.CompatibleInterfaceVppHandler(ctx.Chan, logrus.NewLogger("test"))

	index, err := h.AddMemifInterface("memif1", &vpp_interfaces.MemifLink{
		Id:     1,
		Mode:   vpp_interfaces.MemifLink_IP,
		Secret: "secret",
		Master: true,
        }, 0 )

	if err != nil {
		t.Fatalf("creating memif interface failed: %v", err)
	}
	t.Logf("memif index: %+v", index)

	ifaces, err := h.DumpInterfaces()
	if err != nil {
		t.Fatalf("dumping interfaces failed: %v", err)
	}
	iface, ok := ifaces[index]
	if !ok {
		t.Fatalf("Memif interface not found in dump")
	}
	t.Logf("interface: %+v", iface.Interface)
	if iface.Interface.Name != "memif1" {
		t.Fatalf("expected interface name to be memif1, got %v", iface.Interface.Name)
	}
	if iface.Interface.Type != vpp_interfaces.Interface_MEMIF {
		t.Fatalf("expected interface type to be memif, got %v", iface.Interface.Type)
	}
}

func TestRoutes(t *testing.T) {
	ctx := setupVPP(t)
	defer ctx.teardownVPP()

	ifIndexes := ifaceidx.NewIfaceIndex(logrus.NewLogger("test-if"), "test-if")
	vrfIndexes := vrfidx.NewVRFIndex(logrus.NewLogger("test-vrf"), "test-vrf")
	vrfIndexes.Put("vrf1-ipv4", &vrfidx.VRFMetadata{Index: 0, Protocol: vpp_l3.VrfTable_IPV4})
	vrfIndexes.Put("vrf1-ipv6", &vrfidx.VRFMetadata{Index: 0, Protocol: vpp_l3.VrfTable_IPV6})

	h := l3plugin_vppcalls.CompatibleL3VppHandler(ctx.Chan, ifIndexes, vrfIndexes, logrus.NewLogger("test"))

	routes, err := h.DumpRoutes()
	if err != nil {
		t.Fatalf("dumping routes failed: %v", err)
	}
	t.Logf("%d routes dumped", len(routes))

	var hasIPv4, hasIPv6 bool
	for _, route := range routes {
		t.Logf(" - route: %+v", route.Route)

		ip, _, err := net.ParseCIDR(route.Route.DstNetwork)
		if err != nil {
			t.Fatalf("invalid dst network: %v", route.Route.DstNetwork)
		}
		if ip.To4() == nil {
			hasIPv4 = true
		} else {
			hasIPv6 = true
		}
	}

	if !hasIPv4 || !hasIPv6 {
		t.Fatalf("expected dump to contain both IPv4 and IPv6 routes")
	}
}
<|endoftext|>


<commit_before>package starmap

import (
	"appengine"
	"appengine/memcache"
	"bytes"
	"fmt"
	"geom"
	"image/color"
	"image/png"
	"math"
	"net/http"
	"render"
	"render/style"
	"strings"
)

var smlCircle *style.PointStyle = style.NewPointStyle(0.5, color.White,
	style.CIRCLE)
var midCircle *style.PointStyle = style.NewPointStyle(1, color.White,
	style.CIRCLE)
var lrgCircle *style.PointStyle = style.NewPointStyle(2, color.White,
	style.CIRCLE)
var superCircle *style.PointStyle = style.NewPointStyle(3, color.White,
	style.CIRCLE)

/* create the cache key for a WMS tile */
func createKey(layer string, width, height int, lower,
	upper *geom.Point) string {
	return fmt.Sprintf("%v-%v-%v-%v-%v", layer, width, height, lower, upper)
}

/* WMS getmap handler function */
func getmap(w http.ResponseWriter, r *http.Request) {
	width := intParam("WIDTH", 1024, r)
	height := intParam("HEIGHT", 512, r)
	lower, upper := parseBbox("BBOX", r)
	layer := strParam("LAYERS", "stars", r)
	ctx := appengine.NewContext(r)
	cacheKey := createKey(layer, width, height, lower, upper)
	item, err := memcache.Get(ctx, cacheKey)
	if err == memcache.ErrCacheMiss {
		tile, err := createTile(w, layer, width, height, lower, upper)
		if err != nil {
			doErr(w, err)
			return
		}
		item = &memcache.Item{Key: cacheKey, Value: tile}
		err = memcache.Add(ctx, item)
		if err != nil {
			doErr(w, err)
			return
		}
	} else if err != nil {
		doErr(w, err)
		return
	}

	w.Write(item.Value)
}

/* create a new tile image for request */
func createTile(w http.ResponseWriter, layer string, width, height int,
	lower, upper *geom.Point) ([]byte, error) {
	layer = strings.ToLower(layer)
	if layer == "constellations" {
		return createConstTile(w, width, height, lower, upper)
	} else {
		return createStarTile(w, width, height, lower, upper)
	}
}

/* create a constellation layer tile */
func createConstTile(w http.ResponseWriter, width, height int,
	lower, upper *geom.Point) ([]byte, error) {
	if constelErr != nil {
		return nil, constelErr
	}
	scale := math.Abs(upper.X()-lower.X()) / float64(width)
	txtColor := color.White
	s := style.NewPolyStyle(1, color.White)
	trans := geom.CreateTransform(lower, upper, width, height, geom.STELLAR)
	img := render.CreateTransparent(width, height)
	bbox := geom.NewBBox2D(lower.X(), lower.Y(), upper.X(), upper.Y())
	for _, c := range constelData {
		for _, pi := range c.PolyInfos {
			if bbox.Touches(pi.Geom) {
				render.RenderPoly(img, pi.Geom, trans, s)
				if charsErr == nil && pi.LabelPoint != nil &&
					pi.MaxScale > scale {
					labelPoint := pi.LabelPoint
					pix := trans.TransformXY(labelPoint[0], labelPoint[1])
					render.RenderString(img, chars, 10, pix, c.Name, txtColor)
				}
			}
		}
	}
	var rval bytes.Buffer
	if err := png.Encode(&rval, img); err != nil {
		return nil, err
	}
	return rval.Bytes(), nil
}

/* create a star layer tile */
func createStarTile(w http.ResponseWriter, width, height int,
	lower, upper *geom.Point) ([]byte, error) {
	if dataErr != nil {
		return nil, dataErr
	}
	lowerHash, upperHash := geom.BBoxHash(lower, upper, geom.STELLAR)
	trans := geom.CreateTransform(lower, upper, width, height, geom.STELLAR)
	img := render.Create(width, height, color.Black)
	stars := data.Range(lowerHash, upperHash)
	for _, s := range stars {
		coord, err := geom.UnHash(s.GeoHash, geom.STELLAR)
		if err != nil {
			return nil, err
		}
		pix := trans.Transform(coord)
		mag := s.Magnitude
		style := smlCircle
		var gray uint8
		if mag < -1 {
			style = superCircle
			gray = 255
		} else if mag < 0 {
			style = superCircle
			gray = 200
		} else if mag < 2 {
			style = lrgCircle
			gray = uint8((2.0-mag)*64.0) + 128
		} else if mag < 4 {
			style = midCircle
			gray = uint8((4.0-mag)*64.0) + 128
		} else {
			gray = uint8((6.0-mag)*64.0) + 128
		}
		color := color.RGBA{gray, gray, gray, 255}
		style.Style.Color = color
		render.Render(img, pix, style)
	}
	var rval bytes.Buffer
	if err := png.Encode(&rval, img); err != nil {
		return nil, err
	}
	return rval.Bytes(), nil
}
<commit_msg>labels are color coded<commit_after>package starmap

import (
	"appengine"
	"appengine/memcache"
	"bytes"
	"fmt"
	"geom"
	"image/color"
	"image/png"
	"math"
	"net/http"
	"render"
	"render/style"
	"strings"
)

var smlCircle = style.NewPointStyle(0.5, color.White, style.CIRCLE)
var midCircle = style.NewPointStyle(1, color.White, style.CIRCLE)
var lrgCircle = style.NewPointStyle(2, color.White, style.CIRCLE)
var superCircle = style.NewPointStyle(3, color.White, style.CIRCLE)

var labelColors = map[string]color.Color{
	"Heavenly Waters": color.RGBA{0, 154, 205, 255},
	"Hercules":        color.RGBA{34, 139, 34, 255},
	"Ursa Major":      color.RGBA{100, 149, 237, 255},
	"Perseus":         color.RGBA{225, 58, 58, 255},
	"Orion":           color.RGBA{205, 102, 0, 255},
	"Bayer":           color.RGBA{205, 149, 12, 255},
	"La Caille":       color.RGBA{137, 104, 205, 255},
}

/* create the cache key for a WMS tile */
func createKey(layer string, width, height int, lower,
	upper *geom.Point) string {
	return fmt.Sprintf("%v-%v-%v-%v-%v", layer, width, height, lower, upper)
}

/* WMS getmap handler function */
func getmap(w http.ResponseWriter, r *http.Request) {
	width := intParam("WIDTH", 1024, r)
	height := intParam("HEIGHT", 512, r)
	lower, upper := parseBbox("BBOX", r)
	layer := strParam("LAYERS", "stars", r)
	ctx := appengine.NewContext(r)
	cacheKey := createKey(layer, width, height, lower, upper)
	item, err := memcache.Get(ctx, cacheKey)
	if err == memcache.ErrCacheMiss {
		tile, err := createTile(w, layer, width, height, lower, upper)
		if err != nil {
			doErr(w, err)
			return
		}
		item = &memcache.Item{Key: cacheKey, Value: tile}
		err = memcache.Add(ctx, item)
		if err != nil {
			doErr(w, err)
			return
		}
	} else if err != nil {
		doErr(w, err)
		return
	}

	w.Write(item.Value)
}

/* create a new tile image for request */
func createTile(w http.ResponseWriter, layer string, width, height int,
	lower, upper *geom.Point) ([]byte, error) {
	layer = strings.ToLower(layer)
	if layer == "constellations" {
		return createConstTile(w, width, height, lower, upper)
	} else {
		return createStarTile(w, width, height, lower, upper)
	}
}

/* create a constellation layer tile */
func createConstTile(w http.ResponseWriter, width, height int,
	lower, upper *geom.Point) ([]byte, error) {
	if constelErr != nil {
		return nil, constelErr
	}
	scale := math.Abs(upper.X()-lower.X()) / float64(width)
	s := style.NewPolyStyle(1, color.White)
	trans := geom.CreateTransform(lower, upper, width, height, geom.STELLAR)
	img := render.CreateTransparent(width, height)
	bbox := geom.NewBBox2D(lower.X(), lower.Y(), upper.X(), upper.Y())
	for _, c := range constelData {
		txtColor := labelColors[c.Family]
		if txtColor == nil {
			txtColor = color.White
		}
		for _, pi := range c.PolyInfos {
			if bbox.Touches(pi.Geom) {
				render.RenderPoly(img, pi.Geom, trans, s)
				if charsErr == nil && pi.LabelPoint != nil &&
					pi.MaxScale > scale {
					labelPoint := pi.LabelPoint
					pix := trans.TransformXY(labelPoint[0], labelPoint[1])
					render.RenderString(img, chars, 10, pix, c.Name, txtColor)
				}
			}
		}
	}
	var rval bytes.Buffer
	if err := png.Encode(&rval, img); err != nil {
		return nil, err
	}
	return rval.Bytes(), nil
}

/* create a star layer tile */
func createStarTile(w http.ResponseWriter, width, height int,
	lower, upper *geom.Point) ([]byte, error) {
	if dataErr != nil {
		return nil, dataErr
	}
	lowerHash, upperHash := geom.BBoxHash(lower, upper, geom.STELLAR)
	trans := geom.CreateTransform(lower, upper, width, height, geom.STELLAR)
	img := render.Create(width, height, color.Black)
	stars := data.Range(lowerHash, upperHash)
	for _, s := range stars {
		coord, err := geom.UnHash(s.GeoHash, geom.STELLAR)
		if err != nil {
			return nil, err
		}
		pix := trans.Transform(coord)
		mag := s.Magnitude
		style := smlCircle
		var gray uint8
		if mag < -1 {
			style = superCircle
			gray = 255
		} else if mag < 0 {
			style = superCircle
			gray = 200
		} else if mag < 2 {
			style = lrgCircle
			gray = uint8((2.0-mag)*64.0) + 128
		} else if mag < 4 {
			style = midCircle
			gray = uint8((4.0-mag)*64.0) + 128
		} else {
			gray = uint8((6.0-mag)*64.0) + 128
		}
		color := color.RGBA{gray, gray, gray, 255}
		style.Style.Color = color
		render.Render(img, pix, style)
	}
	var rval bytes.Buffer
	if err := png.Encode(&rval, img); err != nil {
		return nil, err
	}
	return rval.Bytes(), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2020-2022 Snowflake Computing Inc. All rights reserved.

package gosnowflake

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"testing"
)

func openDB(t *testing.T) *sql.DB {
	var db *sql.DB
	var err error

	if db, err = sql.Open("snowflake", dsn); err != nil {
		t.Fatalf("failed to open db. %v, err: %v", dsn, err)
	}
	return db
}

func TestGetQueryID(t *testing.T) {
	db := openDB(t)
	defer db.Close()

	ctx := context.TODO()
	conn, err := db.Conn(ctx)
	if err != nil {
		t.Error(err)
	}

	if err = conn.Raw(func(x interface{}) error {
		rows, err := x.(driver.QueryerContext).QueryContext(ctx, "select 1", nil)
		if err != nil {
			return err
		}
		defer rows.Close()

		if _, err = x.(driver.QueryerContext).QueryContext(ctx, "selectt 1", nil); err == nil {
			t.Fatal("should have failed to execute query")
		}
		if driverErr, ok := err.(*SnowflakeError); ok {
			if driverErr.Number != 1003 {
				t.Fatalf("incorrect error code. expected: 1003, got: %v", driverErr.Number)
			}
			if driverErr.QueryID == "" {
				t.Fatal("should have an associated query ID")
			}
		} else {
			t.Fatal("should have been able to cast to Snowflake Error")
		}
		return nil
	}); err != nil {
		t.Fatalf("failed to prepare statement. err: %v", err)
	}
}

func TestEmitQueryID(t *testing.T) {
	queryIDChan := make(chan string, 1)
	numrows := 100000
	ctx := WithAsyncMode(context.Background())
	ctx = WithQueryIDChan(ctx, queryIDChan)

	goRoutineChan := make(chan string)
	go func(grCh chan string, qIDch chan string) {
		queryID := <-queryIDChan
		grCh <- queryID
	}(goRoutineChan, queryIDChan)

	cnt := 0
	var idx int
	var v string
	runTests(t, dsn, func(dbt *DBTest) {
		rows := dbt.mustQueryContext(ctx, fmt.Sprintf(selectRandomGenerator, numrows))
		defer rows.Close()

		for rows.Next() {
			if err := rows.Scan(&idx, &v); err != nil {
				t.Fatal(err)
			}
			cnt++
		}
		logger.Infof("NextResultSet: %v", rows.NextResultSet())
	})

	queryID := <-goRoutineChan
	if queryID == "" {
		t.Fatal("expected a nonempty query ID")
	}
	if cnt != numrows {
		t.Errorf("number of rows didn't match. expected: %v, got: %v", numrows, cnt)
	}
}

// End-to-end test to fetch result with queryID
func TestE2EFetchResultByID(t *testing.T) {
	db := openDB(t)
	defer db.Close()

	if _, err := db.Exec(`create or replace table test_fetch_result(c1 number,
		c2 string) as select 10, 'z'`); err != nil {
		t.Fatalf("failed to create table: %v", err)
	}

	ctx := context.Background()
	conn, err := db.Conn(ctx)
	if err != nil {
		t.Error(err)
	}
	if err = conn.Raw(func(x interface{}) error {
		stmt, err := x.(driver.ConnPrepareContext).PrepareContext(ctx, "select * from test_fetch_result")
		if err != nil {
			return err
		}

		rows1, err := stmt.(driver.StmtQueryContext).QueryContext(ctx, nil)
		if err != nil {
			return err
		}
		qid := rows1.(SnowflakeResult).GetQueryID()

		newCtx := context.WithValue(context.Background(), fetchResultByID, qid)
		rows2, err := db.QueryContext(newCtx, "")
		if err != nil {
			t.Fatalf("Fetch Query Result by ID failed: %v", err)
		}
		var c1 sql.NullInt64
		var c2 sql.NullString
		for rows2.Next() {
			err = rows2.Scan(&c1, &c2)
		}
		if c1.Int64 != 10 || c2.String != "z" {
			t.Fatalf("Query result is not expected: %v", err)
		}
		return nil
	}); err != nil {
		t.Fatalf("failed to drop table: %v", err)
	}

	if _, err := db.Exec("drop table if exists test_fetch_result"); err != nil {
		t.Fatalf("failed to drop table: %v", err)
	}
}

func TestWithDescribeOnly(t *testing.T) {
	runTests(t, dsn, func(dbt *DBTest) {
		ctx := WithDescribeOnly(context.Background())
		rows := dbt.mustQueryContext(ctx, selectVariousTypes)
		cols, err := rows.Columns()
		if err != nil {
			t.Error(err)
		}
		types, err := rows.ColumnTypes()
		if err != nil {
			t.Error(err)
		}
		for i, col := range cols {
			if types[i].Name() != col {
				t.Fatalf("column name mismatch. expected: %v, got: %v", col, types[i].Name())
			}
		}
		if rows.Next() {
			t.Fatal("there should not be any rows in describe only mode")
		}
	})
}

func TestCallStatement(t *testing.T) {
	runTests(t, dsn, func(dbt *DBTest) {
		in1 := float64(1)
		in2 := string("[2,3]")
		expected := "1 \"[2,3]\" [2,3]"
		var out string

		dbt.db.Exec("ALTER SESSION SET USE_STATEMENT_TYPE_CALL_FOR_STORED_PROC_CALLS = true")

		dbt.mustExec("create or replace procedure " +
			"TEST_SP_CALL_STMT_ENABLED(in1 float, in2 variant) " +
			"returns string language javascript as $$ " +
			"let res = snowflake.execute({sqlText: 'select ? c1, ? c2', binds:[IN1, JSON.stringify(IN2)]}); " +
			"res.next(); " +
			"return res.getColumnValueAsString(1) + ' ' + res.getColumnValueAsString(2) + ' ' + IN2; " +
			"$$;")

		stmt, err := dbt.db.Prepare("call TEST_SP_CALL_STMT_ENABLED(?, to_variant(?))")
		if err != nil {
			dbt.Errorf("failed to prepare query: %v", err)
		}
		defer stmt.Close()
		err = stmt.QueryRow(in1, in2).Scan(&out)
		if err != nil {
			dbt.Errorf("failed to scan: %v", err)
		}

		if expected != out {
			dbt.Errorf("expected: %s, got: %s", expected, out)
		}

		dbt.mustExec("drop procedure if exists TEST_SP_CALL_STMT_ENABLED(float, variant)")
	})
}
<commit_msg>Disable the test for new statement type CALL until server supports the parameter USE_STATEMENT_TYPE_CALL_FOR_STORED_PROC_CALLS (#622)<commit_after>// Copyright (c) 2020-2022 Snowflake Computing Inc. All rights reserved.

package gosnowflake

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"testing"
)

func openDB(t *testing.T) *sql.DB {
	var db *sql.DB
	var err error

	if db, err = sql.Open("snowflake", dsn); err != nil {
		t.Fatalf("failed to open db. %v, err: %v", dsn, err)
	}
	return db
}

func TestGetQueryID(t *testing.T) {
	db := openDB(t)
	defer db.Close()

	ctx := context.TODO()
	conn, err := db.Conn(ctx)
	if err != nil {
		t.Error(err)
	}

	if err = conn.Raw(func(x interface{}) error {
		rows, err := x.(driver.QueryerContext).QueryContext(ctx, "select 1", nil)
		if err != nil {
			return err
		}
		defer rows.Close()

		if _, err = x.(driver.QueryerContext).QueryContext(ctx, "selectt 1", nil); err == nil {
			t.Fatal("should have failed to execute query")
		}
		if driverErr, ok := err.(*SnowflakeError); ok {
			if driverErr.Number != 1003 {
				t.Fatalf("incorrect error code. expected: 1003, got: %v", driverErr.Number)
			}
			if driverErr.QueryID == "" {
				t.Fatal("should have an associated query ID")
			}
		} else {
			t.Fatal("should have been able to cast to Snowflake Error")
		}
		return nil
	}); err != nil {
		t.Fatalf("failed to prepare statement. err: %v", err)
	}
}

func TestEmitQueryID(t *testing.T) {
	queryIDChan := make(chan string, 1)
	numrows := 100000
	ctx := WithAsyncMode(context.Background())
	ctx = WithQueryIDChan(ctx, queryIDChan)

	goRoutineChan := make(chan string)
	go func(grCh chan string, qIDch chan string) {
		queryID := <-queryIDChan
		grCh <- queryID
	}(goRoutineChan, queryIDChan)

	cnt := 0
	var idx int
	var v string
	runTests(t, dsn, func(dbt *DBTest) {
		rows := dbt.mustQueryContext(ctx, fmt.Sprintf(selectRandomGenerator, numrows))
		defer rows.Close()

		for rows.Next() {
			if err := rows.Scan(&idx, &v); err != nil {
				t.Fatal(err)
			}
			cnt++
		}
		logger.Infof("NextResultSet: %v", rows.NextResultSet())
	})

	queryID := <-goRoutineChan
	if queryID == "" {
		t.Fatal("expected a nonempty query ID")
	}
	if cnt != numrows {
		t.Errorf("number of rows didn't match. expected: %v, got: %v", numrows, cnt)
	}
}

// End-to-end test to fetch result with queryID
func TestE2EFetchResultByID(t *testing.T) {
	db := openDB(t)
	defer db.Close()

	if _, err := db.Exec(`create or replace table test_fetch_result(c1 number,
		c2 string) as select 10, 'z'`); err != nil {
		t.Fatalf("failed to create table: %v", err)
	}

	ctx := context.Background()
	conn, err := db.Conn(ctx)
	if err != nil {
		t.Error(err)
	}
	if err = conn.Raw(func(x interface{}) error {
		stmt, err := x.(driver.ConnPrepareContext).PrepareContext(ctx, "select * from test_fetch_result")
		if err != nil {
			return err
		}

		rows1, err := stmt.(driver.StmtQueryContext).QueryContext(ctx, nil)
		if err != nil {
			return err
		}
		qid := rows1.(SnowflakeResult).GetQueryID()

		newCtx := context.WithValue(context.Background(), fetchResultByID, qid)
		rows2, err := db.QueryContext(newCtx, "")
		if err != nil {
			t.Fatalf("Fetch Query Result by ID failed: %v", err)
		}
		var c1 sql.NullInt64
		var c2 sql.NullString
		for rows2.Next() {
			err = rows2.Scan(&c1, &c2)
		}
		if c1.Int64 != 10 || c2.String != "z" {
			t.Fatalf("Query result is not expected: %v", err)
		}
		return nil
	}); err != nil {
		t.Fatalf("failed to drop table: %v", err)
	}

	if _, err := db.Exec("drop table if exists test_fetch_result"); err != nil {
		t.Fatalf("failed to drop table: %v", err)
	}
}

func TestWithDescribeOnly(t *testing.T) {
	runTests(t, dsn, func(dbt *DBTest) {
		ctx := WithDescribeOnly(context.Background())
		rows := dbt.mustQueryContext(ctx, selectVariousTypes)
		cols, err := rows.Columns()
		if err != nil {
			t.Error(err)
		}
		types, err := rows.ColumnTypes()
		if err != nil {
			t.Error(err)
		}
		for i, col := range cols {
			if types[i].Name() != col {
				t.Fatalf("column name mismatch. expected: %v, got: %v", col, types[i].Name())
			}
		}
		if rows.Next() {
			t.Fatal("there should not be any rows in describe only mode")
		}
	})
}

func TestCallStatement(t *testing.T) {
	t.Skip("USE_STATEMENT_TYPE_CALL_FOR_STORED_PROC_CALLS is not yet supported")

	runTests(t, dsn, func(dbt *DBTest) {
		in1 := float64(1)
		in2 := string("[2,3]")
		expected := "1 \"[2,3]\" [2,3]"
		var out string

		dbt.db.Exec("ALTER SESSION SET USE_STATEMENT_TYPE_CALL_FOR_STORED_PROC_CALLS = true")

		dbt.mustExec("create or replace procedure " +
			"TEST_SP_CALL_STMT_ENABLED(in1 float, in2 variant) " +
			"returns string language javascript as $$ " +
			"let res = snowflake.execute({sqlText: 'select ? c1, ? c2', binds:[IN1, JSON.stringify(IN2)]}); " +
			"res.next(); " +
			"return res.getColumnValueAsString(1) + ' ' + res.getColumnValueAsString(2) + ' ' + IN2; " +
			"$$;")

		stmt, err := dbt.db.Prepare("call TEST_SP_CALL_STMT_ENABLED(?, to_variant(?))")
		if err != nil {
			dbt.Errorf("failed to prepare query: %v", err)
		}
		defer stmt.Close()
		err = stmt.QueryRow(in1, in2).Scan(&out)
		if err != nil {
			dbt.Errorf("failed to scan: %v", err)
		}

		if expected != out {
			dbt.Errorf("expected: %s, got: %s", expected, out)
		}

		dbt.mustExec("drop procedure if exists TEST_SP_CALL_STMT_ENABLED(float, variant)")
	})
}
<|endoftext|>


<commit_before>package test

import (
	"fmt"
	"os"
	"os/exec"
	"testing"

	log "github.com/Sirupsen/logrus"
	"github.com/stretchr/testify/assert"

	"github.com/libopenstorage/kvdb"
	"github.com/libopenstorage/kvdb/mem"
	"github.com/libopenstorage/openstorage/api"
	"github.com/libopenstorage/openstorage/volume"
)

// Context maintains current device state. It gets passed into tests
// so that tests can build on other tests' work
type Context struct {
	volume.VolumeDriver
	volID      api.VolumeID
	snapID     api.SnapID
	mountPath  string
	tgtPath    string
	devicePath string
	Filesystem string
}

func NewContext(d volume.VolumeDriver) *Context {
	return &Context{
		VolumeDriver: d,
		volID:        api.BadVolumeID,
		snapID:       api.BadSnapID,
		Filesystem:   string(""),
	}
}

func RunShort(t *testing.T, ctx *Context) {
	create(t, ctx)
	inspect(t, ctx)
	enumerate(t, ctx)
	format(t, ctx)
	attach(t, ctx)
	mount(t, ctx)
	io(t, ctx)
	unmount(t, ctx)
	detach(t, ctx)
	delete(t, ctx)
	runEnd(t, ctx)
}

func Run(t *testing.T, ctx *Context) {
	RunShort(t, ctx)
	RunSnap(t, ctx)
	runEnd(t, ctx)
}

func runEnd(t *testing.T, ctx *Context) {
	detach(t, ctx)
	shutdown(t, ctx)
}

func RunSnap(t *testing.T, ctx *Context) {
	snap(t, ctx)
	snapInspect(t, ctx)
	snapEnumerate(t, ctx)
	snapDiff(t, ctx)
	snapDelete(t, ctx)
}

func create(t *testing.T, ctx *Context) {
	fmt.Println("create")

	volID, err := ctx.Create(
		api.VolumeLocator{Name: "foo"},
		&api.CreateOptions{FailIfExists: false},
		&api.VolumeSpec{
			Size:    1 * 1024 * 1024 * 1024,
			HALevel: 1,
			Format:  api.Filesystem(ctx.Filesystem),
		})

	assert.NoError(t, err, "Failed in Create")
	ctx.volID = volID
}

func inspect(t *testing.T, ctx *Context) {
	fmt.Println("inspect")

	vols, err := ctx.Inspect([]api.VolumeID{ctx.volID})
	assert.NoError(t, err, "Failed in Inspect")
	assert.NotNil(t, vols, "Nil vols")
	assert.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	assert.Equal(t, vols[0].ID, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].ID)

	vols, err = ctx.Inspect([]api.VolumeID{api.VolumeID("shouldNotExist")})
	assert.Equal(t, 0, len(vols), "Expect 0 volume actual %v volumes", len(vols))
}

func enumerate(t *testing.T, ctx *Context) {
	fmt.Println("enumerate")

	vols, err := ctx.Enumerate(api.VolumeLocator{}, nil)
	assert.NoError(t, err, "Failed in Enumerate")
	assert.NotNil(t, vols, "Nil vols")
	assert.Equal(t, 1, len(vols), "Expect 1 volume actual %v volumes", len(vols))
	assert.Equal(t, vols[0].ID, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].ID)

	vols, err = ctx.Enumerate(api.VolumeLocator{Name: "foo"}, nil)
	assert.NoError(t, err, "Failed in Enumerate")
	assert.NotNil(t, vols, "Nil vols")
	assert.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	assert.Equal(t, vols[0].ID, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].ID)

	vols, err = ctx.Enumerate(api.VolumeLocator{Name: "shouldNotExist"}, nil)
	assert.Equal(t, len(vols), 0, "Expect 0 volume actual %v volumes", len(vols))
}

func format(t *testing.T, ctx *Context) {
	fmt.Println("format")

	err := ctx.Format(ctx.volID)
	if err != nil {
		assert.Equal(t, err, volume.ErrNotSupported, "Error on format %v", err)
	}
}

func attach(t *testing.T, ctx *Context) {
	fmt.Println("attach")
	p, err := ctx.Attach(ctx.volID)
	if err != nil {
		assert.Equal(t, err, volume.ErrNotSupported, "Error on attach %v", err)
	}
	ctx.devicePath = p

	p, err = ctx.Attach(ctx.volID)
	if err == nil {
		assert.Equal(t, p, ctx.devicePath, "Multiple calls to attach if not errored should return the same path")
	}
}

func detach(t *testing.T, ctx *Context) {
	fmt.Println("detach")

	err := ctx.Detach(ctx.volID)
	if err != nil {
		assert.Equal(t, ctx.devicePath, "", "Error on detach %s: %v", ctx.devicePath, err)
	}
	ctx.devicePath = ""

	err = ctx.Detach(ctx.volID)
	assert.Error(t, err, "Detaching an already detached device should fail")
}

func mount(t *testing.T, ctx *Context) {
	fmt.Println("mount")

	mountPath := "/mnt/voltest"
	err := os.MkdirAll(mountPath, 0755)

	tgtPath := "/mnt/foo"
	err = os.MkdirAll(tgtPath, 0755)
	assert.NoError(t, err, "Failed in mkdir")

	err = ctx.Mount(ctx.volID, tgtPath)
	assert.NoError(t, err, "Failed in mount")

	ctx.mountPath = mountPath
	ctx.tgtPath = tgtPath
}

func unmount(t *testing.T, ctx *Context) {
	fmt.Println("unmount")

	assert.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	err := ctx.Unmount(ctx.volID, ctx.mountPath)
	assert.NoError(t, err, "Failed in unmount")

	ctx.mountPath = ""
	ctx.tgtPath = ""
}

func shutdown(t *testing.T, ctx *Context) {
	fmt.Println("shutdown")
	ctx.Shutdown()
}

func io(t *testing.T, ctx *Context) {
	fmt.Println("io")
	assert.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	cmd := exec.Command("dd", "if=/dev/urandom", "of=/tmp/xx", "bs=1M", "count=10")
	err := cmd.Run()
	assert.NoError(t, err, "Failed to run dd")

	cmd = exec.Command("dd", "if=/tmp/xx", fmt.Sprintf("of=%s/xx", ctx.mountPath))
	err = cmd.Run()
	assert.NoError(t, err, "Failed to run dd on mountpoint %s/xx", ctx.mountPath)

	cmd = exec.Command("diff", "if=/tmp/xx", fmt.Sprintf("of=%s/xx", ctx.mountPath))
	assert.NoError(t, err, "data mismatch")
}

func detachBad(t *testing.T, ctx *Context) {
	err := ctx.Detach(ctx.volID)
	assert.True(t, (err == nil || err == volume.ErrNotSupported),
		"Detach on mounted device should fail")
}

func deleteBad(t *testing.T, ctx *Context) {
	fmt.Println("deleteBad")
	assert.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	err := ctx.Delete(ctx.volID)
	assert.Error(t, err, "Delete on mounted device must fail")
}

func delete(t *testing.T, ctx *Context) {
	fmt.Println("delete")
	err := ctx.Delete(ctx.volID)
	assert.NoError(t, err, "Delete failed")
	ctx.volID = api.BadVolumeID
}

func snap(t *testing.T, ctx *Context) {
	fmt.Println("snap")
	if ctx.volID == api.BadVolumeID {
		create(t, ctx)
	}
	assert.NotEqual(t, ctx.volID, api.BadVolumeID, "invalid volume ID")
	id, err := ctx.Snapshot(ctx.volID, api.Labels{"oh": "snap"})
	assert.NoError(t, err, "Failed in creating a snapshot")
	ctx.snapID = id
}

func snapInspect(t *testing.T, ctx *Context) {
	fmt.Println("snapInspect")

	snaps, err := ctx.SnapInspect([]api.SnapID{ctx.snapID})
	assert.NoError(t, err, "Failed in Inspect")
	assert.NotNil(t, snaps, "Nil snaps")
	assert.Equal(t, len(snaps), 1, "Expect 1 snaps actual %v snaps", len(snaps))
	assert.Equal(t, snaps[0].ID, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].ID)

	snaps, err = ctx.SnapInspect([]api.SnapID{api.SnapID("shouldNotExist")})
	assert.Equal(t, 0, len(snaps), "Expect 0 snaps actual %v snaps", len(snaps))
}

func snapEnumerate(t *testing.T, ctx *Context) {
	fmt.Println("snapEnumerate")

	snaps, err := ctx.SnapEnumerate(nil, nil)
	assert.NoError(t, err, "Failed in snapEnumerate")
	assert.NotNil(t, snaps, "Nil snaps")
	assert.Equal(t, 1, len(snaps), "Expect 1 snaps actual %v snaps", len(snaps))
	assert.Equal(t, snaps[0].ID, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].ID)
	labels := snaps[0].SnapLabels

	snaps, err = ctx.SnapEnumerate([]api.VolumeID{ctx.volID}, nil)
	assert.NoError(t, err, "Failed in snapEnumerate")
	assert.NotNil(t, snaps, "Nil snaps")
	assert.Equal(t, len(snaps), 1, "Expect 1 snap actual %v snaps", len(snaps))
	assert.Equal(t, snaps[0].ID, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].ID)

	snaps, err = ctx.SnapEnumerate([]api.VolumeID{api.VolumeID("shouldNotExist")}, nil)
	assert.Equal(t, len(snaps), 0, "Expect 0 snap actual %v snaps", len(snaps))

	snaps, err = ctx.SnapEnumerate(nil, labels)
	assert.NoError(t, err, "Failed in snapEnumerate")
	assert.NotNil(t, snaps, "Nil snaps")
	assert.Equal(t, len(snaps), 1, "Expect 1 snap actual %v snaps", len(snaps))
	assert.Equal(t, snaps[0].ID, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].ID)
}

func snapDiff(t *testing.T, ctx *Context) {
	fmt.Println("snapDiff")
}

func snapDelete(t *testing.T, ctx *Context) {
	fmt.Println("snapDelete")
}

func init() {
	kv, err := kvdb.New(mem.Name, "driver_test", []string{}, nil)
	if err != nil {
		log.Panicf("Failed to intialize KVDB")
	}
	err = kvdb.SetInstance(kv)
	if err != nil {
		log.Panicf("Failed to set KVDB instance")
	}
}
<commit_msg>do not require detach to return error on already detached vol<commit_after>package test

import (
	"fmt"
	"os"
	"os/exec"
	"testing"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/stretchr/testify/assert"

	"github.com/libopenstorage/kvdb"
	"github.com/libopenstorage/kvdb/mem"
	"github.com/libopenstorage/openstorage/api"
	"github.com/libopenstorage/openstorage/volume"
)

// Context maintains current device state. It gets passed into tests
// so that tests can build on other tests' work
type Context struct {
	volume.VolumeDriver
	volID      api.VolumeID
	snapID     api.SnapID
	mountPath  string
	tgtPath    string
	devicePath string
	Filesystem string
}

func NewContext(d volume.VolumeDriver) *Context {
	return &Context{
		VolumeDriver: d,
		volID:        api.BadVolumeID,
		snapID:       api.BadSnapID,
		Filesystem:   string(""),
	}
}

func RunShort(t *testing.T, ctx *Context) {
	create(t, ctx)
	inspect(t, ctx)
	enumerate(t, ctx)
	attach(t, ctx)
	format(t, ctx)
	mount(t, ctx)
	io(t, ctx)
	unmount(t, ctx)
	detach(t, ctx)
	delete(t, ctx)
	runEnd(t, ctx)
}

func Run(t *testing.T, ctx *Context) {
	RunShort(t, ctx)
	RunSnap(t, ctx)
	runEnd(t, ctx)
}

func runEnd(t *testing.T, ctx *Context) {
	time.Sleep(time.Second * 2)
	shutdown(t, ctx)
}

func RunSnap(t *testing.T, ctx *Context) {
	snap(t, ctx)
	snapInspect(t, ctx)
	snapEnumerate(t, ctx)
	snapDiff(t, ctx)
	snapDelete(t, ctx)
}

func create(t *testing.T, ctx *Context) {
	fmt.Println("create")

	volID, err := ctx.Create(
		api.VolumeLocator{Name: "foo"},
		&api.CreateOptions{FailIfExists: false},
		&api.VolumeSpec{
			Size:    1 * 1024 * 1024 * 1024,
			HALevel: 1,
			Format:  api.Filesystem(ctx.Filesystem),
		})

	assert.NoError(t, err, "Failed in Create")
	ctx.volID = volID
}

func inspect(t *testing.T, ctx *Context) {
	fmt.Println("inspect")

	vols, err := ctx.Inspect([]api.VolumeID{ctx.volID})
	assert.NoError(t, err, "Failed in Inspect")
	assert.NotNil(t, vols, "Nil vols")
	assert.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	assert.Equal(t, vols[0].ID, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].ID)

	vols, err = ctx.Inspect([]api.VolumeID{api.VolumeID("shouldNotExist")})
	assert.Equal(t, 0, len(vols), "Expect 0 volume actual %v volumes", len(vols))
}

func enumerate(t *testing.T, ctx *Context) {
	fmt.Println("enumerate")

	vols, err := ctx.Enumerate(api.VolumeLocator{}, nil)
	assert.NoError(t, err, "Failed in Enumerate")
	assert.NotNil(t, vols, "Nil vols")
	assert.Equal(t, 1, len(vols), "Expect 1 volume actual %v volumes", len(vols))
	assert.Equal(t, vols[0].ID, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].ID)

	vols, err = ctx.Enumerate(api.VolumeLocator{Name: "foo"}, nil)
	assert.NoError(t, err, "Failed in Enumerate")
	assert.NotNil(t, vols, "Nil vols")
	assert.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	assert.Equal(t, vols[0].ID, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].ID)

	vols, err = ctx.Enumerate(api.VolumeLocator{Name: "shouldNotExist"}, nil)
	assert.Equal(t, len(vols), 0, "Expect 0 volume actual %v volumes", len(vols))
}

func format(t *testing.T, ctx *Context) {
	fmt.Println("format")

	err := ctx.Format(ctx.volID)
	if err != nil {
		assert.Equal(t, err, volume.ErrNotSupported, "Error on format %v", err)
	}
}

func attach(t *testing.T, ctx *Context) {
	fmt.Println("attach")
	p, err := ctx.Attach(ctx.volID)
	if err != nil {
		assert.Equal(t, err, volume.ErrNotSupported, "Error on attach %v", err)
	}
	ctx.devicePath = p

	p, err = ctx.Attach(ctx.volID)
	if err == nil {
		assert.Equal(t, p, ctx.devicePath, "Multiple calls to attach if not errored should return the same path")
	}
}

func detach(t *testing.T, ctx *Context) {
	fmt.Println("detach")
	err := ctx.Detach(ctx.volID)
	if err != nil {
		assert.Equal(t, ctx.devicePath, "", "Error on detach %s: %v", ctx.devicePath, err)
	}
	ctx.devicePath = ""
}

func mount(t *testing.T, ctx *Context) {
	fmt.Println("mount")

	mountPath := "/mnt/voltest"
	err := os.MkdirAll(mountPath, 0755)

	tgtPath := "/mnt/foo"
	err = os.MkdirAll(tgtPath, 0755)
	assert.NoError(t, err, "Failed in mkdir")

	err = ctx.Mount(ctx.volID, tgtPath)
	assert.NoError(t, err, "Failed in mount")

	ctx.mountPath = mountPath
	ctx.tgtPath = tgtPath
}

func unmount(t *testing.T, ctx *Context) {
	fmt.Println("unmount")

	assert.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	err := ctx.Unmount(ctx.volID, ctx.mountPath)
	assert.NoError(t, err, "Failed in unmount")

	ctx.mountPath = ""
	ctx.tgtPath = ""
}

func shutdown(t *testing.T, ctx *Context) {
	fmt.Println("shutdown")
	ctx.Shutdown()
}

func io(t *testing.T, ctx *Context) {
	fmt.Println("io")
	assert.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	cmd := exec.Command("dd", "if=/dev/urandom", "of=/tmp/xx", "bs=1M", "count=10")
	err := cmd.Run()
	assert.NoError(t, err, "Failed to run dd")

	cmd = exec.Command("dd", "if=/tmp/xx", fmt.Sprintf("of=%s/xx", ctx.mountPath))
	err = cmd.Run()
	assert.NoError(t, err, "Failed to run dd on mountpoint %s/xx", ctx.mountPath)

	cmd = exec.Command("diff", "if=/tmp/xx", fmt.Sprintf("of=%s/xx", ctx.mountPath))
	assert.NoError(t, err, "data mismatch")
}

func detachBad(t *testing.T, ctx *Context) {
	err := ctx.Detach(ctx.volID)
	assert.True(t, (err == nil || err == volume.ErrNotSupported),
		"Detach on mounted device should fail")
}

func deleteBad(t *testing.T, ctx *Context) {
	fmt.Println("deleteBad")
	assert.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	err := ctx.Delete(ctx.volID)
	assert.Error(t, err, "Delete on mounted device must fail")
}

func delete(t *testing.T, ctx *Context) {
	fmt.Println("delete")
	err := ctx.Delete(ctx.volID)
	assert.NoError(t, err, "Delete failed")
	ctx.volID = api.BadVolumeID
}

func snap(t *testing.T, ctx *Context) {
	fmt.Println("snap")
	if ctx.volID == api.BadVolumeID {
		create(t, ctx)
	}
	assert.NotEqual(t, ctx.volID, api.BadVolumeID, "invalid volume ID")
	id, err := ctx.Snapshot(ctx.volID, api.Labels{"oh": "snap"})
	assert.NoError(t, err, "Failed in creating a snapshot")
	ctx.snapID = id
}

func snapInspect(t *testing.T, ctx *Context) {
	fmt.Println("snapInspect")

	snaps, err := ctx.SnapInspect([]api.SnapID{ctx.snapID})
	assert.NoError(t, err, "Failed in Inspect")
	assert.NotNil(t, snaps, "Nil snaps")
	assert.Equal(t, len(snaps), 1, "Expect 1 snaps actual %v snaps", len(snaps))
	assert.Equal(t, snaps[0].ID, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].ID)

	snaps, err = ctx.SnapInspect([]api.SnapID{api.SnapID("shouldNotExist")})
	assert.Equal(t, 0, len(snaps), "Expect 0 snaps actual %v snaps", len(snaps))
}

func snapEnumerate(t *testing.T, ctx *Context) {
	fmt.Println("snapEnumerate")

	snaps, err := ctx.SnapEnumerate(nil, nil)
	assert.NoError(t, err, "Failed in snapEnumerate")
	assert.NotNil(t, snaps, "Nil snaps")
	assert.Equal(t, 1, len(snaps), "Expect 1 snaps actual %v snaps", len(snaps))
	assert.Equal(t, snaps[0].ID, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].ID)
	labels := snaps[0].SnapLabels

	snaps, err = ctx.SnapEnumerate([]api.VolumeID{ctx.volID}, nil)
	assert.NoError(t, err, "Failed in snapEnumerate")
	assert.NotNil(t, snaps, "Nil snaps")
	assert.Equal(t, len(snaps), 1, "Expect 1 snap actual %v snaps", len(snaps))
	assert.Equal(t, snaps[0].ID, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].ID)

	snaps, err = ctx.SnapEnumerate([]api.VolumeID{api.VolumeID("shouldNotExist")}, nil)
	assert.Equal(t, len(snaps), 0, "Expect 0 snap actual %v snaps", len(snaps))

	snaps, err = ctx.SnapEnumerate(nil, labels)
	assert.NoError(t, err, "Failed in snapEnumerate")
	assert.NotNil(t, snaps, "Nil snaps")
	assert.Equal(t, len(snaps), 1, "Expect 1 snap actual %v snaps", len(snaps))
	assert.Equal(t, snaps[0].ID, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].ID)
}

func snapDiff(t *testing.T, ctx *Context) {
	fmt.Println("snapDiff")
}

func snapDelete(t *testing.T, ctx *Context) {
	fmt.Println("snapDelete")
}

func init() {
	kv, err := kvdb.New(mem.Name, "driver_test", []string{}, nil)
	if err != nil {
		log.Panicf("Failed to intialize KVDB")
	}
	err = kvdb.SetInstance(kv)
	if err != nil {
		log.Panicf("Failed to set KVDB instance")
	}
}
<|endoftext|>


<commit_before>package storage

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"sync"
	"time"

	"github.com/influxdata/influxql"
	"github.com/influxdata/platform/models"
	"github.com/influxdata/platform/tsdb"
	"github.com/influxdata/platform/tsdb/tsi1"
	"github.com/influxdata/platform/tsdb/tsm1"
	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/zap"
)

// Static objects to prevent small allocs.
var timeBytes = []byte("time")

// ErrEngineClosed is returned when a caller attempts to use the engine while
// it's closed.
var ErrEngineClosed = errors.New("engine is closed")

type Engine struct {
	config   Config
	path     string
	engineID *int // Not used by default.
	nodeID   *int // Not used by default.

	mu                sync.RWMutex
	closing           chan struct{} //closing returns the zero value when the engine is shutting down.
	index             *tsi1.Index
	sfile             *tsdb.SeriesFile
	engine            *tsm1.Engine
	retentionEnforcer *retentionEnforcer

	// Tracks all goroutines started by the Engine.
	wg sync.WaitGroup

	logger *zap.Logger
}

// Option provides a set
type Option func(*Engine)

// WithTSMFilenameFormatter sets a function on the underlying tsm1.Engine to specify
// how TSM files are named.
var WithTSMFilenameFormatter = func(fn tsm1.FormatFileNameFunc) Option {
	return func(e *Engine) {
		e.engine.WithFormatFileNameFunc(fn)
	}
}

// WithEngineID sets an engine id, which can be useful for logging when multiple
// engines are in use.
var WithEngineID = func(id int) Option {
	return func(e *Engine) {
		e.engineID = &id
	}
}

// WithNodeID sets a node id on the engine, which can be useful for logging
// when a system has engines running on multiple nodes.
var WithNodeID = func(id int) Option {
	return func(e *Engine) {
		e.nodeID = &id
	}
}

// WithRetentionEnforcer initialises a retention enforcer on the engine.
// WithRetentionEnforcer must be called after other options to ensure that all
// metrics are labelled correctly.
var WithRetentionEnforcer = func(finder BucketFinder) Option {
	return func(e *Engine) {
		e.retentionEnforcer = newRetentionEnforcer(e, finder)

		if e.engineID != nil {
			e.retentionEnforcer.defaultMetricLabels["engine_id"] = fmt.Sprint(*e.engineID)
		}

		if e.nodeID != nil {
			e.retentionEnforcer.defaultMetricLabels["node_id"] = fmt.Sprint(*e.nodeID)
		}

		// As new labels may have been set, set the new metrics on the enforcer.
		e.retentionEnforcer.retentionMetrics = newRetentionMetrics(e.retentionEnforcer.defaultMetricLabels)
	}
}

// NewEngine initialises a new storage engine, including a series file, index and
// TSM engine.
func NewEngine(path string, c Config, options ...Option) *Engine {
	e := &Engine{
		config: c,
		path:   path,
		sfile:  tsdb.NewSeriesFile(filepath.Join(path, tsdb.SeriesFileDirectory)),
		logger: zap.NewNop(),
	}

	// Initialise index.
	index := tsi1.NewIndex(e.sfile, "remove me", c.Index,
		tsi1.WithPath(filepath.Join(path, tsi1.DefaultIndexDirectoryName)),
	)
	e.index = index

	// Initialise Engine
	// TODO(edd): should just be able to use the config values for data/wal.
	engine := tsm1.NewEngine(0, tsdb.Index(e.index), filepath.Join(path, "data"), filepath.Join(path, "wal"), e.sfile, c.EngineOptions)

	// TODO(edd): Once the tsdb.Engine abstraction is gone, this won't be needed.
	e.engine = engine.(*tsm1.Engine)

	// Apply options.
	for _, option := range options {
		option(e)
	}
	return e
}

// WithLogger sets the logger on the Store. It must be called before Open.
func (e *Engine) WithLogger(log *zap.Logger) {
	fields := []zap.Field{}
	if e.nodeID != nil {
		fields = append(fields, zap.Int("node_id", *e.nodeID))
	}

	if e.engineID != nil {
		fields = append(fields, zap.Int("engine_id", *e.nodeID))
	}
	fields = append(fields, zap.String("service", "storage-engine"))

	e.logger = log.With(fields...)
	e.sfile.WithLogger(e.logger)
	e.index.WithLogger(e.logger)
	e.engine.WithLogger(e.logger)
	e.retentionEnforcer.WithLogger(e.logger)
}

// PrometheusCollectors returns all the prometheus collectors associated with
// the engine and its components.
func (e *Engine) PrometheusCollectors() []prometheus.Collector {
	var metrics []prometheus.Collector
	// TODO(edd): Get prom metrics for TSM.
	// TODO(edd): Get prom metrics for index.
	// TODO(edd): Get prom metrics for series file.
	metrics = append(metrics, e.retentionEnforcer.PrometheusCollectors()...)
	return metrics
}

// Open opens the store and all underlying resources. It returns an error if
// any of the underlying systems fail to open.
func (e *Engine) Open() error {
	e.mu.Lock()
	defer e.mu.Unlock()

	if e.closing != nil {
		return nil // Already open
	}

	if err := e.sfile.Open(); err != nil {
		return err
	}

	if err := e.index.Open(); err != nil {
		return err
	}

	if err := e.engine.Open(); err != nil {
		return err
	}
	e.engine.SetCompactionsEnabled(true) // TODO(edd):is this needed?

	e.closing = make(chan struct{})
	// TODO(edd) background tasks will be run in priority order via a scheduler.
	// For now we will just run on an interval as we only have the retention
	// policy enforcer.
	e.runRetentionEnforcer()

	return nil
}

// runRetentionEnforcer runs the retention enforcer in a separate goroutine.
//
// Currently this just runs on an interval, but in the future we will add the
// ability to reschedule the retention enforcement if there are not enough
// resources available.
func (e *Engine) runRetentionEnforcer() {
	if e.config.RetentionInterval == 0 {
		e.logger.Info("Retention enforcer disabled")
		return // Enforcer disabled.
	}

	if e.config.RetentionInterval < 0 {
		e.logger.Error("Negative retention interval", zap.Int64("interval", e.config.RetentionInterval))
		return
	}

	interval := time.Duration(e.config.RetentionInterval) * time.Second
	logger := e.logger.With(zap.String("component", "retention_enforcer"), zap.Duration("check_interval", interval))
	logger.Info("Starting")

	ticker := time.NewTicker(interval)
	e.wg.Add(1)
	go func() {
		defer e.wg.Done()
		for {
			// It's safe to read closing without a lock because it's never
			// modified if this goroutine is active.
			select {
			case <-e.closing:
				logger.Info("Stopping")
				return
			case <-ticker.C:
				e.retentionEnforcer.run()
			}
		}
	}()
}

// Close closes the store and all underlying resources. It returns an error if
// any of the underlying systems fail to close.
func (e *Engine) Close() error {
	e.mu.RLock()
	if e.closing == nil {
		return nil // Already closed
	}

	close(e.closing)
	e.mu.RUnlock()

	// Wait for any other goroutines to finish.
	e.wg.Wait()

	e.mu.Lock()
	defer e.mu.Unlock()
	e.closing = nil

	if err := e.sfile.Close(); err != nil {
		return err
	}

	if err := e.index.Close(); err != nil {
		return err
	}

	return e.engine.Close()
}

func (e *Engine) CreateSeriesCursor(ctx context.Context, req SeriesCursorRequest, cond influxql.Expr) (SeriesCursor, error) {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return nil, ErrEngineClosed
	}
	// TODO(edd): remove IndexSet
	return newSeriesCursor(req, tsdb.IndexSet{Indexes: []tsdb.Index{e.index}, SeriesFile: e.sfile}, cond)
}

func (e *Engine) CreateCursorIterator(ctx context.Context) (tsdb.CursorIterator, error) {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return nil, ErrEngineClosed
	}
	return e.engine.CreateCursorIterator(ctx)
}

// WritePoints writes the provided points to the engine.
//
// The Engine expects all points to have been correctly validated by the caller.
// WritePoints will however determine if there are any field type conflicts, and
// return an appropriate error in that case.
func (e *Engine) WritePoints(points []models.Point) error {
	collection := tsdb.NewSeriesCollection(points)

	j := 0
	for iter := collection.Iterator(); iter.Next(); {
		tags := iter.Tags()

		if tags.Len() > 0 && bytes.Equal(tags[0].Key, tsdb.FieldKeyTagKeyBytes) && bytes.Equal(tags[0].Value, timeBytes) {
			// Field key "time" is invalid
			if collection.Reason == "" {
				collection.Reason = fmt.Sprintf("invalid field key: input field %q is invalid", timeBytes)
			}
			collection.Dropped++
			collection.DroppedKeys = append(collection.DroppedKeys, iter.Key())
		}

		// Filter out any tags with key equal to "time": they are invalid.
		if tags.Get(timeBytes) != nil {
			if collection.Reason == "" {
				collection.Reason = fmt.Sprintf("invalid tag key: input tag %q on measurement %q is invalid", timeBytes, iter.Name())
			}
			collection.Dropped++
			collection.DroppedKeys = append(collection.DroppedKeys, iter.Key())
			continue
		}

		// Drop any series with invalid unicode characters in the key.
		if e.config.ValidateKeys && !models.ValidKeyTokens(string(iter.Name()), tags) {
			if collection.Reason == "" {
				collection.Reason = fmt.Sprintf("key contains invalid unicode: %q", iter.Key())
			}
			collection.Dropped++
			collection.DroppedKeys = append(collection.DroppedKeys, iter.Key())
			continue
		}

		collection.Copy(j, iter.Index())
		j++
	}
	collection.Truncate(j)

	e.mu.RLock()
	defer e.mu.RUnlock()

	if e.closing == nil {
		return ErrEngineClosed
	}

	// Add new series to the index and series file. Check for partial writes.
	if err := e.index.CreateSeriesListIfNotExists(collection); err != nil {
		// ignore PartialWriteErrors. The collection captures it.
		// TODO(edd/jeff): should we just remove PartialWriteError from the index then?
		if _, ok := err.(tsdb.PartialWriteError); !ok {
			return err
		}
	}

	// Write the points to the cache and WAL.
	if err := e.engine.WritePoints(collection.Points); err != nil {
		return err
	}
	return collection.PartialWriteError()
}

// DeleteSeriesRangeWithPredicate deletes all series data iterated over if fn returns
// true for that series.
func (e *Engine) DeleteSeriesRangeWithPredicate(itr tsdb.SeriesIterator, fn func([]byte, models.Tags) (int64, int64, bool)) error {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return ErrEngineClosed
	}
	return e.engine.DeleteSeriesRangeWithPredicate(itr, fn)
}

// SeriesCardinality returns the number of series in the engine.
func (e *Engine) SeriesCardinality() int64 {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return 0
	}
	return e.index.SeriesN()
}

// Path returns the path of the engine's base directory.
func (e *Engine) Path() string {
	return e.path
}

// ApplyFnToSeriesIDSet allows the caller to apply fn to the SeriesIDSet held
// within the engine's index.
func (e *Engine) ApplyFnToSeriesIDSet(fn func(*tsdb.SeriesIDSet)) {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return
	}
	fn(e.index.SeriesIDSet())
}
<commit_msg>storage: Add MeasurementCardinalityStats and MeasurementStats to Engine<commit_after>package storage

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"sync"
	"time"

	"github.com/influxdata/influxql"
	"github.com/influxdata/platform/models"
	"github.com/influxdata/platform/tsdb"
	"github.com/influxdata/platform/tsdb/tsi1"
	"github.com/influxdata/platform/tsdb/tsm1"
	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/zap"
)

// Static objects to prevent small allocs.
var timeBytes = []byte("time")

// ErrEngineClosed is returned when a caller attempts to use the engine while
// it's closed.
var ErrEngineClosed = errors.New("engine is closed")

type Engine struct {
	config   Config
	path     string
	engineID *int // Not used by default.
	nodeID   *int // Not used by default.

	mu                sync.RWMutex
	closing           chan struct{} //closing returns the zero value when the engine is shutting down.
	index             *tsi1.Index
	sfile             *tsdb.SeriesFile
	engine            *tsm1.Engine
	retentionEnforcer *retentionEnforcer

	// Tracks all goroutines started by the Engine.
	wg sync.WaitGroup

	logger *zap.Logger
}

// Option provides a set
type Option func(*Engine)

// WithTSMFilenameFormatter sets a function on the underlying tsm1.Engine to specify
// how TSM files are named.
var WithTSMFilenameFormatter = func(fn tsm1.FormatFileNameFunc) Option {
	return func(e *Engine) {
		e.engine.WithFormatFileNameFunc(fn)
	}
}

// WithEngineID sets an engine id, which can be useful for logging when multiple
// engines are in use.
var WithEngineID = func(id int) Option {
	return func(e *Engine) {
		e.engineID = &id
	}
}

// WithNodeID sets a node id on the engine, which can be useful for logging
// when a system has engines running on multiple nodes.
var WithNodeID = func(id int) Option {
	return func(e *Engine) {
		e.nodeID = &id
	}
}

// WithRetentionEnforcer initialises a retention enforcer on the engine.
// WithRetentionEnforcer must be called after other options to ensure that all
// metrics are labelled correctly.
var WithRetentionEnforcer = func(finder BucketFinder) Option {
	return func(e *Engine) {
		e.retentionEnforcer = newRetentionEnforcer(e, finder)

		if e.engineID != nil {
			e.retentionEnforcer.defaultMetricLabels["engine_id"] = fmt.Sprint(*e.engineID)
		}

		if e.nodeID != nil {
			e.retentionEnforcer.defaultMetricLabels["node_id"] = fmt.Sprint(*e.nodeID)
		}

		// As new labels may have been set, set the new metrics on the enforcer.
		e.retentionEnforcer.retentionMetrics = newRetentionMetrics(e.retentionEnforcer.defaultMetricLabels)
	}
}

// NewEngine initialises a new storage engine, including a series file, index and
// TSM engine.
func NewEngine(path string, c Config, options ...Option) *Engine {
	e := &Engine{
		config: c,
		path:   path,
		sfile:  tsdb.NewSeriesFile(filepath.Join(path, tsdb.SeriesFileDirectory)),
		logger: zap.NewNop(),
	}

	// Initialise index.
	index := tsi1.NewIndex(e.sfile, "remove me", c.Index,
		tsi1.WithPath(filepath.Join(path, tsi1.DefaultIndexDirectoryName)),
	)
	e.index = index

	// Initialise Engine
	// TODO(edd): should just be able to use the config values for data/wal.
	engine := tsm1.NewEngine(0, tsdb.Index(e.index), filepath.Join(path, "data"), filepath.Join(path, "wal"), e.sfile, c.EngineOptions)

	// TODO(edd): Once the tsdb.Engine abstraction is gone, this won't be needed.
	e.engine = engine.(*tsm1.Engine)

	// Apply options.
	for _, option := range options {
		option(e)
	}
	return e
}

// WithLogger sets the logger on the Store. It must be called before Open.
func (e *Engine) WithLogger(log *zap.Logger) {
	fields := []zap.Field{}
	if e.nodeID != nil {
		fields = append(fields, zap.Int("node_id", *e.nodeID))
	}

	if e.engineID != nil {
		fields = append(fields, zap.Int("engine_id", *e.nodeID))
	}
	fields = append(fields, zap.String("service", "storage-engine"))

	e.logger = log.With(fields...)
	e.sfile.WithLogger(e.logger)
	e.index.WithLogger(e.logger)
	e.engine.WithLogger(e.logger)
	e.retentionEnforcer.WithLogger(e.logger)
}

// PrometheusCollectors returns all the prometheus collectors associated with
// the engine and its components.
func (e *Engine) PrometheusCollectors() []prometheus.Collector {
	var metrics []prometheus.Collector
	// TODO(edd): Get prom metrics for TSM.
	// TODO(edd): Get prom metrics for index.
	// TODO(edd): Get prom metrics for series file.
	metrics = append(metrics, e.retentionEnforcer.PrometheusCollectors()...)
	return metrics
}

// Open opens the store and all underlying resources. It returns an error if
// any of the underlying systems fail to open.
func (e *Engine) Open() error {
	e.mu.Lock()
	defer e.mu.Unlock()

	if e.closing != nil {
		return nil // Already open
	}

	if err := e.sfile.Open(); err != nil {
		return err
	}

	if err := e.index.Open(); err != nil {
		return err
	}

	if err := e.engine.Open(); err != nil {
		return err
	}
	e.engine.SetCompactionsEnabled(true) // TODO(edd):is this needed?

	e.closing = make(chan struct{})
	// TODO(edd) background tasks will be run in priority order via a scheduler.
	// For now we will just run on an interval as we only have the retention
	// policy enforcer.
	e.runRetentionEnforcer()

	return nil
}

// runRetentionEnforcer runs the retention enforcer in a separate goroutine.
//
// Currently this just runs on an interval, but in the future we will add the
// ability to reschedule the retention enforcement if there are not enough
// resources available.
func (e *Engine) runRetentionEnforcer() {
	if e.config.RetentionInterval == 0 {
		e.logger.Info("Retention enforcer disabled")
		return // Enforcer disabled.
	}

	if e.config.RetentionInterval < 0 {
		e.logger.Error("Negative retention interval", zap.Int64("interval", e.config.RetentionInterval))
		return
	}

	interval := time.Duration(e.config.RetentionInterval) * time.Second
	logger := e.logger.With(zap.String("component", "retention_enforcer"), zap.Duration("check_interval", interval))
	logger.Info("Starting")

	ticker := time.NewTicker(interval)
	e.wg.Add(1)
	go func() {
		defer e.wg.Done()
		for {
			// It's safe to read closing without a lock because it's never
			// modified if this goroutine is active.
			select {
			case <-e.closing:
				logger.Info("Stopping")
				return
			case <-ticker.C:
				e.retentionEnforcer.run()
			}
		}
	}()
}

// Close closes the store and all underlying resources. It returns an error if
// any of the underlying systems fail to close.
func (e *Engine) Close() error {
	e.mu.RLock()
	if e.closing == nil {
		return nil // Already closed
	}

	close(e.closing)
	e.mu.RUnlock()

	// Wait for any other goroutines to finish.
	e.wg.Wait()

	e.mu.Lock()
	defer e.mu.Unlock()
	e.closing = nil

	if err := e.sfile.Close(); err != nil {
		return err
	}

	if err := e.index.Close(); err != nil {
		return err
	}

	return e.engine.Close()
}

func (e *Engine) CreateSeriesCursor(ctx context.Context, req SeriesCursorRequest, cond influxql.Expr) (SeriesCursor, error) {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return nil, ErrEngineClosed
	}
	// TODO(edd): remove IndexSet
	return newSeriesCursor(req, tsdb.IndexSet{Indexes: []tsdb.Index{e.index}, SeriesFile: e.sfile}, cond)
}

func (e *Engine) CreateCursorIterator(ctx context.Context) (tsdb.CursorIterator, error) {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return nil, ErrEngineClosed
	}
	return e.engine.CreateCursorIterator(ctx)
}

// WritePoints writes the provided points to the engine.
//
// The Engine expects all points to have been correctly validated by the caller.
// WritePoints will however determine if there are any field type conflicts, and
// return an appropriate error in that case.
func (e *Engine) WritePoints(points []models.Point) error {
	collection := tsdb.NewSeriesCollection(points)

	j := 0
	for iter := collection.Iterator(); iter.Next(); {
		tags := iter.Tags()

		if tags.Len() > 0 && bytes.Equal(tags[0].Key, tsdb.FieldKeyTagKeyBytes) && bytes.Equal(tags[0].Value, timeBytes) {
			// Field key "time" is invalid
			if collection.Reason == "" {
				collection.Reason = fmt.Sprintf("invalid field key: input field %q is invalid", timeBytes)
			}
			collection.Dropped++
			collection.DroppedKeys = append(collection.DroppedKeys, iter.Key())
		}

		// Filter out any tags with key equal to "time": they are invalid.
		if tags.Get(timeBytes) != nil {
			if collection.Reason == "" {
				collection.Reason = fmt.Sprintf("invalid tag key: input tag %q on measurement %q is invalid", timeBytes, iter.Name())
			}
			collection.Dropped++
			collection.DroppedKeys = append(collection.DroppedKeys, iter.Key())
			continue
		}

		// Drop any series with invalid unicode characters in the key.
		if e.config.ValidateKeys && !models.ValidKeyTokens(string(iter.Name()), tags) {
			if collection.Reason == "" {
				collection.Reason = fmt.Sprintf("key contains invalid unicode: %q", iter.Key())
			}
			collection.Dropped++
			collection.DroppedKeys = append(collection.DroppedKeys, iter.Key())
			continue
		}

		collection.Copy(j, iter.Index())
		j++
	}
	collection.Truncate(j)

	e.mu.RLock()
	defer e.mu.RUnlock()

	if e.closing == nil {
		return ErrEngineClosed
	}

	// Add new series to the index and series file. Check for partial writes.
	if err := e.index.CreateSeriesListIfNotExists(collection); err != nil {
		// ignore PartialWriteErrors. The collection captures it.
		// TODO(edd/jeff): should we just remove PartialWriteError from the index then?
		if _, ok := err.(tsdb.PartialWriteError); !ok {
			return err
		}
	}

	// Write the points to the cache and WAL.
	if err := e.engine.WritePoints(collection.Points); err != nil {
		return err
	}
	return collection.PartialWriteError()
}

// DeleteSeriesRangeWithPredicate deletes all series data iterated over if fn returns
// true for that series.
func (e *Engine) DeleteSeriesRangeWithPredicate(itr tsdb.SeriesIterator, fn func([]byte, models.Tags) (int64, int64, bool)) error {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return ErrEngineClosed
	}
	return e.engine.DeleteSeriesRangeWithPredicate(itr, fn)
}

// SeriesCardinality returns the number of series in the engine.
func (e *Engine) SeriesCardinality() int64 {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return 0
	}
	return e.index.SeriesN()
}

// Path returns the path of the engine's base directory.
func (e *Engine) Path() string {
	return e.path
}

// ApplyFnToSeriesIDSet allows the caller to apply fn to the SeriesIDSet held
// within the engine's index.
func (e *Engine) ApplyFnToSeriesIDSet(fn func(*tsdb.SeriesIDSet)) {
	e.mu.RLock()
	defer e.mu.RUnlock()
	if e.closing == nil {
		return
	}
	fn(e.index.SeriesIDSet())
}

// MeasurementCardinalityStats returns cardinality stats for all measurements.
func (e *Engine) MeasurementCardinalityStats() tsi1.MeasurementCardinalityStats {
	return e.index.MeasurementCardinalityStats()
}

// MeasurementStats returns the current measurement stats for the engine.
func (e *Engine) MeasurementStats() (tsm1.MeasurementStats, error) {
	return e.engine.MeasurementStats()
}
<|endoftext|>


<commit_before>package stream

import (
	"net"
	"os"

	"github.com/bahadley/esp/log"
	"github.com/bahadley/esp/operator"
)

const (
	defaultIngestAddr = "localhost"
	defaultIngestPort = "22221"

	envIngestAddr = "ESP_ADDR"
	envIngestPort = "ESP_PORT"
)

var (
	IngestAddr *net.UDPAddr
)

func Ingress() {
	conn, err := net.ListenUDP("udp", IngestAddr)
	if err != nil {
		log.Error.Fatal(err.Error())
	}
	defer conn.Close()

	log.Info.Printf("Listening for sensor tuples (%s UDP) ...",
		IngestAddr.String())

	go Egress()
	go operator.Ingest()

	buf := make([]byte, 128, 1024)
	for {
		n, caddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			log.Warning.Println(err.Error())
			continue
		}

		msg := string(buf[0:n])
		log.Info.Printf("Rx(%s): %s", caddr, msg)
		operator.IngestChan <- msg
	}
}

func init() {
	// Build the UDP address that we will listen on.
	addr := os.Getenv(envIngestAddr)
	if len(addr) == 0 {
		addr = defaultIngestAddr
	}

	port := os.Getenv(envIngestPort)
	if len(port) == 0 {
		port = defaultIngestPort
	}

	var err error
	IngestAddr, err = net.ResolveUDPAddr("udp", addr+":"+port)
	if err != nil {
		log.Error.Fatal(err.Error())
	}
}
<commit_msg>Changed references to ingest to ingress<commit_after>package stream

import (
	"net"
	"os"

	"github.com/bahadley/esp/log"
	"github.com/bahadley/esp/operator"
)

const (
	defaultIngressAddr = "localhost"
	defaultIngressPort = "22221"

	envIngressAddr = "ESP_ADDR"
	envIngressPort = "ESP_PORT"
)

var (
	IngressAddr *net.UDPAddr
)

func Ingress() {
	conn, err := net.ListenUDP("udp", IngressAddr)
	if err != nil {
		log.Error.Fatal(err.Error())
	}
	defer conn.Close()

	log.Info.Printf("Listening for sensor tuples (%s UDP) ...",
		IngressAddr.String())

	go Egress()
	go operator.Ingest()

	buf := make([]byte, 128, 1024)
	for {
		n, caddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			log.Warning.Println(err.Error())
			continue
		}

		msg := string(buf[0:n])
		log.Info.Printf("Rx(%s): %s", caddr, msg)
		operator.IngestChan <- msg
	}
}

func init() {
	// Build the UDP address that we will listen on.
	addr := os.Getenv(envIngressAddr)
	if len(addr) == 0 {
		addr = defaultIngressAddr
	}

	port := os.Getenv(envIngressPort)
	if len(port) == 0 {
		port = defaultIngressPort
	}

	var err error
	IngressAddr, err = net.ResolveUDPAddr("udp", addr+":"+port)
	if err != nil {
		log.Error.Fatal(err.Error())
	}
}
<|endoftext|>


<commit_before>// This package implements a provisioner for Packer that executes
// Puppet on the remote machine, configured to apply a local manifest
// versus connecting to a Puppet master.
package puppetmasterless

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

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	ctx                 interpolate.Context

	// The command used to execute Puppet.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Additional arguments to pass when executing Puppet
	ExtraArguments []string `mapstructure:"extra_arguments"`

	// Additional facts to set when executing Puppet
	Facter map[string]string

	// Path to a hiera configuration file to upload and use.
	HieraConfigPath string `mapstructure:"hiera_config_path"`

	// An array of local paths of modules to upload.
	ModulePaths []string `mapstructure:"module_paths"`

	// The main manifest file to apply to kick off the entire thing.
	ManifestFile string `mapstructure:"manifest_file"`

	// A directory of manifest files that will be uploaded to the remote
	// machine.
	ManifestDir string `mapstructure:"manifest_dir"`

	// If true, `sudo` will NOT be used to execute Puppet.
	PreventSudo bool `mapstructure:"prevent_sudo"`

	// The directory where files will be uploaded. Packer requires write
	// permissions in this directory.
	StagingDir string `mapstructure:"staging_directory"`

	// The directory from which the command will be executed.
	// Packer requires the directory to exist when running puppet.
	WorkingDir string `mapstructure:"working_directory"`
}

type Provisioner struct {
	config Config
}

type ExecuteTemplate struct {
	WorkingDir      string
	FacterVars      string
	HieraConfigPath string
	ModulePath      string
	ManifestFile    string
	ManifestDir     string
	Sudo            bool
	ExtraArguments  string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	// Set some defaults
	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "cd {{.WorkingDir}} && " +
			"{{.FacterVars}} {{if .Sudo}} sudo -E {{end}}" +
			"puppet apply --verbose --modulepath='{{.ModulePath}}' " +
			"{{if ne .HieraConfigPath \"\"}}--hiera_config='{{.HieraConfigPath}}' {{end}}" +
			"{{if ne .ManifestDir \"\"}}--manifestdir='{{.ManifestDir}}' {{end}}" +
			"--detailed-exitcodes " +
			"{{.ExtraArguments}} " +
			"{{.ManifestFile}}"
	}

	if p.config.StagingDir == "" {
		p.config.StagingDir = "/tmp/packer-puppet-masterless"
	}

	if p.config.WorkingDir == "" {
		p.config.WorkingDir = p.config.StagingDir
	}

	if p.config.Facter == nil {
		p.config.Facter = make(map[string]string)
	}
	p.config.Facter["packer_build_name"] = p.config.PackerBuildName
	p.config.Facter["packer_builder_type"] = p.config.PackerBuilderType

	// Validation
	var errs *packer.MultiError
	if p.config.HieraConfigPath != "" {
		info, err := os.Stat(p.config.HieraConfigPath)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path is invalid: %s", err))
		} else if info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path must point to a file"))
		}
	}

	if p.config.ManifestDir != "" {
		info, err := os.Stat(p.config.ManifestDir)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_dir is invalid: %s", err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_dir must point to a directory"))
		}
	}

	if p.config.ManifestFile == "" {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("A manifest_file must be specified."))
	} else {
		_, err := os.Stat(p.config.ManifestFile)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_file is invalid: %s", err))
		}
	}

	for i, path := range p.config.ModulePaths {
		info, err := os.Stat(path)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] is invalid: %s", i, err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] must point to a directory", i))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Provisioning with Puppet...")
	ui.Message("Creating Puppet staging directory...")
	if err := p.createDir(ui, comm, p.config.StagingDir); err != nil {
		return fmt.Errorf("Error creating staging directory: %s", err)
	}

	// Upload hiera config if set
	remoteHieraConfigPath := ""
	if p.config.HieraConfigPath != "" {
		var err error
		remoteHieraConfigPath, err = p.uploadHieraConfig(ui, comm)
		if err != nil {
			return fmt.Errorf("Error uploading hiera config: %s", err)
		}
	}

	// Upload manifest dir if set
	remoteManifestDir := ""
	if p.config.ManifestDir != "" {
		ui.Message(fmt.Sprintf(
			"Uploading manifest directory from: %s", p.config.ManifestDir))
		remoteManifestDir = fmt.Sprintf("%s/manifests", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteManifestDir, p.config.ManifestDir)
		if err != nil {
			return fmt.Errorf("Error uploading manifest dir: %s", err)
		}
	}

	// Upload all modules
	modulePaths := make([]string, 0, len(p.config.ModulePaths))
	for i, path := range p.config.ModulePaths {
		ui.Message(fmt.Sprintf("Uploading local modules from: %s", path))
		targetPath := fmt.Sprintf("%s/module-%d", p.config.StagingDir, i)
		if err := p.uploadDirectory(ui, comm, targetPath, path); err != nil {
			return fmt.Errorf("Error uploading modules: %s", err)
		}

		modulePaths = append(modulePaths, targetPath)
	}

	// Upload manifests
	remoteManifestFile, err := p.uploadManifests(ui, comm)
	if err != nil {
		return fmt.Errorf("Error uploading manifests: %s", err)
	}

	// Compile the facter variables
	facterVars := make([]string, 0, len(p.config.Facter))
	for k, v := range p.config.Facter {
		facterVars = append(facterVars, fmt.Sprintf("FACTER_%s='%s'", k, v))
	}

	// Execute Puppet
	p.config.ctx.Data = &ExecuteTemplate{
		FacterVars:      strings.Join(facterVars, " "),
		HieraConfigPath: remoteHieraConfigPath,
		ManifestDir:     remoteManifestDir,
		ManifestFile:    remoteManifestFile,
		ModulePath:      strings.Join(modulePaths, ":"),
		Sudo:            !p.config.PreventSudo,
		WorkingDir:      p.config.WorkingDir,
		ExtraArguments:  strings.Join(p.config.ExtraArguments, " "),
	}
	command, err := interpolate.Render(p.config.ExecuteCommand, &p.config.ctx)
	if err != nil {
		return err
	}

	cmd := &packer.RemoteCmd{
		Command: command,
	}

	ui.Message(fmt.Sprintf("Running Puppet: %s", command))
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 && cmd.ExitStatus != 2 {
		return fmt.Errorf("Puppet exited with a non-zero exit status: %d", cmd.ExitStatus)
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

func (p *Provisioner) uploadHieraConfig(ui packer.Ui, comm packer.Communicator) (string, error) {
	ui.Message("Uploading hiera configuration...")
	f, err := os.Open(p.config.HieraConfigPath)
	if err != nil {
		return "", err
	}
	defer f.Close()

	path := fmt.Sprintf("%s/hiera.yaml", p.config.StagingDir)
	if err := comm.Upload(path, f, nil); err != nil {
		return "", err
	}

	return path, nil
}

func (p *Provisioner) uploadManifests(ui packer.Ui, comm packer.Communicator) (string, error) {
	// Create the remote manifests directory...
	ui.Message("Uploading manifests...")
	remoteManifestsPath := fmt.Sprintf("%s/manifests", p.config.StagingDir)
	if err := p.createDir(ui, comm, remoteManifestsPath); err != nil {
		return "", fmt.Errorf("Error creating manifests directory: %s", err)
	}

	// NOTE! manifest_file may either be a directory or a file, as puppet apply
	// now accepts either one.

	fi, err := os.Stat(p.config.ManifestFile)
	if err != nil {
		return "", fmt.Errorf("Error inspecting manifest file: %s", err)
	}

	if fi.IsDir() {
		// If manifest_file is a directory we'll upload the whole thing
		ui.Message(fmt.Sprintf(
			"Uploading manifest directory from: %s", p.config.ManifestFile))

		remoteManifestDir := fmt.Sprintf("%s/manifests", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteManifestDir, p.config.ManifestFile)
		if err != nil {
			return "", fmt.Errorf("Error uploading manifest dir: %s", err)
		}
		return remoteManifestDir, nil
	} else {
		// Otherwise manifest_file is a file and we'll upload it
		ui.Message(fmt.Sprintf(
			"Uploading manifest file from: %s", p.config.ManifestFile))

		f, err := os.Open(p.config.ManifestFile)
		if err != nil {
			return "", err
		}
		defer f.Close()

		manifestFilename := filepath.Base(p.config.ManifestFile)
		remoteManifestFile := fmt.Sprintf("%s/%s", remoteManifestsPath, manifestFilename)
		if err := comm.Upload(remoteManifestFile, f, nil); err != nil {
			return "", err
		}
		return remoteManifestFile, nil
	}
}

func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("mkdir -p '%s'", dir),
	}

	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status.")
	}

	return nil
}

func (p *Provisioner) uploadDirectory(ui packer.Ui, comm packer.Communicator, dst string, src string) error {
	if err := p.createDir(ui, comm, dst); err != nil {
		return err
	}

	// Make sure there is a trailing "/" so that the directory isn't
	// created on the other side.
	if src[len(src)-1] != '/' {
		src = src + "/"
	}

	return comm.UploadDir(dst, src, nil)
}
<commit_msg>Fixing the bug found in the tests<commit_after>// This package implements a provisioner for Packer that executes
// Puppet on the remote machine, configured to apply a local manifest
// versus connecting to a Puppet master.
package puppetmasterless

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

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	ctx                 interpolate.Context

	// The command used to execute Puppet.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Additional arguments to pass when executing Puppet
	ExtraArguments []string `mapstructure:"extra_arguments"`

	// Additional facts to set when executing Puppet
	Facter map[string]string

	// Path to a hiera configuration file to upload and use.
	HieraConfigPath string `mapstructure:"hiera_config_path"`

	// An array of local paths of modules to upload.
	ModulePaths []string `mapstructure:"module_paths"`

	// The main manifest file to apply to kick off the entire thing.
	ManifestFile string `mapstructure:"manifest_file"`

	// A directory of manifest files that will be uploaded to the remote
	// machine.
	ManifestDir string `mapstructure:"manifest_dir"`

	// If true, `sudo` will NOT be used to execute Puppet.
	PreventSudo bool `mapstructure:"prevent_sudo"`

	// The directory where files will be uploaded. Packer requires write
	// permissions in this directory.
	StagingDir string `mapstructure:"staging_directory"`

	// The directory from which the command will be executed.
	// Packer requires the directory to exist when running puppet.
	WorkingDir string `mapstructure:"working_directory"`
}

type Provisioner struct {
	config Config
}

type ExecuteTemplate struct {
	WorkingDir      string
	FacterVars      string
	HieraConfigPath string
	ModulePath      string
	ManifestFile    string
	ManifestDir     string
	Sudo            bool
	ExtraArguments  string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	// Set some defaults
	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "cd {{.WorkingDir}} && " +
			"{{.FacterVars}} {{if .Sudo}} sudo -E {{end}}" +
			"puppet apply --verbose --modulepath='{{.ModulePath}}' " +
			"{{if ne .HieraConfigPath \"\"}}--hiera_config='{{.HieraConfigPath}}' {{end}}" +
			"{{if ne .ManifestDir \"\"}}--manifestdir='{{.ManifestDir}}' {{end}}" +
			"--detailed-exitcodes " +
			"{{if ne .ExtraArguments \"\"}}{{.ExtraArguments}} {{end}}" +
			"{{.ManifestFile}}"
	}

	if p.config.StagingDir == "" {
		p.config.StagingDir = "/tmp/packer-puppet-masterless"
	}

	if p.config.WorkingDir == "" {
		p.config.WorkingDir = p.config.StagingDir
	}

	if p.config.Facter == nil {
		p.config.Facter = make(map[string]string)
	}
	p.config.Facter["packer_build_name"] = p.config.PackerBuildName
	p.config.Facter["packer_builder_type"] = p.config.PackerBuilderType

	// Validation
	var errs *packer.MultiError
	if p.config.HieraConfigPath != "" {
		info, err := os.Stat(p.config.HieraConfigPath)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path is invalid: %s", err))
		} else if info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path must point to a file"))
		}
	}

	if p.config.ManifestDir != "" {
		info, err := os.Stat(p.config.ManifestDir)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_dir is invalid: %s", err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_dir must point to a directory"))
		}
	}

	if p.config.ManifestFile == "" {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("A manifest_file must be specified."))
	} else {
		_, err := os.Stat(p.config.ManifestFile)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_file is invalid: %s", err))
		}
	}

	for i, path := range p.config.ModulePaths {
		info, err := os.Stat(path)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] is invalid: %s", i, err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] must point to a directory", i))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Provisioning with Puppet...")
	ui.Message("Creating Puppet staging directory...")
	if err := p.createDir(ui, comm, p.config.StagingDir); err != nil {
		return fmt.Errorf("Error creating staging directory: %s", err)
	}

	// Upload hiera config if set
	remoteHieraConfigPath := ""
	if p.config.HieraConfigPath != "" {
		var err error
		remoteHieraConfigPath, err = p.uploadHieraConfig(ui, comm)
		if err != nil {
			return fmt.Errorf("Error uploading hiera config: %s", err)
		}
	}

	// Upload manifest dir if set
	remoteManifestDir := ""
	if p.config.ManifestDir != "" {
		ui.Message(fmt.Sprintf(
			"Uploading manifest directory from: %s", p.config.ManifestDir))
		remoteManifestDir = fmt.Sprintf("%s/manifests", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteManifestDir, p.config.ManifestDir)
		if err != nil {
			return fmt.Errorf("Error uploading manifest dir: %s", err)
		}
	}

	// Upload all modules
	modulePaths := make([]string, 0, len(p.config.ModulePaths))
	for i, path := range p.config.ModulePaths {
		ui.Message(fmt.Sprintf("Uploading local modules from: %s", path))
		targetPath := fmt.Sprintf("%s/module-%d", p.config.StagingDir, i)
		if err := p.uploadDirectory(ui, comm, targetPath, path); err != nil {
			return fmt.Errorf("Error uploading modules: %s", err)
		}

		modulePaths = append(modulePaths, targetPath)
	}

	// Upload manifests
	remoteManifestFile, err := p.uploadManifests(ui, comm)
	if err != nil {
		return fmt.Errorf("Error uploading manifests: %s", err)
	}

	// Compile the facter variables
	facterVars := make([]string, 0, len(p.config.Facter))
	for k, v := range p.config.Facter {
		facterVars = append(facterVars, fmt.Sprintf("FACTER_%s='%s'", k, v))
	}

	// Execute Puppet
	p.config.ctx.Data = &ExecuteTemplate{
		FacterVars:      strings.Join(facterVars, " "),
		HieraConfigPath: remoteHieraConfigPath,
		ManifestDir:     remoteManifestDir,
		ManifestFile:    remoteManifestFile,
		ModulePath:      strings.Join(modulePaths, ":"),
		Sudo:            !p.config.PreventSudo,
		WorkingDir:      p.config.WorkingDir,
		ExtraArguments:  strings.Join(p.config.ExtraArguments, " "),
	}
	command, err := interpolate.Render(p.config.ExecuteCommand, &p.config.ctx)
	if err != nil {
		return err
	}

	cmd := &packer.RemoteCmd{
		Command: command,
	}

	ui.Message(fmt.Sprintf("Running Puppet: %s", command))
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 && cmd.ExitStatus != 2 {
		return fmt.Errorf("Puppet exited with a non-zero exit status: %d", cmd.ExitStatus)
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

func (p *Provisioner) uploadHieraConfig(ui packer.Ui, comm packer.Communicator) (string, error) {
	ui.Message("Uploading hiera configuration...")
	f, err := os.Open(p.config.HieraConfigPath)
	if err != nil {
		return "", err
	}
	defer f.Close()

	path := fmt.Sprintf("%s/hiera.yaml", p.config.StagingDir)
	if err := comm.Upload(path, f, nil); err != nil {
		return "", err
	}

	return path, nil
}

func (p *Provisioner) uploadManifests(ui packer.Ui, comm packer.Communicator) (string, error) {
	// Create the remote manifests directory...
	ui.Message("Uploading manifests...")
	remoteManifestsPath := fmt.Sprintf("%s/manifests", p.config.StagingDir)
	if err := p.createDir(ui, comm, remoteManifestsPath); err != nil {
		return "", fmt.Errorf("Error creating manifests directory: %s", err)
	}

	// NOTE! manifest_file may either be a directory or a file, as puppet apply
	// now accepts either one.

	fi, err := os.Stat(p.config.ManifestFile)
	if err != nil {
		return "", fmt.Errorf("Error inspecting manifest file: %s", err)
	}

	if fi.IsDir() {
		// If manifest_file is a directory we'll upload the whole thing
		ui.Message(fmt.Sprintf(
			"Uploading manifest directory from: %s", p.config.ManifestFile))

		remoteManifestDir := fmt.Sprintf("%s/manifests", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteManifestDir, p.config.ManifestFile)
		if err != nil {
			return "", fmt.Errorf("Error uploading manifest dir: %s", err)
		}
		return remoteManifestDir, nil
	} else {
		// Otherwise manifest_file is a file and we'll upload it
		ui.Message(fmt.Sprintf(
			"Uploading manifest file from: %s", p.config.ManifestFile))

		f, err := os.Open(p.config.ManifestFile)
		if err != nil {
			return "", err
		}
		defer f.Close()

		manifestFilename := filepath.Base(p.config.ManifestFile)
		remoteManifestFile := fmt.Sprintf("%s/%s", remoteManifestsPath, manifestFilename)
		if err := comm.Upload(remoteManifestFile, f, nil); err != nil {
			return "", err
		}
		return remoteManifestFile, nil
	}
}

func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("mkdir -p '%s'", dir),
	}

	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status.")
	}

	return nil
}

func (p *Provisioner) uploadDirectory(ui packer.Ui, comm packer.Communicator, dst string, src string) error {
	if err := p.createDir(ui, comm, dst); err != nil {
		return err
	}

	// Make sure there is a trailing "/" so that the directory isn't
	// created on the other side.
	if src[len(src)-1] != '/' {
		src = src + "/"
	}

	return comm.UploadDir(dst, src, nil)
}
<|endoftext|>


<commit_before>package ui

import "testing"

type dummyRenderer struct {
}

func (dummyRenderer) Render(bb *BufferBuilder) {
	bb.WriteString("xy", "1")
}

func TestRender(t *testing.T) {
	b := Render(dummyRenderer{}, 10)
	if b.Width != 10 {
		t.Errorf("Rendered Buffer has Width %d, want %d", b.Width, 10)
	}
	if eq, _ := CompareCells(b.Lines[0], []Cell{{"x", 1, "1"}, {"y", 1, "1"}}); !eq {
		t.Errorf("Rendered Buffer has unexpected content")
	}
}
<commit_msg>edit/ui: Add test cases against renderers.<commit_after>package ui

import (
	"testing"

	"github.com/elves/elvish/tt"
)

type dummyRenderer struct{}

func (dummyRenderer) Render(bb *BufferBuilder) { bb.WriteString("dummy", "1") }

var Args = tt.Args

func TestRender(t *testing.T) {
	tt.Test(t, tt.Fn("Render", Render), tt.Table{
		Args(dummyRenderer{}, 10).
			Rets(NewBufferBuilder(10).WriteString("dummy", "1").Buffer()),

		Args(NewStringRenderer("string"), 10).
			Rets(NewBufferBuilder(10).WriteString("string", "").Buffer()),
		Args(NewStringRenderer("string"), 3).
			Rets(NewBufferBuilder(3).WriteString("str", "").Buffer()),

		Args(NewModeLineRenderer("M", "f"), 10).
			Rets(
				NewBufferBuilder(10).
					WriteString("M", styleForMode.String()).
					WriteSpaces(1, "").
					WriteString("f", styleForFilter.String()).
					SetDotToCursor().
					Buffer()),

		// Width left for scrollbar is 5
		Args(NewModeLineWithScrollBarRenderer(NewModeLineRenderer("M", "f"), 5, 0, 1), 10).
			Rets(
				NewBufferBuilder(10).
					WriteString("M", styleForMode.String()).
					WriteSpaces(1, "").
					WriteString("f", styleForFilter.String()).
					SetDotToCursor().
					WriteSpaces(1, "").
					Write(' ', styleForScrollBarThumb.String()).
					WriteString("━━━━", styleForScrollBarArea.String()).
					Buffer()),
	})
}
<|endoftext|>


<commit_before>// Copyright 2014 The lime Authors.
// Use of this source code is governed by a 2-clause
// BSD-style license that can be found in the LICENSE file.

package render
<commit_msg>Add some tests for render/view.<commit_after>// Copyright 2014 The lime Authors.
// Use of this source code is governed by a 2-clause
// BSD-style license that can be found in the LICENSE file.

package render

import (
	"github.com/quarnster/util/text"
	"reflect"
	"testing"
)

func TestViewRegionsCull(t *testing.T) {
	tests := []struct {
		regions []text.Region
		cull text.Region
		exp []text.Region
	}{
		{
			[]text.Region{{100, 200}},
			text.Region{0, 50},
			[]text.Region{},
		},
		{
			[]text.Region{{100, 100}},
			text.Region{100, 100},
			[]text.Region{{100, 100}},
		},
		{
			[]text.Region{{100, 100}},
			text.Region{95, 105},
			[]text.Region{{100, 100}},
		},
		{
			[]text.Region{{100, 100}},
			text.Region{95, 100},
			[]text.Region{{100, 100}},
		},
		{
			[]text.Region{{100, 200}},
			text.Region{150, 150},
			[]text.Region{{150, 150}},
		},
		{
			[]text.Region{{100, 200}},
			text.Region{90, 100},
			[]text.Region{{100, 100}},
		},
		{
			[]text.Region{{100, 200}},
			text.Region{100, 150},
			[]text.Region{{100, 150}},
		},
		{
			[]text.Region{{100, 200}},
			text.Region{150, 175},
			[]text.Region{{150, 175}},
		},
		{
			[]text.Region{{100, 200}},
			text.Region{0, 150},
			[]text.Region{{100, 150}},
		},
		{
			[]text.Region{{100, 200}},
			text.Region{150, 250},
			[]text.Region{{150, 200}},
		},
		{
			[]text.Region{{100, 200}},
			text.Region{0, 250},
			[]text.Region{{100, 200}},
		},
		{
			[]text.Region{{100, 200}, {300, 400}},
			text.Region{0, 500},
			[]text.Region{{100, 200}, {300, 400}},
		},
		{
			[]text.Region{{100, 200}, {300, 400}},
			text.Region{150, 350},
			[]text.Region{{150, 200}, {300, 350}},
		},
		{
			[]text.Region{{100, 200}, {300, 400}},
			text.Region{150, 250},
			[]text.Region{{150, 200}},
		},
		{
			[]text.Region{{100, 200}, {300, 400}},
			text.Region{250, 350},
			[]text.Region{{300, 350}},
		},
	}

	for i, test := range tests {
		vr := ViewRegions{}
		vr.Regions.AddAll(test.regions)

		vr.Cull(test.cull)

		r := vr.Regions.Regions()

		if !reflect.DeepEqual(r, test.exp) {
			t.Errorf("Test %d: Expected %s, but got %s", i, test.exp, r)
		}
	}
}

func TestViewRegionsClone(t *testing.T) {
	vr := ViewRegions{
		Scope:   "testScope",
		Icon:    "testIcon",
		Flags:   100,
	}
	vr.Regions.AddAll([]text.Region{{0, 0}, {120, 300}, {24, 34}, {45, 40}})

	c := vr.Clone()
	if !reflect.DeepEqual(c, vr) {
		t.Errorf("Expected %+v, but got %+v", vr, c)
	}
}
<|endoftext|>