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

// Package boltdb implements a store.KVStore on top of BoltDB. It supports the
// following options:
//
// "bucket" (string): the name of BoltDB bucket to use, defaults to "bleve".
//
// "nosync" (bool): if true, set boltdb.DB.NoSync to true. It speeds up index
// operations in exchange of losing integrity guarantees if indexation aborts
// without closing the index. Use it when rebuilding indexes from zero.
package boltdb

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"

	"github.com/blevesearch/bleve/index/store"
	"github.com/blevesearch/bleve/registry"
	bolt "github.com/etcd-io/bbolt"
)

const (
	Name                    = "boltdb"
	defaultCompactBatchSize = 100
)

type Store struct {
	path        string
	bucket      string
	db          *bolt.DB
	noSync      bool
	fillPercent float64
	mo          store.MergeOperator
}

func New(mo store.MergeOperator, config map[string]interface{}) (store.KVStore, error) {
	path, ok := config["path"].(string)
	if !ok {
		return nil, fmt.Errorf("must specify path")
	}
	if path == "" {
		return nil, os.ErrInvalid
	}

	bucket, ok := config["bucket"].(string)
	if !ok {
		bucket = "bleve"
	}

	noSync, _ := config["nosync"].(bool)

	fillPercent, ok := config["fillPercent"].(float64)
	if !ok {
		fillPercent = bolt.DefaultFillPercent
	}

	bo := &bolt.Options{}
	ro, ok := config["read_only"].(bool)
	if ok {
		bo.ReadOnly = ro
	}

	db, err := bolt.Open(path, 0600, bo)
	if err != nil {
		return nil, err
	}
	db.NoSync = noSync

	if !bo.ReadOnly {
		err = db.Update(func(tx *bolt.Tx) error {
			_, err := tx.CreateBucketIfNotExists([]byte(bucket))

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

	rv := Store{
		path:        path,
		bucket:      bucket,
		db:          db,
		mo:          mo,
		noSync:      noSync,
		fillPercent: fillPercent,
	}
	return &rv, nil
}

func (bs *Store) Close() error {
	return bs.db.Close()
}

func (bs *Store) Reader() (store.KVReader, error) {
	tx, err := bs.db.Begin(false)
	if err != nil {
		return nil, err
	}
	return &Reader{
		store:  bs,
		tx:     tx,
		bucket: tx.Bucket([]byte(bs.bucket)),
	}, nil
}

func (bs *Store) Writer() (store.KVWriter, error) {
	return &Writer{
		store: bs,
	}, nil
}

func (bs *Store) Stats() json.Marshaler {
	return &stats{
		s: bs,
	}
}

// CompactWithBatchSize removes DictionaryTerm entries with a count of zero (in batchSize batches)
// Removing entries is a workaround for github issue #374.
func (bs *Store) CompactWithBatchSize(batchSize int) error {
	for {
		cnt := 0
		err := bs.db.Batch(func(tx *bolt.Tx) error {
			c := tx.Bucket([]byte(bs.bucket)).Cursor()
			prefix := []byte("d")

			for k, v := c.Seek(prefix); bytes.HasPrefix(k, prefix); k, v = c.Next() {
				if bytes.Equal(v, []byte{0}) {
					cnt++
					if err := c.Delete(); err != nil {
						return err
					}
					if cnt == batchSize {
						break
					}
				}

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

		if cnt == 0 {
			break
		}
	}
	return nil
}

// Compact calls CompactWithBatchSize with a default batch size of 100.  This is a workaround
// for github issue #374.
func (bs *Store) Compact() error {
	return bs.CompactWithBatchSize(defaultCompactBatchSize)
}

func init() {
	registry.RegisterKVStore(Name, New)
}
<commit_msg>add support for configuring the boltdb initialMmapSize option (#1135)<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 		http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package boltdb implements a store.KVStore on top of BoltDB. It supports the
// following options:
//
// "bucket" (string): the name of BoltDB bucket to use, defaults to "bleve".
//
// "nosync" (bool): if true, set boltdb.DB.NoSync to true. It speeds up index
// operations in exchange of losing integrity guarantees if indexation aborts
// without closing the index. Use it when rebuilding indexes from zero.
package boltdb

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"

	"github.com/blevesearch/bleve/index/store"
	"github.com/blevesearch/bleve/registry"
	bolt "github.com/etcd-io/bbolt"
)

const (
	Name                    = "boltdb"
	defaultCompactBatchSize = 100
)

type Store struct {
	path        string
	bucket      string
	db          *bolt.DB
	noSync      bool
	fillPercent float64
	mo          store.MergeOperator
}

func New(mo store.MergeOperator, config map[string]interface{}) (store.KVStore, error) {
	path, ok := config["path"].(string)
	if !ok {
		return nil, fmt.Errorf("must specify path")
	}
	if path == "" {
		return nil, os.ErrInvalid
	}

	bucket, ok := config["bucket"].(string)
	if !ok {
		bucket = "bleve"
	}

	noSync, _ := config["nosync"].(bool)

	fillPercent, ok := config["fillPercent"].(float64)
	if !ok {
		fillPercent = bolt.DefaultFillPercent
	}

	bo := &bolt.Options{}
	ro, ok := config["read_only"].(bool)
	if ok {
		bo.ReadOnly = ro
	}

	if initialMmapSize, ok := config["initialMmapSize"].(int); ok {
		bo.InitialMmapSize = initialMmapSize
	} else if initialMmapSize, ok := config["initialMmapSize"].(float64); ok {
		bo.InitialMmapSize = int(initialMmapSize)
	}

	db, err := bolt.Open(path, 0600, bo)
	if err != nil {
		return nil, err
	}
	db.NoSync = noSync

	if !bo.ReadOnly {
		err = db.Update(func(tx *bolt.Tx) error {
			_, err := tx.CreateBucketIfNotExists([]byte(bucket))

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

	rv := Store{
		path:        path,
		bucket:      bucket,
		db:          db,
		mo:          mo,
		noSync:      noSync,
		fillPercent: fillPercent,
	}
	return &rv, nil
}

func (bs *Store) Close() error {
	return bs.db.Close()
}

func (bs *Store) Reader() (store.KVReader, error) {
	tx, err := bs.db.Begin(false)
	if err != nil {
		return nil, err
	}
	return &Reader{
		store:  bs,
		tx:     tx,
		bucket: tx.Bucket([]byte(bs.bucket)),
	}, nil
}

func (bs *Store) Writer() (store.KVWriter, error) {
	return &Writer{
		store: bs,
	}, nil
}

func (bs *Store) Stats() json.Marshaler {
	return &stats{
		s: bs,
	}
}

// CompactWithBatchSize removes DictionaryTerm entries with a count of zero (in batchSize batches)
// Removing entries is a workaround for github issue #374.
func (bs *Store) CompactWithBatchSize(batchSize int) error {
	for {
		cnt := 0
		err := bs.db.Batch(func(tx *bolt.Tx) error {
			c := tx.Bucket([]byte(bs.bucket)).Cursor()
			prefix := []byte("d")

			for k, v := c.Seek(prefix); bytes.HasPrefix(k, prefix); k, v = c.Next() {
				if bytes.Equal(v, []byte{0}) {
					cnt++
					if err := c.Delete(); err != nil {
						return err
					}
					if cnt == batchSize {
						break
					}
				}

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

		if cnt == 0 {
			break
		}
	}
	return nil
}

// Compact calls CompactWithBatchSize with a default batch size of 100.  This is a workaround
// for github issue #374.
func (bs *Store) Compact() error {
	return bs.CompactWithBatchSize(defaultCompactBatchSize)
}

func init() {
	registry.RegisterKVStore(Name, New)
}
<|endoftext|>


<commit_before>package parse

import (
	"strings"
	"testing"
)

func TestStandardAddress(t *testing.T) {
	for i, test := range []struct {
		input      string
		host, port string
		shouldErr  bool
	}{
		{`localhost`, "localhost", "", false},
		{`localhost:1234`, "localhost", "1234", false},
		{`localhost:`, "localhost", "", false},
		{`0.0.0.0`, "0.0.0.0", "", false},
		{`127.0.0.1:1234`, "127.0.0.1", "1234", false},
		{`:1234`, "", "1234", false},
		{`[::1]`, "::1", "", false},
		{`[::1]:1234`, "::1", "1234", false},
		{`:`, "", "", false},
		{`localhost:http`, "localhost", "http", false},
		{`localhost:https`, "localhost", "https", false},
		{`:http`, "", "http", false},
		{`:https`, "", "https", false},
		{`http://localhost`, "localhost", "http", false},
		{`https://localhost`, "localhost", "https", false},
		{`http://127.0.0.1`, "127.0.0.1", "http", false},
		{`https://127.0.0.1`, "127.0.0.1", "https", false},
		{`http://[::1]`, "::1", "http", false},
		{`http://localhost:1234`, "localhost", "1234", false},
		{`https://127.0.0.1:1234`, "127.0.0.1", "1234", false},
		{`http://[::1]:1234`, "::1", "1234", false},
		{``, "", "", false},
		{`::1`, "::1", "", true},
		{`localhost::`, "localhost::", "", true},
		{`#$%@`, "#$%@", "", true},
	} {
		host, port, err := standardAddress(test.input)

		if err != nil && !test.shouldErr {
			t.Errorf("Test %d: Expected no error, but had error: %v", i, err)
		}
		if err == nil && test.shouldErr {
			t.Errorf("Test %d: Expected error, but had none", i)
		}

		if host != test.host {
			t.Errorf("Test %d: Expected host '%s', got '%s'", i, test.host, host)
		}

		if port != test.port {
			t.Errorf("Test %d: Expected port '%s', got '%s'", i, test.port, port)
		}
	}
}

func TestParseOne(t *testing.T) {
	testParseOne := func(input string, shouldErr bool) (multiServerBlock, error) {
		p := testParser(input)
		p.Next()
		err := p.parseOne()
		return p.block, err
	}

	// Set up some bogus directives for testing
	ValidDirectives = map[string]struct{}{
		"dir1": struct{}{},
		"dir2": struct{}{},
		"dir3": struct{}{},
	}

	for i, test := range []struct {
		input     string
		shouldErr bool
		addresses []address
		tokens    map[string]int // map of directive name to number of tokens expected
	}{
		{`localhost`, false, []address{
			{"localhost", ""},
		}, map[string]int{}},

		{`localhost
		  dir1`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 1,
		}},

		{`localhost:1234
		  dir1 foo bar`, false, []address{
			{"localhost", "1234"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost {
		    dir1
		  }`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 1,
		}},

		{`localhost:1234 {
		    dir1 foo bar
		    dir2
		  }`, false, []address{
			{"localhost", "1234"},
		}, map[string]int{
			"dir1": 3,
			"dir2": 1,
		}},

		{`http://localhost https://localhost
		  dir1 foo bar`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost https://localhost {
		    dir1 foo bar
		  }`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost, https://localhost {
		    dir1 foo bar
		  }`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost, {
		  }`, true, []address{
			{"localhost", "http"},
		}, map[string]int{}},

		{`host1:80, http://host2.com
		  dir1 foo bar
		  dir2 baz`, false, []address{
			{"host1", "80"},
			{"host2.com", "http"},
		}, map[string]int{
			"dir1": 3,
			"dir2": 2,
		}},

		{`http://host1.com,
		  http://host2.com,
		  https://host3.com`, false, []address{
			{"host1.com", "http"},
			{"host2.com", "http"},
			{"host3.com", "https"},
		}, map[string]int{}},

		{`http://host1.com:1234, https://host2.com
		  dir1 foo {
		    bar baz
		  }
		  dir2`, false, []address{
			{"host1.com", "1234"},
			{"host2.com", "https"},
		}, map[string]int{
			"dir1": 6,
			"dir2": 1,
		}},

		{`127.0.0.1
		  dir1 {
		    bar baz
		  }
		  dir2 {
		    foo bar
		  }`, false, []address{
			{"127.0.0.1", ""},
		}, map[string]int{
			"dir1": 5,
			"dir2": 5,
		}},

		{`127.0.0.1
		  unknown_directive`, true, []address{
			{"127.0.0.1", ""},
		}, map[string]int{}},

		{`localhost
		  dir1 {
		    foo`, true, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		  }`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		    nested {
		      foo
		    }
		  }
		  dir2 foo bar`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 7,
			"dir2": 3,
		}},

		{``, false, []address{}, map[string]int{}},
	} {
		result, err := testParseOne(test.input, test.shouldErr)

		if test.shouldErr && err == nil {
			t.Errorf("Test %d: Expected an error, but didn't get one", i)
		}
		if !test.shouldErr && err != nil {
			t.Errorf("Test %d: Expected no error, but got: %v", i, err)
		}

		if len(result.addresses) != len(test.addresses) {
			t.Errorf("Test %d: Expected %d addresses, got %d",
				i, len(test.addresses), len(result.addresses))
			continue
		}
		for j, addr := range result.addresses {
			if addr.host != test.addresses[j].host {
				t.Errorf("Test %d, address %d: Expected host to be '%s', but was '%s'",
					i, j, test.addresses[j].host, addr.host)
			}
			if addr.port != test.addresses[j].port {
				t.Errorf("Test %d, address %d: Expected port to be '%s', but was '%s'",
					i, j, test.addresses[j].port, addr.port)
			}
		}

		if len(result.tokens) != len(test.tokens) {
			t.Errorf("Test %d: Expected %d directives, had %d",
				i, len(test.tokens), len(result.tokens))
			continue
		}
		for directive, tokens := range result.tokens {
			if len(tokens) != test.tokens[directive] {
				t.Errorf("Test %d, directive '%s': Expected %d tokens, counted %d",
					i, directive, test.tokens[directive], len(tokens))
				continue
			}
		}
	}

}

func testParser(input string) parser {
	buf := strings.NewReader(input)
	p := parser{Dispenser: NewDispenser("Test", buf)}
	return p
}
<commit_msg>More parser tests<commit_after>package parse

import (
	"reflect"
	"strings"
	"testing"
)

func TestStandardAddress(t *testing.T) {
	for i, test := range []struct {
		input      string
		host, port string
		shouldErr  bool
	}{
		{`localhost`, "localhost", "", false},
		{`localhost:1234`, "localhost", "1234", false},
		{`localhost:`, "localhost", "", false},
		{`0.0.0.0`, "0.0.0.0", "", false},
		{`127.0.0.1:1234`, "127.0.0.1", "1234", false},
		{`:1234`, "", "1234", false},
		{`[::1]`, "::1", "", false},
		{`[::1]:1234`, "::1", "1234", false},
		{`:`, "", "", false},
		{`localhost:http`, "localhost", "http", false},
		{`localhost:https`, "localhost", "https", false},
		{`:http`, "", "http", false},
		{`:https`, "", "https", false},
		{`http://localhost`, "localhost", "http", false},
		{`https://localhost`, "localhost", "https", false},
		{`http://127.0.0.1`, "127.0.0.1", "http", false},
		{`https://127.0.0.1`, "127.0.0.1", "https", false},
		{`http://[::1]`, "::1", "http", false},
		{`http://localhost:1234`, "localhost", "1234", false},
		{`https://127.0.0.1:1234`, "127.0.0.1", "1234", false},
		{`http://[::1]:1234`, "::1", "1234", false},
		{``, "", "", false},
		{`::1`, "::1", "", true},
		{`localhost::`, "localhost::", "", true},
		{`#$%@`, "#$%@", "", true},
	} {
		host, port, err := standardAddress(test.input)

		if err != nil && !test.shouldErr {
			t.Errorf("Test %d: Expected no error, but had error: %v", i, err)
		}
		if err == nil && test.shouldErr {
			t.Errorf("Test %d: Expected error, but had none", i)
		}

		if host != test.host {
			t.Errorf("Test %d: Expected host '%s', got '%s'", i, test.host, host)
		}

		if port != test.port {
			t.Errorf("Test %d: Expected port '%s', got '%s'", i, test.port, port)
		}
	}
}

func TestParseOne(t *testing.T) {
	setupParseTests()

	testParseOne := func(input string) (multiServerBlock, error) {
		p := testParser(input)
		p.Next()
		err := p.parseOne()
		return p.block, err
	}

	for i, test := range []struct {
		input     string
		shouldErr bool
		addresses []address
		tokens    map[string]int // map of directive name to number of tokens expected
	}{
		{`localhost`, false, []address{
			{"localhost", ""},
		}, map[string]int{}},

		{`localhost
		  dir1`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 1,
		}},

		{`localhost:1234
		  dir1 foo bar`, false, []address{
			{"localhost", "1234"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost {
		    dir1
		  }`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 1,
		}},

		{`localhost:1234 {
		    dir1 foo bar
		    dir2
		  }`, false, []address{
			{"localhost", "1234"},
		}, map[string]int{
			"dir1": 3,
			"dir2": 1,
		}},

		{`http://localhost https://localhost
		  dir1 foo bar`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost https://localhost {
		    dir1 foo bar
		  }`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost, https://localhost {
		    dir1 foo bar
		  }`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost, {
		  }`, true, []address{
			{"localhost", "http"},
		}, map[string]int{}},

		{`host1:80, http://host2.com
		  dir1 foo bar
		  dir2 baz`, false, []address{
			{"host1", "80"},
			{"host2.com", "http"},
		}, map[string]int{
			"dir1": 3,
			"dir2": 2,
		}},

		{`http://host1.com,
		  http://host2.com,
		  https://host3.com`, false, []address{
			{"host1.com", "http"},
			{"host2.com", "http"},
			{"host3.com", "https"},
		}, map[string]int{}},

		{`http://host1.com:1234, https://host2.com
		  dir1 foo {
		    bar baz
		  }
		  dir2`, false, []address{
			{"host1.com", "1234"},
			{"host2.com", "https"},
		}, map[string]int{
			"dir1": 6,
			"dir2": 1,
		}},

		{`127.0.0.1
		  dir1 {
		    bar baz
		  }
		  dir2 {
		    foo bar
		  }`, false, []address{
			{"127.0.0.1", ""},
		}, map[string]int{
			"dir1": 5,
			"dir2": 5,
		}},

		{`127.0.0.1
		  unknown_directive`, true, []address{
			{"127.0.0.1", ""},
		}, map[string]int{}},

		{`localhost
		  dir1 {
		    foo`, true, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		  }`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		    nested {
		      foo
		    }
		  }
		  dir2 foo bar`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 7,
			"dir2": 3,
		}},

		{``, false, []address{}, map[string]int{}},
	} {
		result, err := testParseOne(test.input)

		if test.shouldErr && err == nil {
			t.Errorf("Test %d: Expected an error, but didn't get one", i)
		}
		if !test.shouldErr && err != nil {
			t.Errorf("Test %d: Expected no error, but got: %v", i, err)
		}

		if len(result.addresses) != len(test.addresses) {
			t.Errorf("Test %d: Expected %d addresses, got %d",
				i, len(test.addresses), len(result.addresses))
			continue
		}
		for j, addr := range result.addresses {
			if addr.host != test.addresses[j].host {
				t.Errorf("Test %d, address %d: Expected host to be '%s', but was '%s'",
					i, j, test.addresses[j].host, addr.host)
			}
			if addr.port != test.addresses[j].port {
				t.Errorf("Test %d, address %d: Expected port to be '%s', but was '%s'",
					i, j, test.addresses[j].port, addr.port)
			}
		}

		if len(result.tokens) != len(test.tokens) {
			t.Errorf("Test %d: Expected %d directives, had %d",
				i, len(test.tokens), len(result.tokens))
			continue
		}
		for directive, tokens := range result.tokens {
			if len(tokens) != test.tokens[directive] {
				t.Errorf("Test %d, directive '%s': Expected %d tokens, counted %d",
					i, directive, test.tokens[directive], len(tokens))
				continue
			}
		}
	}
}

func TestParseAll(t *testing.T) {
	setupParseTests()

	for i, test := range []struct {
		input     string
		shouldErr bool
		addresses []address // one per expected server block, in order
	}{
		{`localhost`, false, []address{
			{"localhost", ""},
		}},

		{`localhost:1234`, false, []address{
			{"localhost", "1234"},
		}},

		{`localhost:1234 {
		  }
		  localhost:2015 {
		  }`, false, []address{
			{"localhost", "1234"},
			{"localhost", "2015"},
		}},

		{`localhost:1234, http://host2`, false, []address{
			{"localhost", "1234"},
			{"host2", "http"},
		}},

		{`http://host1.com, http://host2.com {
		  }
		  https://host3.com, https://host4.com {
		  }`, false, []address{
			{"host1.com", "http"},
			{"host2.com", "http"},
			{"host3.com", "https"},
			{"host4.com", "https"},
		}},
	} {
		p := testParser(test.input)
		blocks, err := p.parseAll()

		if test.shouldErr && err == nil {
			t.Errorf("Test %d: Expected an error, but didn't get one", i)
		}
		if !test.shouldErr && err != nil {
			t.Errorf("Test %d: Expected no error, but got: %v", i, err)
		}

		if len(blocks) != len(test.addresses) {
			t.Errorf("Test %d: Expected %d server blocks, got %d",
				i, len(test.addresses), len(blocks))
			continue
		}
		for j, block := range blocks {
			if block.Host != test.addresses[j].host {
				t.Errorf("Test %d, block %d: Expected host to be '%s', but was '%s'",
					i, j, test.addresses[j].host, block.Host)
			}
			if block.Port != test.addresses[j].port {
				t.Errorf("Test %d, block %d: Expected port to be '%s', but was '%s'",
					i, j, test.addresses[j].port, block.Port)
			}
		}
	}

	// Exploding the server blocks that have more than one address should replicate/share tokens
	p := testParser(`host1 {
	                   dir1 foo bar
	                 }

	                 host2, host3 {
	                   dir2 foo bar
	                   dir3 foo {
	                     bar
	                   }
	                 }`)
	blocks, err := p.parseAll()
	if err != nil {
		t.Fatal("Expected there to not be an error, but there was: %v", err)
	}

	if !reflect.DeepEqual(blocks[1].Tokens, blocks[2].Tokens) {
		t.Errorf("Expected host2 and host3 to have same tokens, but they didn't.\nhost2 Block: %v\nhost3 Block: %v",
			blocks[1].Tokens, blocks[2].Tokens)
	}
}

func setupParseTests() {
	// Set up some bogus directives for testing
	ValidDirectives = map[string]struct{}{
		"dir1": struct{}{},
		"dir2": struct{}{},
		"dir3": struct{}{},
	}
}

func testParser(input string) parser {
	buf := strings.NewReader(input)
	p := parser{Dispenser: NewDispenser("Test", buf)}
	return p
}
<|endoftext|>


<commit_before>// go-qrcode
// Copyright 2014 Tom Harwood

package main

import (
	"flag"
	"fmt"
	"os"
	"strings"

	qrcode "github.com/skip2/go-qrcode"
)

func main() {
	outFile := flag.String("o", "", "out PNG file prefix, empty for stdout")
	size := flag.Int("s", 256, "image size (pixel)")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `qrcode -- QR Code encoder in Go
https://github.com/skip2/go-qrcode

Flags:
`)
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, `
Usage:
  1. Arguments except for flags are joined by " " and used to generate QR code.
     Default output is STDOUT, pipe to imagemagick command "display" to display
     on any X server.

       qrcode hello word | display

  2. Save to file if "display" not available:

       qrcode "homepage: https://github.com/skip2/go-qrcode" > out.png

`)
	}
	flag.Parse()

	if *size <= 0 {
		checkError(fmt.Errorf("Error: value of -s should > 0"))
	}
	if len(flag.Args()) == 0 {
		flag.Usage()
		checkError(fmt.Errorf("Error: no content given"))
	}

	content := strings.Join(flag.Args(), " ")

	var err error
	var q *qrcode.QRCode
	q, err = qrcode.New(content, qrcode.Highest)
	checkError(err)

	var png []byte
	png, err = q.PNG(*size)
	checkError(err)

	if *outFile == "" {
		os.Stdout.Write(png)
	} else {
		var fh *os.File
		fh, err = os.Create(*outFile + ".png")
		checkError(err)
		defer fh.Close()
		fh.Write(png)
	}
}

func checkError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		os.Exit(0)
	}
}
<commit_msg>Extra flag -t to print as 2D char map on STDOUT.<commit_after>// go-qrcode
// Copyright 2014 Tom Harwood

package main

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"strings"

	qrcode "github.com/skip2/go-qrcode"
)

func main() {
	outFile := flag.String("o", "", "out PNG file prefix, empty for stdout")
	size := flag.Int("s", 256, "image size (pixel)")
	textArt := flag.Bool("t", false, "print as text-art on stdout")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `qrcode -- QR Code encoder in Go
https://github.com/skip2/go-qrcode

Flags:
`)
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, `
Usage:
  1. Arguments except for flags are joined by " " and used to generate QR code.
     Default output is STDOUT, pipe to imagemagick command "display" to display
     on any X server.

       qrcode hello word | display

  2. Save to file if "display" not available:

       qrcode "homepage: https://github.com/skip2/go-qrcode" > out.png

`)
	}
	flag.Parse()

	if *size <= 0 {
		checkError(fmt.Errorf("Error: value of -s should > 0"))
	}
	if len(flag.Args()) == 0 {
		flag.Usage()
		checkError(fmt.Errorf("Error: no content given"))
	}

	content := strings.Join(flag.Args(), " ")

	var err error
	var q *qrcode.QRCode
	q, err = qrcode.New(content, qrcode.Highest)
	checkError(err)

	if *textArt {
		art := qr2String(q, true)
		fmt.Println(art)
		return
	}

	var png []byte
	png, err = q.PNG(*size)
	checkError(err)

	if *outFile == "" {
		os.Stdout.Write(png)
	} else {
		var fh *os.File
		fh, err = os.Create(*outFile + ".png")
		checkError(err)
		defer fh.Close()
		fh.Write(png)
	}
}

func checkError(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		os.Exit(1)
	}
}

// qr2String produces a multi-line string that forms a QR-code image.
func qr2String(qr *qrcode.QRCode, inverseColor bool) string {
	bits := qr.Bitmap()
	var buf bytes.Buffer
	for y := range bits {
		for x := range bits[y] {
			if bits[y][x] != inverseColor {
				buf.WriteString("  ")
			} else {
				buf.WriteString("■■")
			}
		}
		buf.WriteString("\n")
	}
	return buf.String()
}
<|endoftext|>


<commit_before>package queue

import (
	"container/heap"
	"math/big"

	peer "github.com/jbenet/go-ipfs/peer"
	ks "github.com/jbenet/go-ipfs/routing/keyspace"
	u "github.com/jbenet/go-ipfs/util"
)

// peerDistance tracks a peer and its distance to something else.
type peerDistance struct {
	// the peer
	peer *peer.Peer

	// big.Int for XOR metric
	distance *big.Int
}

// distancePQ implements heap.Interface and PeerQueue
type distancePQ struct {
	// from is the Key this PQ measures against
	from ks.Key

	// peers is a heap of peerDistance items
	peers []*peerDistance
}

func (pq distancePQ) Len() int {
	return len(pq.peers)
}

func (pq distancePQ) Less(i, j int) bool {
	return -1 == pq.peers[i].distance.Cmp(pq.peers[j].distance)
}

func (pq distancePQ) Swap(i, j int) {
	p := pq.peers
	p[i], p[j] = p[j], p[i]
}

func (pq *distancePQ) Push(x interface{}) {
	item := x.(*peerDistance)
	pq.peers = append(pq.peers, item)
}

func (pq *distancePQ) Pop() interface{} {
	old := pq.peers
	n := len(old)
	item := old[n-1]
	pq.peers = old[0 : n-1]
	return item
}

func (pq *distancePQ) Enqueue(p *peer.Peer) {
	distance := ks.XORKeySpace.Key(p.ID).Distance(pq.from)

	heap.Push(pq, &peerDistance{
		peer:     p,
		distance: distance,
	})
}

func (pq *distancePQ) Dequeue() *peer.Peer {
	if len(pq.peers) < 1 {
		panic("called Dequeue on an empty PeerQueue")
		// will panic internally anyway, but we can help debug here
	}

	o := heap.Pop(pq)
	p := o.(*peerDistance)
	return p.peer
}

// NewXORDistancePQ returns a PeerQueue which maintains its peers sorted
// in terms of their distances to each other in an XORKeySpace (i.e. using
// XOR as a metric of distance).
func NewXORDistancePQ(fromKey u.Key) PeerQueue {
	return &distancePQ{
		from:  ks.XORKeySpace.Key([]byte(fromKey)),
		peers: []*peerDistance{},
	}
}
<commit_msg>sync safety to pq<commit_after>package queue

import (
	"container/heap"
	"math/big"
	"sync"

	peer "github.com/jbenet/go-ipfs/peer"
	ks "github.com/jbenet/go-ipfs/routing/keyspace"
	u "github.com/jbenet/go-ipfs/util"
)

// peerDistance tracks a peer and its distance to something else.
type peerDistance struct {
	// the peer
	peer *peer.Peer

	// big.Int for XOR metric
	distance *big.Int
}

// distancePQ implements heap.Interface and PeerQueue
type distancePQ struct {
	// from is the Key this PQ measures against
	from ks.Key

	// peers is a heap of peerDistance items
	peers []*peerDistance

	sync.RWMutex
}

func (pq *distancePQ) Len() int {
	return len(pq.peers)
}

func (pq *distancePQ) Less(i, j int) bool {
	return -1 == pq.peers[i].distance.Cmp(pq.peers[j].distance)
}

func (pq *distancePQ) Swap(i, j int) {
	p := pq.peers
	p[i], p[j] = p[j], p[i]
}

func (pq *distancePQ) Push(x interface{}) {
	item := x.(*peerDistance)
	pq.peers = append(pq.peers, item)
}

func (pq *distancePQ) Pop() interface{} {
	old := pq.peers
	n := len(old)
	item := old[n-1]
	pq.peers = old[0 : n-1]
	return item
}

func (pq *distancePQ) Enqueue(p *peer.Peer) {
	pq.Lock()
	defer pq.Unlock()

	distance := ks.XORKeySpace.Key(p.ID).Distance(pq.from)

	heap.Push(pq, &peerDistance{
		peer:     p,
		distance: distance,
	})
}

func (pq *distancePQ) Dequeue() *peer.Peer {
	pq.Lock()
	defer pq.Unlock()

	if len(pq.peers) < 1 {
		panic("called Dequeue on an empty PeerQueue")
		// will panic internally anyway, but we can help debug here
	}

	o := heap.Pop(pq)
	p := o.(*peerDistance)
	return p.peer
}

// NewXORDistancePQ returns a PeerQueue which maintains its peers sorted
// in terms of their distances to each other in an XORKeySpace (i.e. using
// XOR as a metric of distance).
func NewXORDistancePQ(fromKey u.Key) PeerQueue {
	return &distancePQ{
		from:  ks.XORKeySpace.Key([]byte(fromKey)),
		peers: []*peerDistance{},
	}
}
<|endoftext|>


<commit_before>package docgen

import (
	"fmt"
	"strings"
)

func (c *Context) Markdown() string {
	out := `### Variables
| Name | Type |
|------|------|
`
	for ident, v := range c.Variables {
		if v.Kind == "func" {
			continue
		}
		if v.Kind == "operator" {
			continue
		}
		out += fmt.Sprintf("| %v | %v |\n", ident, link(v))
	}

	out += `
### Functions
| Name | Return type |
|------|-------------|
`
	for ident, v := range c.Variables {
		if v.Kind == "func" {
			args := make([]string, len(v.Arguments))
			for i, arg := range v.Arguments {
				args[i] = link(arg)
			}
			out += fmt.Sprintf("| %v(%v) | %v |\n", ident, strings.Join(args, ", "), link(v.Return))
		}
	}

	out += "\n### Types\n"
	for name, t := range c.Types {
		out += fmt.Sprintf("#### %v\n", name)
		out += fields(t)
		out += "\n"
	}

	return out
}

func link(t *Type) string {
	if t == nil {
		return "nil"
	}
	if t.Name != "" {
		return fmt.Sprintf("[%v](#%v)", t.Name, t.Name)
	}
	if t.Kind == "array" {
		return fmt.Sprintf("array(%v)", link(t.Type))
	}
	if t.Kind == "map" {
		return fmt.Sprintf("map(%v => %v)", link(t.Key), link(t.Type))
	}
	return fmt.Sprintf("`%v`", t.Kind)
}

func fields(t *Type) string {
	out := ""
	foundFields := false
	for ident, v := range t.Fields {
		if v.Kind != "func" {
			if !foundFields {
				out += "| Field | Type |\n|---|---|\n"
			}
			foundFields = true

			out += fmt.Sprintf("| %v | %v |\n", ident, link(v))
		}
	}
	foundMethod := false
	for ident, v := range t.Fields {
		if v.Kind == "func" {
			if !foundMethod {
				out += "\n| Method | Returns |\n|---|---|\n"
			}
			foundMethod = true

			args := make([]string, len(v.Arguments))
			for i, arg := range v.Arguments {
				args[i] = link(arg)
			}
			out += fmt.Sprintf("| %v(%v) | %v |\n", ident, strings.Join(args, ", "), link(v.Return))
		}
	}
	return out
}
<commit_msg>Output sorted fields in markdown<commit_after>package docgen

import (
	"fmt"
	"sort"
	"strings"
)

func (c *Context) Markdown() string {
	var variables []string
	for name := range c.Variables {
		variables = append(variables, string(name))
	}

	var types []string
	for name := range c.Types {
		types = append(types, string(name))
	}

	sort.Strings(variables)
	sort.Strings(types)

	out := `### Variables
| Name | Type |
|------|------|
`
	for _, name := range variables {
		v := c.Variables[Identifier(name)]
		if v.Kind == "func" {
			continue
		}
		if v.Kind == "operator" {
			continue
		}
		out += fmt.Sprintf("| %v | %v |\n", name, link(v))
	}

	out += `
### Functions
| Name | Return type |
|------|-------------|
`
	for _, name := range variables {
		v := c.Variables[Identifier(name)]
		if v.Kind == "func" {
			args := make([]string, len(v.Arguments))
			for i, arg := range v.Arguments {
				args[i] = link(arg)
			}
			out += fmt.Sprintf("| %v(%v) | %v |\n", name, strings.Join(args, ", "), link(v.Return))
		}
	}

	out += "\n### Types\n"
	for _, name := range types {
		t := c.Types[TypeName(name)]
		out += fmt.Sprintf("#### %v\n", name)
		out += fields(t)
		out += "\n"
	}

	return out
}

func link(t *Type) string {
	if t == nil {
		return "nil"
	}
	if t.Name != "" {
		return fmt.Sprintf("[%v](#%v)", t.Name, t.Name)
	}
	if t.Kind == "array" {
		return fmt.Sprintf("array(%v)", link(t.Type))
	}
	if t.Kind == "map" {
		return fmt.Sprintf("map(%v => %v)", link(t.Key), link(t.Type))
	}
	return fmt.Sprintf("`%v`", t.Kind)
}

func fields(t *Type) string {
	var fields []string
	for field := range t.Fields {
		fields = append(fields, string(field))
	}
	sort.Strings(fields)

	out := ""
	foundFields := false
	for _, name := range fields {
		v := t.Fields[Identifier(name)]
		if v.Kind != "func" {
			if !foundFields {
				out += "| Field | Type |\n|---|---|\n"
			}
			foundFields = true

			out += fmt.Sprintf("| %v | %v |\n", name, link(v))
		}
	}
	foundMethod := false
	for _, name := range fields {
		v := t.Fields[Identifier(name)]
		if v.Kind == "func" {
			if !foundMethod {
				out += "\n| Method | Returns |\n|---|---|\n"
			}
			foundMethod = true

			args := make([]string, len(v.Arguments))
			for i, arg := range v.Arguments {
				args[i] = link(arg)
			}
			out += fmt.Sprintf("| %v(%v) | %v |\n", name, strings.Join(args, ", "), link(v.Return))
		}
	}
	return out
}
<|endoftext|>


<commit_before>package clever

import (
	"code.google.com/p/goauth2/oauth"
	"encoding/json"
	"fmt"
	mock "github.com/Clever/clever-go/mock"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
)

var dummytransport = &oauth.Transport{
	Token: &oauth.Token{AccessToken: "doesntmatter"},
}

func TestBasicAuthTransport(t *testing.T) {
	bat := &BasicAuthTransport{"user", "pass"}
	client := bat.Client()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Basic dXNlcjpwYXNz" {
			t.Fatal("unexpected auth header")
		}
	}))
	if _, err := client.Get(ts.URL); err != nil {
		t.Fatal(err)
	}
}

func TestQueryDistricts(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/districts", nil)

	if !results.Next() {
		t.Fatal("Found no districts")
	}
	district0 := District{}
	if err := results.Scan(&district0); err != nil {
		t.Fatalf("Error retrieving district: %s\n", err)
	}

	resp := DistrictResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/districts/%s", district0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving district: %s\n", err)
	}

	expectedDistrict0 := District{
		Id:   "51a5a56312ec00cc5100007e",
		Name: "test district",
	}
	if !reflect.DeepEqual(expectedDistrict0, district0) {
		t.Fatal("District did not match expected.")
	}
}

func TestQuerySchools(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/schools", nil)
	if !results.Next() {
		t.Fatal("Found no schools")
	}
	school0 := School{}
	if err := results.Scan(&school0); err != nil {
		t.Fatalf("Error retrieving school: %s\n", err)
	}

	resp := SchoolResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/schools/%s", school0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving school: %s\n", err)
	}

	expectedSchool0 := School{
		Created:      "2012-11-06T00:00:00Z",
		District:     "51a5a56312ec00cc5100007e",
		HighGrade:    "8",
		Id:           "4fee004cca2e43cf27000002",
		LastModified: "2012-11-07T00:44:53.079Z",
		Location: Location{
			Address: "42139 Fadel Mountains",
			City:    "Thomasfort",
			State:   "AP",
			Zip:     "61397-3760",
		},
		LowGrade:     "6",
		Name:         "Clever Preparatory",
		NcesId:       "94755881",
		Phone:        "(527) 825-2248",
		SchoolNumber: "2559",
		SisId:        "2559",
		StateId:      "23",
	}
	if !reflect.DeepEqual(expectedSchool0, school0) {
		t.Fatalf("School did not match expected.")
	}
}

func TestQueryTeachers(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/teachers", nil)
	if !results.Next() {
		t.Fatal("Found no teachers")
	}
	teacher0 := Teacher{}
	if err := results.Scan(&teacher0); err != nil {
		t.Fatalf("Error retrieving teacher: %s\n", err)
	}

	resp := TeacherResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/teachers/%s", teacher0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving teachers: %s\n", err)
	}

	expectedTeacher0 := Teacher{
		Created:      "2013-05-29T06:51:25.139Z",
		District:     "51a5a56312ec00cc5100007e",
		Email:        "t.teacher2@ga4edu.org",
		Id:           "51a5a56d4867bbdf51053c34",
		LastModified: "2013-05-29T06:51:25.145Z",
		Name: Name{
			First:  "Test",
			Middle: "",
			Last:   "Teacher",
		},
		School:        "4fee004cca2e43cf27000002",
		SisId:         "56",
		TeacherNumber: "100",
		Title:         "Math Teacher",
	}
	if !reflect.DeepEqual(expectedTeacher0, teacher0) {
		t.Fatalf("Teacher did not match expected.")
	}
}

func TestQueryStudents(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/students", nil)
	if !results.Next() {
		t.Fatal("Found no students")
	}
	student0 := Student{}
	if err := results.Scan(&student0); err != nil {
		t.Fatalf("Error retrieving student: %s\n", err)
	}

	resp := StudentResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/students/%s", student0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving students: %s\n", err)
	}

	expectedStudent0 := Student{
		District:          "51a5a56312ec00cc5100007e",
		Dob:               "12/12/1998",
		FrlStatus:         "Paid",
		Gender:            "M",
		Grade:             "9",
		HispanicEthnicity: "N",
		Race:              "Caucasian",
		School:            "4fee004cca2e43cf27000002",
		SisId:             "1",
		StateId:           "2237504",
		StudentNumber:     "24772",
		Location: Location{
			Address: "",
			City:    "",
			State:   "",
			Zip:     "",
		},
		Name: Name{
			First:  "John",
			Middle: "",
			Last:   "Doe",
		},
		LastModified: "2013-05-29T06:51:28.006Z",
		Created:      "2013-05-29T06:51:27.997Z",
		Email:        "john.doe@ga4edu.org",
		Id:           "51a5a56f4867bbdf51054054",
	}
	if !reflect.DeepEqual(expectedStudent0, student0) {
		t.Fatalf("Student did not match expected.")
	}
}

func TestQuerySections(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/sections", nil)
	if !results.Next() {
		t.Fatal("Found no sections")
	}
	section0 := Section{}
	if err := results.Scan(&section0); err != nil {
		t.Fatalf("Error retrieving section: %s\n", err)
	}

	resp := SectionResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/sections/%s", section0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving sections: %s\n", err)
	}

	expectedSection0 := Section{
		CourseName:   "test course",
		CourseNumber: "12345",
		Created:      "2013-05-29T06:51:27.997Z",
		District:     "51a5a56312ec00cc5100007e",
		Grade:        "K",
		Id:           "51a5a56f4867bbdf51054054",
		LastModified: "2013-05-29T06:51:28.006Z",
		Name:         "test section",
		School:       "4fee004cca2e43cf27000002",
		SisId:        "1",
		Students:     []string{"51a5a56f4867bbdf51054054"},
		Subject:      "math",
		Teacher:      "51a5a56d4867bbdf51053c34",
		Term: Term{
			Name:      "",
			StartDate: "",
			EndDate:   "",
		},
	}
	if !reflect.DeepEqual(expectedSection0, section0) {
		t.Fatalf("Section did not match expected.")
	}
}

func postDistrictTest(req *http.Request, params map[string]string) error {
	district := District{}
	if err := json.NewDecoder(req.Body).Decode(&district); err != nil {
		return fmt.Errorf("{\"error\":\"%s\"}", err.Error())
	}
	if district.Name != "new name district" {
		return fmt.Errorf("{\"error\":\"%s\"}", "Error formatting post request")
	}
	return nil
}

func TestPostRequest(t *testing.T) {
	clever := New(mock.NewMock(postDistrictTest, "./data"))
	resp := map[string]string{}
	district1 := District{
		Name: "new name district",
	}
	if err := clever.Request("POST", "/v1.1/districts", nil, district1, &resp); err != nil {
		t.Fatalf("Error posting district: %s\n", err)
	}
}

func TestQueryAll(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	result := clever.QueryAll("/v1.1/sections", nil)

	count := 0
	for result.Next() {
		section := Section{}
		result.Scan(&section)
		count++
	}
	if count != 2 {
		t.Fatalf("Did not get both section pages.")
	}
}

func TestTooManyRequestsError(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	result := clever.QueryAll("/mock/rate/limiter", nil)
	result.Next()
	if result.Error() == nil {
		t.Fatalf("Http response 429 (TooManyRequests) did not trigger an error as expected.")
	} else if _, ok := result.Error().(*TooManyRequestsError); !ok {
		t.Fatalf("Http response 429 (TooManyRequests) did not generate the expected error.")
	}
}

func TestHandlesErrors(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	result := clever.QueryAll("/mock/error", nil)
	result.Next()
	if result.Error() == nil {
		t.Fatalf("error endpoint did not trigger an error as expected")
	} else if result.Error().Error() != "there was an error (1337)" {
		t.Fatalf("error endpoint did not generate the expected error: %s", result.Error().Error())
	}
}
<commit_msg>add header test<commit_after>package clever

import (
	"code.google.com/p/goauth2/oauth"
	"encoding/json"
	"fmt"
	mock "github.com/Clever/clever-go/mock"
	"net/http"
	"net/http/httptest"
	"reflect"
	"testing"
)

var dummytransport = &oauth.Transport{
	Token: &oauth.Token{AccessToken: "doesntmatter"},
}

func TestBasicAuthTransport(t *testing.T) {
	bat := &BasicAuthTransport{"user", "pass"}
	client := bat.Client()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Basic dXNlcjpwYXNz" {
			t.Fatal("unexpected auth header")
		}
	}))
	if _, err := client.Get(ts.URL); err != nil {
		t.Fatal(err)
	}
}

func TestQueryDistricts(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/districts", nil)

	if !results.Next() {
		t.Fatal("Found no districts")
	}
	district0 := District{}
	if err := results.Scan(&district0); err != nil {
		t.Fatalf("Error retrieving district: %s\n", err)
	}

	resp := DistrictResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/districts/%s", district0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving district: %s\n", err)
	}

	expectedDistrict0 := District{
		Id:   "51a5a56312ec00cc5100007e",
		Name: "test district",
	}
	if !reflect.DeepEqual(expectedDistrict0, district0) {
		t.Fatal("District did not match expected.")
	}
}

func TestQuerySchools(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/schools", nil)
	if !results.Next() {
		t.Fatal("Found no schools")
	}
	school0 := School{}
	if err := results.Scan(&school0); err != nil {
		t.Fatalf("Error retrieving school: %s\n", err)
	}

	resp := SchoolResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/schools/%s", school0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving school: %s\n", err)
	}

	expectedSchool0 := School{
		Created:      "2012-11-06T00:00:00Z",
		District:     "51a5a56312ec00cc5100007e",
		HighGrade:    "8",
		Id:           "4fee004cca2e43cf27000002",
		LastModified: "2012-11-07T00:44:53.079Z",
		Location: Location{
			Address: "42139 Fadel Mountains",
			City:    "Thomasfort",
			State:   "AP",
			Zip:     "61397-3760",
		},
		LowGrade:     "6",
		Name:         "Clever Preparatory",
		NcesId:       "94755881",
		Phone:        "(527) 825-2248",
		SchoolNumber: "2559",
		SisId:        "2559",
		StateId:      "23",
	}
	if !reflect.DeepEqual(expectedSchool0, school0) {
		t.Fatalf("School did not match expected.")
	}
}

func TestQueryTeachers(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/teachers", nil)
	if !results.Next() {
		t.Fatal("Found no teachers")
	}
	teacher0 := Teacher{}
	if err := results.Scan(&teacher0); err != nil {
		t.Fatalf("Error retrieving teacher: %s\n", err)
	}

	resp := TeacherResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/teachers/%s", teacher0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving teachers: %s\n", err)
	}

	expectedTeacher0 := Teacher{
		Created:      "2013-05-29T06:51:25.139Z",
		District:     "51a5a56312ec00cc5100007e",
		Email:        "t.teacher2@ga4edu.org",
		Id:           "51a5a56d4867bbdf51053c34",
		LastModified: "2013-05-29T06:51:25.145Z",
		Name: Name{
			First:  "Test",
			Middle: "",
			Last:   "Teacher",
		},
		School:        "4fee004cca2e43cf27000002",
		SisId:         "56",
		TeacherNumber: "100",
		Title:         "Math Teacher",
	}
	if !reflect.DeepEqual(expectedTeacher0, teacher0) {
		t.Fatalf("Teacher did not match expected.")
	}
}

func TestQueryStudents(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/students", nil)
	if !results.Next() {
		t.Fatal("Found no students")
	}
	student0 := Student{}
	if err := results.Scan(&student0); err != nil {
		t.Fatalf("Error retrieving student: %s\n", err)
	}

	resp := StudentResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/students/%s", student0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving students: %s\n", err)
	}

	expectedStudent0 := Student{
		District:          "51a5a56312ec00cc5100007e",
		Dob:               "12/12/1998",
		FrlStatus:         "Paid",
		Gender:            "M",
		Grade:             "9",
		HispanicEthnicity: "N",
		Race:              "Caucasian",
		School:            "4fee004cca2e43cf27000002",
		SisId:             "1",
		StateId:           "2237504",
		StudentNumber:     "24772",
		Location: Location{
			Address: "",
			City:    "",
			State:   "",
			Zip:     "",
		},
		Name: Name{
			First:  "John",
			Middle: "",
			Last:   "Doe",
		},
		LastModified: "2013-05-29T06:51:28.006Z",
		Created:      "2013-05-29T06:51:27.997Z",
		Email:        "john.doe@ga4edu.org",
		Id:           "51a5a56f4867bbdf51054054",
	}
	if !reflect.DeepEqual(expectedStudent0, student0) {
		t.Fatalf("Student did not match expected.")
	}
}

func TestQuerySections(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	results := clever.QueryAll("/v1.1/sections", nil)
	if !results.Next() {
		t.Fatal("Found no sections")
	}
	section0 := Section{}
	if err := results.Scan(&section0); err != nil {
		t.Fatalf("Error retrieving section: %s\n", err)
	}

	resp := SectionResp{}
	if err := clever.Query(fmt.Sprintf("/v1.1/sections/%s", section0.Id), nil, &resp); err != nil {
		t.Fatalf("Error retrieving sections: %s\n", err)
	}

	expectedSection0 := Section{
		CourseName:   "test course",
		CourseNumber: "12345",
		Created:      "2013-05-29T06:51:27.997Z",
		District:     "51a5a56312ec00cc5100007e",
		Grade:        "K",
		Id:           "51a5a56f4867bbdf51054054",
		LastModified: "2013-05-29T06:51:28.006Z",
		Name:         "test section",
		School:       "4fee004cca2e43cf27000002",
		SisId:        "1",
		Students:     []string{"51a5a56f4867bbdf51054054"},
		Subject:      "math",
		Teacher:      "51a5a56d4867bbdf51053c34",
		Term: Term{
			Name:      "",
			StartDate: "",
			EndDate:   "",
		},
	}
	if !reflect.DeepEqual(expectedSection0, section0) {
		t.Fatalf("Section did not match expected.")
	}
}

func postDistrictTest(req *http.Request, params map[string]string) error {
	district := District{}
	if err := json.NewDecoder(req.Body).Decode(&district); err != nil {
		return fmt.Errorf("{\"error\":\"%s\"}", err.Error())
	}
	if req.Header.Get("Content-Type") != "application/json" {
		return fmt.Errorf("{\"error\":\"%s\"}", "Error formatting post request header")
	}
	if district.Name != "new name district" {
		return fmt.Errorf("{\"error\":\"%s\"}", "Error formatting post request body")
	}
	return nil
}

func TestPostRequest(t *testing.T) {
	clever := New(mock.NewMock(postDistrictTest, "./data"))
	resp := map[string]string{}
	district1 := District{
		Name: "new name district",
	}
	if err := clever.Request("POST", "/v1.1/districts", nil, district1, &resp); err != nil {
		t.Fatalf("Error posting district: %s\n", err)
	}
}

func TestQueryAll(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	result := clever.QueryAll("/v1.1/sections", nil)

	count := 0
	for result.Next() {
		section := Section{}
		result.Scan(&section)
		count++
	}
	if count != 2 {
		t.Fatalf("Did not get both section pages.")
	}
}

func TestTooManyRequestsError(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	result := clever.QueryAll("/mock/rate/limiter", nil)
	result.Next()
	if result.Error() == nil {
		t.Fatalf("Http response 429 (TooManyRequests) did not trigger an error as expected.")
	} else if _, ok := result.Error().(*TooManyRequestsError); !ok {
		t.Fatalf("Http response 429 (TooManyRequests) did not generate the expected error.")
	}
}

func TestHandlesErrors(t *testing.T) {
	clever := New(mock.NewMock(nil, "./data"))
	result := clever.QueryAll("/mock/error", nil)
	result.Next()
	if result.Error() == nil {
		t.Fatalf("error endpoint did not trigger an error as expected")
	} else if result.Error().Error() != "there was an error (1337)" {
		t.Fatalf("error endpoint did not generate the expected error: %s", result.Error().Error())
	}
}
<|endoftext|>


<commit_before>package dioder

import (
	"image/color"
	"reflect"
	"testing"
)

var (
	piBlasterFile = "/tmp/pi-blaster"
	pinConfiguration = Pins{
		Red: 1,
		Green: 2,
		Blue: 3,
	}
)

func TestNew(t *testing.T) {
	d := New(Pins{}, "")

	if reflect.TypeOf(d).String() != "dioder.Dioder" {
		t.Errorf("Got wrong dioder.Dioder object: %s", reflect.TypeOf(d))
	}
}

func TestDioderGetCurrentColor(t *testing.T) {
	dioder := New(pinConfiguration, piBlasterFile)

	colorSet := dioder.GetCurrentColor()

	if colorSet.A != 0 {
		t.Error("Opacity is not 0")
	}

	if colorSet.R != 0 {
		t.Error("Red is not 0")
	}

	if colorSet.G != 0 {
		t.Error("Green is not 0")
	}

	if colorSet.B != 0 {
		t.Error("Blue is not 0")
	}
}

func TestDioderPinConfiguration(t *testing.T) {
	dioder := New(pinConfiguration, piBlasterFile)

	if dioder.PinConfiguration != pinConfiguration {
		t.Errorf("Pins are not correctly configured. Gave %s, got %s", pinConfiguration, dioder.PinConfiguration)
	}

	if dioder.PiBlaster != piBlasterFile {
		t.Errorf("Pi-Blaster file wrong configured. Gave %s, got %s", piBlasterFile, dioder.PiBlaster)
	}
}

func TestDioderSetAll(t *testing.T) {
	d := New(Pins{}, piBlasterFile)

	//All values at zero
	d.SetAll(color.RGBA{})

	if d.ColorConfiguration.A != 0 {
		t.Error("Opacity is not correct")
	}

	if d.ColorConfiguration.R != 0 {
		t.Error("Red is not correct")
	}

	if d.ColorConfiguration.G != 0 {
		t.Error("Green is not correct")
	}

	if d.ColorConfiguration.B != 0 {
		t.Error("Blue is not correct")
	}

	d.SetAll(color.RGBA{255, 255, 255, 0})
	if d.ColorConfiguration.R != 255 {
		t.Errorf("Red is not correct after applying opacity of zero. Got: %d\n", d.ColorConfiguration.R)
	}

	if d.ColorConfiguration.G != 255 {
		t.Errorf("Green is not correct after applying opacity of zero. Got: %d\n", d.ColorConfiguration.G)
	}

	if d.ColorConfiguration.B != 255 {
		t.Errorf("Blue is not correct after applying opacity of zero. Got: %d\n", d.ColorConfiguration.B)
	}
}

func TestDioderSetPins(t *testing.T) {
	d := New(pinConfiguration, piBlasterFile)

	if d.PinConfiguration.Blue != 3 {
		t.Error("Blue pin is not correct")
	}

	if d.PinConfiguration.Green != 2 {
		t.Error("Green pin is not correct")
	}

	if d.PinConfiguration.Red != 1 {
		t.Error("Red pin is not correct")
	}
}

func TestDioderTurnOff(t *testing.T) {
	d := New(pinConfiguration, piBlasterFile)

	configuration := d.ColorConfiguration

	d.TurnOff()

	if d.ColorConfiguration != configuration {
		t.Errorf("Didn't saved the current settings, had %d - got %d", configuration, d.ColorConfiguration)
	}
}

func TestDioderTurnOn(t *testing.T) {
	d := New(pinConfiguration, piBlasterFile)

	d.TurnOff()
	d.TurnOn()

	if d.ColorConfiguration.A != 100 {
		t.Errorf("Value for opacity is wrong. Expected 0, got %d", d.ColorConfiguration.A)
	}
	if d.ColorConfiguration.R != 255 {
		t.Errorf("Value for red is wrong. Expected 0, got %d", d.ColorConfiguration.R)
	}
	if d.ColorConfiguration.G != 255 {
		t.Errorf("Value for green is wrong. Expected 0, got %d", d.ColorConfiguration.G)
	}
	if d.ColorConfiguration.B != 255 {
		t.Errorf("Value for blue is wrong. Expected 0, got %d", d.ColorConfiguration.B)
	}
}

func TestDioder_SetChannelInteger(t *testing.T) {
	d := New(pinConfiguration, piBlasterFile)

	err := d.SetChannelInteger(255, 1)
	if err != nil {
		t.Errorf("Error: %s", err)
	}
}
<commit_msg>Test Dioder.SetColor()<commit_after>package dioder

import (
	"image/color"
	"reflect"
	"testing"
)

var (
	piBlasterFile = "/tmp/pi-blaster"
	pinConfiguration = Pins{
		Red: 1,
		Green: 2,
		Blue: 3,
	}
)

func TestNew(t *testing.T) {
	d := New(Pins{}, "")

	if reflect.TypeOf(d).String() != "dioder.Dioder" {
		t.Errorf("Got wrong dioder.Dioder object: %s", reflect.TypeOf(d))
	}
}

func TestDioderGetCurrentColor(t *testing.T) {
	dioder := New(pinConfiguration, piBlasterFile)

	colorSet := dioder.GetCurrentColor()

	if colorSet.A != 0 {
		t.Error("Opacity is not 0")
	}

	if colorSet.R != 0 {
		t.Error("Red is not 0")
	}

	if colorSet.G != 0 {
		t.Error("Green is not 0")
	}

	if colorSet.B != 0 {
		t.Error("Blue is not 0")
	}
}

func TestDioderPinConfiguration(t *testing.T) {
	dioder := New(pinConfiguration, piBlasterFile)

	if dioder.PinConfiguration != pinConfiguration {
		t.Errorf("Pins are not correctly configured. Gave %s, got %s", pinConfiguration, dioder.PinConfiguration)
	}

	if dioder.PiBlaster != piBlasterFile {
		t.Errorf("Pi-Blaster file wrong configured. Gave %s, got %s", piBlasterFile, dioder.PiBlaster)
	}
}

func TestDioderSetAll(t *testing.T) {
	d := New(Pins{}, piBlasterFile)

	//All values at zero
	d.SetAll(color.RGBA{})

	if d.ColorConfiguration.A != 0 {
		t.Error("Opacity is not correct")
	}

	if d.ColorConfiguration.R != 0 {
		t.Error("Red is not correct")
	}

	if d.ColorConfiguration.G != 0 {
		t.Error("Green is not correct")
	}

	if d.ColorConfiguration.B != 0 {
		t.Error("Blue is not correct")
	}

	d.SetAll(color.RGBA{255, 255, 255, 0})
	if d.ColorConfiguration.R != 255 {
		t.Errorf("Red is not correct after applying opacity of zero. Got: %d\n", d.ColorConfiguration.R)
	}

	if d.ColorConfiguration.G != 255 {
		t.Errorf("Green is not correct after applying opacity of zero. Got: %d\n", d.ColorConfiguration.G)
	}

	if d.ColorConfiguration.B != 255 {
		t.Errorf("Blue is not correct after applying opacity of zero. Got: %d\n", d.ColorConfiguration.B)
	}
}

func TestDioderSetPins(t *testing.T) {
	d := New(pinConfiguration, piBlasterFile)

	if d.PinConfiguration.Blue != 3 {
		t.Error("Blue pin is not correct")
	}

	if d.PinConfiguration.Green != 2 {
		t.Error("Green pin is not correct")
	}

	if d.PinConfiguration.Red != 1 {
		t.Error("Red pin is not correct")
	}
}

func TestDioderTurnOff(t *testing.T) {
	d := New(pinConfiguration, piBlasterFile)

	configuration := d.ColorConfiguration

	d.TurnOff()

	if d.ColorConfiguration != configuration {
		t.Errorf("Didn't saved the current settings, had %d - got %d", configuration, d.ColorConfiguration)
	}
}

func TestDioderTurnOn(t *testing.T) {
	d := New(pinConfiguration, piBlasterFile)

	d.TurnOff()
	d.TurnOn()

	if d.ColorConfiguration.A != 100 {
		t.Errorf("Value for opacity is wrong. Expected 0, got %d", d.ColorConfiguration.A)
	}
	if d.ColorConfiguration.R != 255 {
		t.Errorf("Value for red is wrong. Expected 0, got %d", d.ColorConfiguration.R)
	}
	if d.ColorConfiguration.G != 255 {
		t.Errorf("Value for green is wrong. Expected 0, got %d", d.ColorConfiguration.G)
	}
	if d.ColorConfiguration.B != 255 {
		t.Errorf("Value for blue is wrong. Expected 0, got %d", d.ColorConfiguration.B)
	}
}

func TestDioder_SetChannelInteger(t *testing.T) {
	d := New(pinConfiguration, piBlasterFile)

	err := d.SetChannelInteger(255, 1)
	if err != nil {
		t.Errorf("Error: %s", err)
	}
}

func TestDioder_SetColor(t *testing.T) {
	d := New(pinConfiguration, "nonExistentFile")

	err := d.SetColor(0, 0)
	if err == nil {
		t.Error("Called with an non existent filename. SetColor() should return an error")
	}

	d.PiBlaster = piBlasterFile

	err = d.SetColor(0, 0)
	if err != nil {
		t.Errorf("Returned an error: %s", err)
	}
}
<|endoftext|>


<commit_before>package workload

import (
	"fmt"
	"io"
	"math/rand"

	"golang.org/x/net/context"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pfs/pfsutil"
	"github.com/pachyderm/pachyderm/src/pps"
	"github.com/pachyderm/pachyderm/src/pps/ppsutil"
)

func RunWorkload(
	pfsClient pfs.APIClient,
	ppsClient pps.APIClient,
	rand *rand.Rand,
	size int,
) error {
	worker := newWorker(rand)
	for i := 0; i < size; i++ {
		if err := worker.work(pfsClient, ppsClient); err != nil {
			return err
		}
	}
	for _, job := range worker.startedJobs {
		jobInfo, err := ppsClient.InspectJob(
			context.Background(),
			&pps.InspectJobRequest{
				Job:        job,
				BlockState: true,
			},
		)
		if err != nil {
			return err
		}
		if jobInfo.State != pps.JobState_JOB_STATE_SUCCESS {
			return fmt.Errorf("job %s failed", job.Id)
		}
	}
	return nil
}

type worker struct {
	repos       []*pfs.Repo
	finished    []*pfs.Commit
	started     []*pfs.Commit
	files       []*pfs.File
	startedJobs []*pps.Job
	jobs        []*pps.Job
	pipelines   []*pps.Pipeline
	rand        *rand.Rand
}

func newWorker(rand *rand.Rand) *worker {
	return &worker{
		rand: rand,
	}
}

const (
	repo     float64 = .02
	commit           = .1
	file             = .9
	job              = .98
	pipeline         = 1.0
)

const maxStartedCommits = 6
const maxStartedJobs = 6

func (w *worker) work(pfsClient pfs.APIClient, ppsClient pps.APIClient) error {
	opt := w.rand.Float64()
	switch {
	case opt < repo:
		repoName := w.randString(10)
		if err := pfsutil.CreateRepo(pfsClient, repoName); err != nil {
			return err
		}
		w.repos = append(w.repos, &pfs.Repo{Name: repoName})
		commit, err := pfsutil.StartCommit(pfsClient, repoName, "")
		if err != nil {
			return err
		}
		w.started = append(w.started, commit)
	case opt < commit:
		if len(w.started) >= maxStartedCommits || len(w.finished) == 0 {
			if len(w.started) == 0 {
				return nil
			}
			i := w.rand.Intn(len(w.started))
			commit := w.started[i]
			if err := pfsutil.FinishCommit(pfsClient, commit.Repo.Name, commit.Id); err != nil {
				return err
			}
			w.started = append(w.started[:i], w.started[i+1:]...)
			w.finished = append(w.finished, commit)
		} else {
			if len(w.finished) == 0 {
				return nil
			}
			commit := w.finished[w.rand.Intn(len(w.finished))]
			commit, err := pfsutil.StartCommit(pfsClient, commit.Repo.Name, commit.Id)
			if err != nil {
				return err
			}
			w.started = append(w.started, commit)
		}
	case opt < file:
		if len(w.started) == 0 {
			return nil
		}
		commit := w.started[w.rand.Intn(len(w.started))]
		if _, err := pfsutil.PutFile(pfsClient, commit.Repo.Name, commit.Id, w.randString(10), 0, w.reader()); err != nil {
			return err
		}
	case opt < job:
		if len(w.finished) == 0 {
			return nil
		}
		if len(w.startedJobs) >= maxStartedJobs {
			job := w.startedJobs[0]
			w.startedJobs = w.startedJobs[1:]
			jobInfo, err := ppsClient.InspectJob(
				context.Background(),
				&pps.InspectJobRequest{
					Job:        job,
					BlockState: true,
				},
			)
			if err != nil {
				return err
			}
			if jobInfo.State != pps.JobState_JOB_STATE_SUCCESS {
				return fmt.Errorf("job %s failed", job.Id)
			}
			w.jobs = append(w.jobs, job)
		} else {
			inputs := [5]string{}
			var inputCommits []*pfs.Commit
			for i := range inputs {
				randI := w.rand.Intn(len(w.finished))
				inputs[i] = w.finished[randI].Repo.Name
				inputCommits = append(inputCommits, w.finished[randI])
			}
			var parentJobID string
			if len(w.jobs) > 0 {
				parentJobID = w.jobs[w.rand.Intn(len(w.jobs))].Id
			}
			outFilename := w.randString(10)
			job, err := ppsutil.CreateJob(
				ppsClient,
				"",
				[]string{"bash"},
				w.grepCmd(inputs, outFilename),
				1,
				inputCommits,
				parentJobID,
			)
			if err != nil {
				return err
			}
			w.startedJobs = append(w.startedJobs, job)
		}
	case opt < pipeline:
		if len(w.repos) == 0 {
			return nil
		}
		inputs := [5]string{}
		var inputRepos []*pfs.Repo
		for i := range inputs {
			randI := w.rand.Intn(len(w.repos))
			inputs[i] = w.repos[randI].Name
			inputRepos = append(inputRepos, w.repos[randI])
		}
		pipelineName := w.randString(10)
		outFilename := w.randString(10)
		if err := ppsutil.CreatePipeline(
			ppsClient,
			pipelineName,
			"",
			[]string{"bash"},
			w.grepCmd(inputs, outFilename),
			1,
			inputRepos,
		); err != nil {
			return err
		}
		w.pipelines = append(w.pipelines, ppsutil.NewPipeline(pipelineName))
	}
	return nil
}

const letters = "abcdefghijklmnopqrstuvwxyz"
const lettersAndSpaces = "abcdefghijklmnopqrstuvwxyz      "

func (w *worker) randString(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = letters[w.rand.Intn(len(letters))]
	}
	return string(b)
}

type reader struct {
	rand  *rand.Rand
	lines int
}

func NewReader(rand *rand.Rand, lines int) io.Reader {
	return &reader{
		rand:  rand,
		lines: lines,
	}
}

func (r *reader) Read(p []byte) (int, error) {
	for i := range p {
		if i%128 == 127 {
			p[i] = '\n'
			r.lines--
		} else {
			p[i] = lettersAndSpaces[r.rand.Intn(len(lettersAndSpaces))]
		}
	}
	p[len(p)-1] = '\n'
	if r.lines <= 0 {
		return len(p), io.EOF
	}
	return len(p), nil
}

func (w *worker) reader() io.Reader {
	return NewReader(w.rand, 1000)
}

func (w *worker) grepCmd(inputs [5]string, outFilename string) string {
	return fmt.Sprintf(
		"grep %s /pfs/{%s,%s,%s,%s,%s}/* >/pfs/out/%s; true",
		w.randString(4),
		inputs[0],
		inputs[1],
		inputs[2],
		inputs[3],
		inputs[4],
		outFilename,
	)
}
<commit_msg>Make workload send unique inputs.<commit_after>package workload

import (
	"fmt"
	"io"
	"math/rand"

	"golang.org/x/net/context"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pfs/pfsutil"
	"github.com/pachyderm/pachyderm/src/pps"
	"github.com/pachyderm/pachyderm/src/pps/ppsutil"
)

func RunWorkload(
	pfsClient pfs.APIClient,
	ppsClient pps.APIClient,
	rand *rand.Rand,
	size int,
) error {
	worker := newWorker(rand)
	for i := 0; i < size; i++ {
		if err := worker.work(pfsClient, ppsClient); err != nil {
			return err
		}
	}
	for _, job := range worker.startedJobs {
		jobInfo, err := ppsClient.InspectJob(
			context.Background(),
			&pps.InspectJobRequest{
				Job:        job,
				BlockState: true,
			},
		)
		if err != nil {
			return err
		}
		if jobInfo.State != pps.JobState_JOB_STATE_SUCCESS {
			return fmt.Errorf("job %s failed", job.Id)
		}
	}
	return nil
}

type worker struct {
	repos       []*pfs.Repo
	finished    []*pfs.Commit
	started     []*pfs.Commit
	files       []*pfs.File
	startedJobs []*pps.Job
	jobs        []*pps.Job
	pipelines   []*pps.Pipeline
	rand        *rand.Rand
}

func newWorker(rand *rand.Rand) *worker {
	return &worker{
		rand: rand,
	}
}

const (
	repo     float64 = .02
	commit           = .1
	file             = .9
	job              = .98
	pipeline         = 1.0
)

const maxStartedCommits = 6
const maxStartedJobs = 6

func (w *worker) work(pfsClient pfs.APIClient, ppsClient pps.APIClient) error {
	opt := w.rand.Float64()
	switch {
	case opt < repo:
		repoName := w.randString(10)
		if err := pfsutil.CreateRepo(pfsClient, repoName); err != nil {
			return err
		}
		w.repos = append(w.repos, &pfs.Repo{Name: repoName})
		commit, err := pfsutil.StartCommit(pfsClient, repoName, "")
		if err != nil {
			return err
		}
		w.started = append(w.started, commit)
	case opt < commit:
		if len(w.started) >= maxStartedCommits || len(w.finished) == 0 {
			if len(w.started) == 0 {
				return nil
			}
			i := w.rand.Intn(len(w.started))
			commit := w.started[i]
			if err := pfsutil.FinishCommit(pfsClient, commit.Repo.Name, commit.Id); err != nil {
				return err
			}
			w.started = append(w.started[:i], w.started[i+1:]...)
			w.finished = append(w.finished, commit)
		} else {
			if len(w.finished) == 0 {
				return nil
			}
			commit := w.finished[w.rand.Intn(len(w.finished))]
			commit, err := pfsutil.StartCommit(pfsClient, commit.Repo.Name, commit.Id)
			if err != nil {
				return err
			}
			w.started = append(w.started, commit)
		}
	case opt < file:
		if len(w.started) == 0 {
			return nil
		}
		commit := w.started[w.rand.Intn(len(w.started))]
		if _, err := pfsutil.PutFile(pfsClient, commit.Repo.Name, commit.Id, w.randString(10), 0, w.reader()); err != nil {
			return err
		}
	case opt < job:
		if len(w.startedJobs) >= maxStartedJobs {
			job := w.startedJobs[0]
			w.startedJobs = w.startedJobs[1:]
			jobInfo, err := ppsClient.InspectJob(
				context.Background(),
				&pps.InspectJobRequest{
					Job:        job,
					BlockState: true,
				},
			)
			if err != nil {
				return err
			}
			if jobInfo.State != pps.JobState_JOB_STATE_SUCCESS {
				return fmt.Errorf("job %s failed", job.Id)
			}
			w.jobs = append(w.jobs, job)
		} else {
			if len(w.finished) == 0 {
				return nil
			}
			inputs := [5]string{}
			var inputCommits []*pfs.Commit
			repoSet := make(map[string]bool)
			for i := range inputs {
				commit := w.finished[w.rand.Intn(len(w.finished))]
				if _, ok := repoSet[commit.Repo.Name]; ok {
					continue
				}
				repoSet[commit.Repo.Name] = true
				inputs[i] = commit.Repo.Name
				inputCommits = append(inputCommits, commit)
			}
			var parentJobID string
			if len(w.jobs) > 0 {
				parentJobID = w.jobs[w.rand.Intn(len(w.jobs))].Id
			}
			outFilename := w.randString(10)
			job, err := ppsutil.CreateJob(
				ppsClient,
				"",
				[]string{"bash"},
				w.grepCmd(inputs, outFilename),
				1,
				inputCommits,
				parentJobID,
			)
			if err != nil {
				return err
			}
			w.startedJobs = append(w.startedJobs, job)
		}
	case opt < pipeline:
		if len(w.repos) == 0 {
			return nil
		}
		inputs := [5]string{}
		var inputRepos []*pfs.Repo
		repoSet := make(map[string]bool)
		for i := range inputs {
			repo := w.repos[w.rand.Intn(len(w.repos))]
			if _, ok := repoSet[repo.Name]; ok {
				continue
			}
			repoSet[repo.Name] = true
			inputs[i] = repo.Name
			inputRepos = append(inputRepos, repo)
		}
		pipelineName := w.randString(10)
		outFilename := w.randString(10)
		if err := ppsutil.CreatePipeline(
			ppsClient,
			pipelineName,
			"",
			[]string{"bash"},
			w.grepCmd(inputs, outFilename),
			1,
			inputRepos,
		); err != nil {
			return err
		}
		w.pipelines = append(w.pipelines, ppsutil.NewPipeline(pipelineName))
	}
	return nil
}

const letters = "abcdefghijklmnopqrstuvwxyz"
const lettersAndSpaces = "abcdefghijklmnopqrstuvwxyz      "

func (w *worker) randString(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = letters[w.rand.Intn(len(letters))]
	}
	return string(b)
}

type reader struct {
	rand  *rand.Rand
	lines int
}

func NewReader(rand *rand.Rand, lines int) io.Reader {
	return &reader{
		rand:  rand,
		lines: lines,
	}
}

func (r *reader) Read(p []byte) (int, error) {
	for i := range p {
		if i%128 == 127 {
			p[i] = '\n'
			r.lines--
		} else {
			p[i] = lettersAndSpaces[r.rand.Intn(len(lettersAndSpaces))]
		}
	}
	p[len(p)-1] = '\n'
	if r.lines <= 0 {
		return len(p), io.EOF
	}
	return len(p), nil
}

func (w *worker) reader() io.Reader {
	return NewReader(w.rand, 1000)
}

func (w *worker) grepCmd(inputs [5]string, outFilename string) string {
	return fmt.Sprintf(
		"grep %s /pfs/{%s,%s,%s,%s,%s}/* >/pfs/out/%s; true",
		w.randString(4),
		inputs[0],
		inputs[1],
		inputs[2],
		inputs[3],
		inputs[4],
		outFilename,
	)
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"os"
	"strings"
	"text/tabwriter"

	gflag "github.com/jessevdk/go-flags"
)

func (cli *HyperClient) HyperCmdList(args ...string) error {
	var opts struct {
		Pod string `short:"p" long:"pod" value-name:"\"\"" description:"only list the specified pod"`
		VM  string `short:"m" long:"vm" value-name:"\"\"" description:"only list resources on the specified vm"`
	}

	var parser = gflag.NewParser(&opts, gflag.Default|gflag.IgnoreUnknown)
	parser.Usage = "list [OPTIONS] [pod|container|vm]\n\nlist all pods or container information"
	args, err := parser.ParseArgs(args)
	if err != nil {
		if !strings.Contains(err.Error(), "Usage") {
			return err
		} else {
			return nil
		}
	}
	var item string
	if len(args) == 0 {
		item = "pod"
	} else {
		item = args[0]
	}

	if item != "pod" && item != "vm" && item != "container" {
		return fmt.Errorf("Error, the %s can not support %s list!", os.Args[0], item)
	}

	remoteInfo, err := cli.client.List(item, opts.Pod, opts.VM)
	if err != nil {
		return err
	}

	var (
		vmResponse        = []string{}
		podResponse       = []string{}
		containerResponse = []string{}
	)
	if remoteInfo.Exists("item") {
		item = remoteInfo.Get("item")
	}
	if remoteInfo.Exists("Error") {
		return fmt.Errorf("Found an error while getting %s list: %s", item, remoteInfo.Get("Error"))
	}

	w := tabwriter.NewWriter(cli.out, 20, 1, 3, ' ', 0)
	if item == "vm" {
		vmResponse = remoteInfo.GetList("vmData")
	}
	if item == "pod" {
		podResponse = remoteInfo.GetList("podData")
	}
	if item == "container" {
		containerResponse = remoteInfo.GetList("cData")
	}

	//fmt.Printf("Item is %s\n", item)
	if item == "vm" {
		fmt.Fprintln(w, "VM name\tStatus")
		for _, vm := range vmResponse {
			fields := strings.Split(vm, ":")
			fmt.Fprintf(w, "%s\t%s\n", fields[0], fields[2])
		}
	}

	if item == "pod" {
		fmt.Fprintln(w, "POD ID\tPOD Name\tVM name\tStatus")
		for _, p := range podResponse {
			fields := strings.Split(p, ":")
			fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", fields[0], fields[1], fields[2], fields[3])
		}
	}

	if item == "container" {
		fmt.Fprintln(w, "Container ID\tName\tPOD ID\tStatus")
		for _, c := range containerResponse {
			fields := strings.Split(c, ":")
			name := fields[1]
			if len(name) > 0 {
				if name[0] == '/' {
					name = name[1:]
				}
			}
			fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", fields[0], name, fields[2], fields[3])
		}
	}
	w.Flush()
	return nil
}
<commit_msg>[hyperctl] Add support of quiet param<commit_after>package client

import (
	"fmt"
	"os"
	"strings"
	"text/tabwriter"

	gflag "github.com/jessevdk/go-flags"
)

func (cli *HyperClient) HyperCmdList(args ...string) error {
	var opts struct {
		Pod   string `short:"p" long:"pod" value-name:"\"\"" description:"only list the specified pod"`
		VM    string `short:"m" long:"vm" value-name:"\"\"" description:"only list resources on the specified vm"`
		Quiet bool   `short:"q" long:"quiet" value-name:"\"\"" description:"Quiet mode"`
	}

	var parser = gflag.NewParser(&opts, gflag.Default|gflag.IgnoreUnknown)
	parser.Usage = "list [OPTIONS] [pod|container|vm]\n\nlist all pods or container information"
	args, err := parser.ParseArgs(args)
	if err != nil {
		if !strings.Contains(err.Error(), "Usage") {
			return err
		} else {
			return nil
		}
	}
	var item string
	if len(args) == 0 {
		item = "pod"
	} else {
		item = args[0]
	}

	if item != "pod" && item != "vm" && item != "container" {
		return fmt.Errorf("Error, the %s can not support %s list!", os.Args[0], item)
	}

	remoteInfo, err := cli.client.List(item, opts.Pod, opts.VM)
	if err != nil {
		return err
	}

	var (
		vmResponse        = []string{}
		podResponse       = []string{}
		containerResponse = []string{}
	)
	if remoteInfo.Exists("item") {
		item = remoteInfo.Get("item")
	}
	if remoteInfo.Exists("Error") {
		return fmt.Errorf("Found an error while getting %s list: %s", item, remoteInfo.Get("Error"))
	}

	w := tabwriter.NewWriter(cli.out, 20, 1, 3, ' ', 0)
	if item == "vm" {
		vmResponse = remoteInfo.GetList("vmData")
	}
	if item == "pod" {
		podResponse = remoteInfo.GetList("podData")
	}
	if item == "container" {
		containerResponse = remoteInfo.GetList("cData")
	}

	//fmt.Printf("Item is %s\n", item)
	if item == "vm" {
		if !opts.Quiet {
			fmt.Fprintln(w, "VM name\tStatus")
		}

		for _, vm := range vmResponse {
			fields := strings.Split(vm, ":")
			if opts.Quiet {
				fmt.Fprintf(w, "%s\n", fields[0])
			} else {
				fmt.Fprintf(w, "%s\t%s\n", fields[0], fields[2])
			}
		}
	}

	if item == "pod" {
		if !opts.Quiet {
			fmt.Fprintln(w, "POD ID\tPOD Name\tVM name\tStatus")
		}

		for _, p := range podResponse {
			fields := strings.Split(p, ":")
			if opts.Quiet {
				fmt.Fprintf(w, "%s\n", fields[0])
			} else {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", fields[0], fields[1], fields[2], fields[3])
			}
		}
	}

	if item == "container" {
		if !opts.Quiet {
			fmt.Fprintln(w, "Container ID\tName\tPOD ID\tStatus")
		}

		for _, c := range containerResponse {
			fields := strings.Split(c, ":")
			if opts.Quiet {
				fmt.Fprintf(w, "%s\n", fields[0])
			} else {
				name := fields[1]
				if len(name) > 0 {
					if name[0] == '/' {
						name = name[1:]
					}
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", fields[0], name, fields[2], fields[3])
			}
		}
	}
	w.Flush()
	return nil
}
<|endoftext|>


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

// A binary to produce LaTeX documents representing Merkle trees.
// The generated document should be fed into xelatex, and the Forest package
// must be available.
//
// Usage: go run main.go | xelatex
// This should generate a PDF file called treetek.pdf containing a drawing of
// the tree.
//
package main

import (
	"flag"
	"fmt"
	"log"
	"math/bits"
	"strings"

	"github.com/google/trillian/merkle"
	"github.com/google/trillian/merkle/compact"
)

const (
	preamble = `
% Hash-tree
% Author: treetex
\documentclass[convert]{standalone}
\usepackage[dvipsnames]{xcolor}
\usepackage{forest}

\begin{document}

% Change colours here:
\definecolor{inclusion}{rgb}{1,0.5,0.5}
\definecolor{inclusion_ephemeral}{rgb}{1,0.7,0.7}
\definecolor{perfect}{rgb}{1,0.9,0.5}
\definecolor{target}{rgb}{0.5,0.5,0.9}
\definecolor{target_path}{rgb}{0.7,0.7,0.9}
\definecolor{mega}{rgb}{0.9,0.9,0.9}
\definecolor{range0}{rgb}{0.3,0.9,0.3}
\definecolor{range1}{rgb}{0.3,0.3,0.9}
\definecolor{range2}{rgb}{0.9,0.3,0.9}

\forestset{
	% This defines a new "edge" style for drawing the perfect subtrees.
	% Rather than simply drawing a line representing an edge, this draws a
	% triangle between the labelled anchors on the given nodes.
	% See "Anchors" section in the Forest manual for more details:
	%  http://mirrors.ibiblio.org/CTAN/graphics/pgf/contrib/forest/forest-doc.pdf
	perfect/.style={edge path={%
		\noexpand\path[fill=mega, \forestoption{edge}]
				(.parent first)--(!u.children)--(.parent last)--cycle
				\forestoption{edge label};
		}
	},
}
\begin{forest}
`

	postfix = `\end{forest}
\end{document}
`

	// Maximum number of ranges to allow.
	maxRanges = 3
)

var (
	treeSize  = flag.Uint64("tree_size", 23, "Size of tree to produce")
	inclusion = flag.Int64("inclusion", -1, "Leaf index to show inclusion proof")
	megaMode  = flag.Uint("megamode_threshold", 4, "Treat perfect trees larger than this many layers as a single entity")
	ranges    = flag.String("ranges", "", "Comma-separated Open-Closed ranges of the form L:R")

	// nInfo holds nodeInfo data for the tree.
	nInfo = make(map[compact.NodeID]nodeInfo)
)

// nodeInfo represents the style to be applied to a tree node.
type nodeInfo struct {
	incProof     bool
	incPath      bool
	target       bool
	perfectRoot  bool
	ephemeral    bool
	leaf         bool
	rangeIndices []int
}

// String returns a string containing Forest attributes suitable for
// rendering the node, given its type.
func (n nodeInfo) String() string {
	attr := make([]string, 0, 4)

	// Figure out which colour to fill with:
	fill := "white"
	if n.perfectRoot {
		attr = append(attr, "line width=4pt")
	}
	if n.incProof {
		fill = "inclusion"
		if n.ephemeral {
			fill = "inclusion_ephemeral"
			attr = append(attr, "draw, dotted")
		}
	}
	if n.target {
		fill = "target"
	}
	if n.incPath {
		fill = "target_path"
	}

	if len(n.rangeIndices) == 1 {
		// For nodes in a single range just change the fill.
		fill = fmt.Sprintf("range%d!50", n.rangeIndices[0])
	} else {
		// Otherwise, we need to be a bit cleverer, and use the shading feature.
		for i, ri := range n.rangeIndices {
			pos := []string{"left", "right", "middle"}[i]
			attr = append(attr, fmt.Sprintf("%s color=range%d!50", pos, ri))
		}
	}
	attr = append(attr, "fill="+fill)

	if !n.ephemeral {
		attr = append(attr, "draw")
	}
	if !n.leaf {
		attr = append(attr, "circle, minimum size=3em")
	} else {
		attr = append(attr, "minimum size=1.5em")
	}
	return strings.Join(attr, ", ")
}

// modifyNodeInfo applies f to the nodeInfo associated with node id.
func modifyNodeInfo(id compact.NodeID, f func(*nodeInfo)) {
	n := nInfo[id] // Note: Returns an empty nodeInfo if id is not found.
	f(&n)
	nInfo[id] = n
}

// perfectMega renders a large perfect subtree as a single entity.
func perfectMega(prefix string, height uint, leafIndex uint64) {
	stLeaves := uint64(1) << height
	stWidth := float32(stLeaves) / float32(*treeSize)
	fmt.Printf("%s [%d\\dots%d, edge label={node[midway, above]{%d}}, perfect, tier=leaf, minimum width=%f\\linewidth ]\n", prefix, leafIndex, leafIndex+stLeaves, stLeaves, stWidth)

	// Create some hidden nodes to preseve the tier spacings:
	fmt.Printf("%s", prefix)
	for i := int(height - 2); i > 0; i-- {
		fmt.Printf(" [, no edge, tier=%d ", i)
		defer fmt.Printf(" ] ")
	}
}

// perfect renders a perfect subtree.
func perfect(prefix string, height uint, index uint64) {
	perfectInner(prefix, height, index, true)
}

// drawLeaf emits TeX code to render a leaf.
func drawLeaf(prefix string, index uint64) {
	a := nInfo[compact.NewNodeID(0, index)]
	fmt.Printf("%s [%d, %s, tier=leaf]\n", prefix, index, a.String())
}

// openInnerNode renders TeX code to open an internal node.
// The caller may emit any number of child nodes before calling the returned
// func to close the node.
// Returns a func to be called to close the node.
func openInnerNode(prefix string, id compact.NodeID) func() {
	attr := nInfo[id].String()
	fmt.Printf("%s [%d.%d, %s, tier=%d\n", prefix, id.Level, id.Index, attr, id.Index-1)
	return func() { fmt.Printf("%s ]\n", prefix) }
}

// perfectInner renders the nodes of a perfect internal subtree.
func perfectInner(prefix string, level uint, index uint64, top bool) {
	id := compact.NewNodeID(level, index)
	modifyNodeInfo(id, func(n *nodeInfo) {
		n.leaf = id.Level == 0
		n.perfectRoot = top
	})

	if level == 0 {
		drawLeaf(prefix, index)
		return
	}
	c := openInnerNode(prefix, id)
	childIndex := index << 1
	if level > *megaMode {
		perfectMega(prefix, level, index<<level)
	} else {
		perfectInner(prefix+" ", level-1, childIndex, false)
		perfectInner(prefix+" ", level-1, childIndex+1, false)
	}
	c()
}

// renderTree renders a tree node and recurses if necessary.
func renderTree(prefix string, treeSize, index uint64) {
	// Look at the bit of the treeSize corresponding to the current level:
	height := uint(bits.Len64(treeSize) - 1)
	b := uint64(1) << height
	rest := treeSize - b
	// left child is a perfect subtree.

	// if there's a right-hand child, then we'll emit this node to be the
	// parent. (Otherwise we'll just keep quiet, and recurse down - this is how
	// we arrange for leaves to always be on the bottom level.)
	if rest > 0 {
		childHeight := height + 1
		id := compact.NewNodeID(childHeight, index>>childHeight)
		modifyNodeInfo(id, func(n *nodeInfo) { n.ephemeral = true })
		c := openInnerNode(prefix, id)
		defer c()
	}
	perfect(prefix+" ", height, index>>height)
	index += b
	renderTree(prefix+" ", rest, index)
}

// parseRanges parses and validates a string of comma-separates open-closed
// ranges of the form L:R.
// Returns the parsed ranges, or an error if there's a problem.
func parseRanges(ranges string, treeSize uint64) ([][2]uint64, error) {
	rangePairs := strings.Split(ranges, ",")
	numRanges := len(rangePairs)
	if num, max := numRanges, maxRanges; num > max {
		return nil, fmt.Errorf("too many ranges %d, must be %d or fewer", num, max)
	}
	ret := make([][2]uint64, 0, numRanges)
	for _, rng := range rangePairs {
		lr := strings.Split(rng, ":")
		if len(lr) != 2 {
			return nil, fmt.Errorf("specified range %q is invalid", rng)
		}
		var l, r uint64
		if _, err := fmt.Sscanf(rng, "%d:%d", &l, &r); err != nil {
			return nil, fmt.Errorf("range %q is malformed: %s", rng, err)
		}
		switch {
		case r > treeSize:
			return nil, fmt.Errorf("range %q extends past end of tree (%d)", lr, treeSize)
		case l > r:
			return nil, fmt.Errorf("range elements in %q are out of order", rng)
		}
		ret = append(ret, [2]uint64{l, r})
	}
	return ret, nil
}

// modifyRangeNodeInfo sets style info for nodes affected by ranges.
// This includes leaves and perfect subtree roots.
// TODO(al): Figure out what, if anything, to do to make this show ranges
// which are inside the perfect meganodes.
func modifyRangeNodeInfo() error {
	rng, err := parseRanges(*ranges, *treeSize)
	if err != nil {
		return err
	}
	for ri, lr := range rng {
		l, r := lr[0], lr[1]
		// Set leaves:
		for i := l; i < r; i++ {
			id := compact.NewNodeID(0, i)
			modifyNodeInfo(id, func(n *nodeInfo) { n.rangeIndices = append(n.rangeIndices, ri) })
		}

		// Now perfect roots which comprise the range:
		maskL, maskR := compact.Decompose(l, r)
		p := l
		// Do left perfect subtree roots:
		nBitsL := uint(bits.Len64(maskL))
		for i, bit := uint(0), uint64(1); i < nBitsL; i, bit = i+1, bit<<1 {
			if maskL&bit != 0 {
				if i > 0 {
					id := compact.NewNodeID(i, p>>i)
					modifyNodeInfo(id, func(n *nodeInfo) { n.rangeIndices = append(n.rangeIndices, ri) })
				}
				p += bit
			}
		}
		// Do right perfect subtree roots:
		nBitsR := uint(bits.Len64(maskR))
		for i, bit := nBitsR, uint64(1<<nBitsR); i > 1; i, bit = i-1, bit>>1 {
			if maskR&bit != 0 {
				id := compact.NewNodeID(i, p>>i)
				modifyNodeInfo(id, func(n *nodeInfo) { n.rangeIndices = append(n.rangeIndices, ri) })
				p += bit
			}
		}
	}
	return nil
}

// Whee - here we go!
func main() {
	// TODO(al): check flag validity.
	flag.Parse()
	height := uint(bits.Len64(*treeSize-1)) + 1

	if *inclusion > 0 {
		leafID := compact.NewNodeID(0, uint64(*inclusion))
		modifyNodeInfo(leafID, func(n *nodeInfo) { n.target = true })
		nf, err := merkle.CalcInclusionProofNodeAddresses(int64(*treeSize), *inclusion, int64(*treeSize))
		if err != nil {
			log.Fatalf("Failed to calculate inclusion proof addresses: %s", err)
		}
		for _, n := range nf {
			modifyNodeInfo(n.ID, func(n *nodeInfo) { n.incProof = true })
		}
		for h, i := uint(0), leafID.Index; h < height; h, i = h+1, i>>1 {
			id := compact.NewNodeID(h, i)
			modifyNodeInfo(id, func(n *nodeInfo) { n.incPath = true })
		}
	}

	if len(*ranges) > 0 {
		if err := modifyRangeNodeInfo(); err != nil {
			log.Fatalf("Failed to modify range node styles: %s", err)
		}
	}

	// TODO(al): structify this into a util, and add ability to output to an
	// arbitrary stream.
	fmt.Print(preamble)
	renderTree("", *treeSize, 0)
	fmt.Print(postfix)
}
<commit_msg>[treetex] Fix treetex (#1982)<commit_after>// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// A binary to produce LaTeX documents representing Merkle trees.
// The generated document should be fed into xelatex, and the Forest package
// must be available.
//
// Usage: go run main.go | xelatex
// This should generate a PDF file called treetek.pdf containing a drawing of
// the tree.
//
package main

import (
	"flag"
	"fmt"
	"log"
	"math/bits"
	"strings"

	"github.com/google/trillian/merkle"
	"github.com/google/trillian/merkle/compact"
)

const (
	preamble = `
% Hash-tree
% Author: treetex
\documentclass[convert]{standalone}
\usepackage[dvipsnames]{xcolor}
\usepackage{forest}

\begin{document}

% Change colours here:
\definecolor{inclusion}{rgb}{1,0.5,0.5}
\definecolor{inclusion_ephemeral}{rgb}{1,0.7,0.7}
\definecolor{perfect}{rgb}{1,0.9,0.5}
\definecolor{target}{rgb}{0.5,0.5,0.9}
\definecolor{target_path}{rgb}{0.7,0.7,0.9}
\definecolor{mega}{rgb}{0.9,0.9,0.9}
\definecolor{range0}{rgb}{0.3,0.9,0.3}
\definecolor{range1}{rgb}{0.3,0.3,0.9}
\definecolor{range2}{rgb}{0.9,0.3,0.9}

\forestset{
	% This defines a new "edge" style for drawing the perfect subtrees.
	% Rather than simply drawing a line representing an edge, this draws a
	% triangle between the labelled anchors on the given nodes.
	% See "Anchors" section in the Forest manual for more details:
	%  http://mirrors.ibiblio.org/CTAN/graphics/pgf/contrib/forest/forest-doc.pdf
	perfect/.style={edge path={%
		\noexpand\path[fill=mega, \forestoption{edge}]
				(.parent first)--(!u.children)--(.parent last)--cycle
				\forestoption{edge label};
		}
	},
}
\begin{forest}
`

	postfix = `\end{forest}
\end{document}
`

	// Maximum number of ranges to allow.
	maxRanges = 3
)

var (
	treeSize  = flag.Uint64("tree_size", 23, "Size of tree to produce")
	inclusion = flag.Int64("inclusion", -1, "Leaf index to show inclusion proof")
	megaMode  = flag.Uint("megamode_threshold", 4, "Treat perfect trees larger than this many layers as a single entity")
	ranges    = flag.String("ranges", "", "Comma-separated Open-Closed ranges of the form L:R")

	// nInfo holds nodeInfo data for the tree.
	nInfo = make(map[compact.NodeID]nodeInfo)
)

// nodeInfo represents the style to be applied to a tree node.
type nodeInfo struct {
	incProof     bool
	incPath      bool
	target       bool
	perfectRoot  bool
	ephemeral    bool
	leaf         bool
	rangeIndices []int
}

// String returns a string containing Forest attributes suitable for
// rendering the node, given its type.
func (n nodeInfo) String() string {
	attr := make([]string, 0, 4)

	// Figure out which colour to fill with:
	fill := "white"
	if n.perfectRoot {
		attr = append(attr, "line width=4pt")
	}
	if n.incProof {
		fill = "inclusion"
		if n.ephemeral {
			fill = "inclusion_ephemeral"
			attr = append(attr, "draw, dotted")
		}
	}
	if n.target {
		fill = "target"
	}
	if n.incPath {
		fill = "target_path"
	}

	if len(n.rangeIndices) == 1 {
		// For nodes in a single range just change the fill.
		fill = fmt.Sprintf("range%d!50", n.rangeIndices[0])
	} else {
		// Otherwise, we need to be a bit cleverer, and use the shading feature.
		for i, ri := range n.rangeIndices {
			pos := []string{"left", "right", "middle"}[i]
			attr = append(attr, fmt.Sprintf("%s color=range%d!50", pos, ri))
		}
	}
	attr = append(attr, "fill="+fill)

	if !n.ephemeral {
		attr = append(attr, "draw")
	}
	if !n.leaf {
		attr = append(attr, "circle, minimum size=3em")
	} else {
		attr = append(attr, "minimum size=1.5em")
	}
	return strings.Join(attr, ", ")
}

// modifyNodeInfo applies f to the nodeInfo associated with node id.
func modifyNodeInfo(id compact.NodeID, f func(*nodeInfo)) {
	n := nInfo[id] // Note: Returns an empty nodeInfo if id is not found.
	f(&n)
	nInfo[id] = n
}

// perfectMega renders a large perfect subtree as a single entity.
func perfectMega(prefix string, height uint, leafIndex uint64) {
	stLeaves := uint64(1) << height
	stWidth := float32(stLeaves) / float32(*treeSize)
	fmt.Printf("%s [%d\\dots%d, edge label={node[midway, above]{%d}}, perfect, tier=leaf, minimum width=%f\\linewidth ]\n", prefix, leafIndex, leafIndex+stLeaves, stLeaves, stWidth)

	// Create some hidden nodes to preseve the tier spacings:
	fmt.Printf("%s", prefix)
	for i := int(height - 2); i > 0; i-- {
		fmt.Printf(" [, no edge, tier=%d ", i)
		defer fmt.Printf(" ] ")
	}
}

// perfect renders a perfect subtree.
func perfect(prefix string, height uint, index uint64) {
	perfectInner(prefix, height, index, true)
}

// drawLeaf emits TeX code to render a leaf.
func drawLeaf(prefix string, index uint64) {
	a := nInfo[compact.NewNodeID(0, index)]
	fmt.Printf("%s [%d, %s, tier=leaf]\n", prefix, index, a.String())
}

// openInnerNode renders TeX code to open an internal node.
// The caller may emit any number of child nodes before calling the returned
// func to close the node.
// Returns a func to be called to close the node.
func openInnerNode(prefix string, id compact.NodeID) func() {
	attr := nInfo[id].String()
	fmt.Printf("%s [%d.%d, %s, tier=%[2]d\n", prefix, id.Level, id.Index, attr)
	return func() { fmt.Printf("%s ]\n", prefix) }
}

// perfectInner renders the nodes of a perfect internal subtree.
func perfectInner(prefix string, level uint, index uint64, top bool) {
	id := compact.NewNodeID(level, index)
	modifyNodeInfo(id, func(n *nodeInfo) {
		n.leaf = id.Level == 0
		n.perfectRoot = top
	})

	if level == 0 {
		drawLeaf(prefix, index)
		return
	}
	c := openInnerNode(prefix, id)
	childIndex := index << 1
	if level > *megaMode {
		perfectMega(prefix, level, index<<level)
	} else {
		perfectInner(prefix+" ", level-1, childIndex, false)
		perfectInner(prefix+" ", level-1, childIndex+1, false)
	}
	c()
}

// renderTree renders a tree node and recurses if necessary.
func renderTree(prefix string, treeSize, index uint64) {
	if treeSize == 0 {
		return
	}
	// Look at the bit of the treeSize corresponding to the current level:
	height := uint(bits.Len64(treeSize) - 1)
	b := uint64(1) << height
	rest := treeSize - b
	// left child is a perfect subtree.

	// if there's a right-hand child, then we'll emit this node to be the
	// parent. (Otherwise we'll just keep quiet, and recurse down - this is how
	// we arrange for leaves to always be on the bottom level.)
	if rest > 0 {
		childHeight := height + 1
		id := compact.NewNodeID(childHeight, index>>childHeight)
		modifyNodeInfo(id, func(n *nodeInfo) { n.ephemeral = true })
		c := openInnerNode(prefix, id)
		defer c()
	}
	perfect(prefix+" ", height, index>>height)
	index += b
	renderTree(prefix+" ", rest, index)
}

// parseRanges parses and validates a string of comma-separates open-closed
// ranges of the form L:R.
// Returns the parsed ranges, or an error if there's a problem.
func parseRanges(ranges string, treeSize uint64) ([][2]uint64, error) {
	rangePairs := strings.Split(ranges, ",")
	numRanges := len(rangePairs)
	if num, max := numRanges, maxRanges; num > max {
		return nil, fmt.Errorf("too many ranges %d, must be %d or fewer", num, max)
	}
	ret := make([][2]uint64, 0, numRanges)
	for _, rng := range rangePairs {
		lr := strings.Split(rng, ":")
		if len(lr) != 2 {
			return nil, fmt.Errorf("specified range %q is invalid", rng)
		}
		var l, r uint64
		if _, err := fmt.Sscanf(rng, "%d:%d", &l, &r); err != nil {
			return nil, fmt.Errorf("range %q is malformed: %s", rng, err)
		}
		switch {
		case r > treeSize:
			return nil, fmt.Errorf("range %q extends past end of tree (%d)", lr, treeSize)
		case l > r:
			return nil, fmt.Errorf("range elements in %q are out of order", rng)
		}
		ret = append(ret, [2]uint64{l, r})
	}
	return ret, nil
}

// modifyRangeNodeInfo sets style info for nodes affected by ranges.
// This includes leaves and perfect subtree roots.
// TODO(al): Figure out what, if anything, to do to make this show ranges
// which are inside the perfect meganodes.
func modifyRangeNodeInfo() error {
	rng, err := parseRanges(*ranges, *treeSize)
	if err != nil {
		return err
	}
	for ri, lr := range rng {
		l, r := lr[0], lr[1]
		// Set leaves:
		for i := l; i < r; i++ {
			id := compact.NewNodeID(0, i)
			modifyNodeInfo(id, func(n *nodeInfo) { n.rangeIndices = append(n.rangeIndices, ri) })
		}

		// Now perfect roots which comprise the range:
		maskL, maskR := compact.Decompose(l, r)
		p := l
		// Do left perfect subtree roots:
		nBitsL := uint(bits.Len64(maskL))
		for i, bit := uint(0), uint64(1); i < nBitsL; i, bit = i+1, bit<<1 {
			if maskL&bit != 0 {
				if i > 0 {
					id := compact.NewNodeID(i, p>>i)
					modifyNodeInfo(id, func(n *nodeInfo) { n.rangeIndices = append(n.rangeIndices, ri) })
				}
				p += bit
			}
		}
		// Do right perfect subtree roots:
		nBitsR := uint(bits.Len64(maskR))
		for i, bit := nBitsR, uint64(1<<nBitsR); i > 1; i, bit = i-1, bit>>1 {
			if maskR&bit != 0 {
				id := compact.NewNodeID(i, p>>i)
				modifyNodeInfo(id, func(n *nodeInfo) { n.rangeIndices = append(n.rangeIndices, ri) })
				p += bit
			}
		}
	}
	return nil
}

// Whee - here we go!
func main() {
	// TODO(al): check flag validity.
	flag.Parse()
	height := uint(bits.Len64(*treeSize-1)) + 1

	if *inclusion > 0 {
		leafID := compact.NewNodeID(0, uint64(*inclusion))
		modifyNodeInfo(leafID, func(n *nodeInfo) { n.target = true })
		nf, err := merkle.CalcInclusionProofNodeAddresses(int64(*treeSize), *inclusion, int64(*treeSize))
		if err != nil {
			log.Fatalf("Failed to calculate inclusion proof addresses: %s", err)
		}
		for _, n := range nf {
			modifyNodeInfo(n.ID, func(n *nodeInfo) { n.incProof = true })
		}
		for h, i := uint(0), leafID.Index; h < height; h, i = h+1, i>>1 {
			id := compact.NewNodeID(h, i)
			modifyNodeInfo(id, func(n *nodeInfo) { n.incPath = true })
		}
	}

	if len(*ranges) > 0 {
		if err := modifyRangeNodeInfo(); err != nil {
			log.Fatalf("Failed to modify range node styles: %s", err)
		}
	}

	// TODO(al): structify this into a util, and add ability to output to an
	// arbitrary stream.
	fmt.Print(preamble)
	renderTree("", *treeSize, 0)
	fmt.Print(postfix)
}
<|endoftext|>