<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package server

import (
	"fmt"
	"io"

	"github.com/spf13/cobra"

	genericapiserver "k8s.io/apiserver/pkg/server"
	genericoptions "k8s.io/apiserver/pkg/server/options"
	"k8s.io/sample-apiserver/pkg/apis/wardle/v1alpha1"
	"k8s.io/sample-apiserver/pkg/apiserver"
)

const defaultEtcdPathPrefix = "/registry/wardle.kubernetes.io"

type WardleServerOptions struct {
	RecommendedOptions *genericoptions.RecommendedOptions

	StdOut io.Writer
	StdErr io.Writer
}

func NewWardleServerOptions(out, errOut io.Writer) *WardleServerOptions {
	o := &WardleServerOptions{
		RecommendedOptions: genericoptions.NewRecommendedOptions(defaultEtcdPathPrefix, apiserver.Scheme, apiserver.Codecs.LegacyCodec(v1alpha1.SchemeGroupVersion)),

		StdOut: out,
		StdErr: errOut,
	}
	o.RecommendedOptions.SecureServing.ServingOptions.BindPort = 443

	return o
}

// NewCommandStartMaster provides a CLI handler for 'start master' command
func NewCommandStartWardleServer(out, errOut io.Writer, stopCh <-chan struct{}) *cobra.Command {
	o := NewWardleServerOptions(out, errOut)

	cmd := &cobra.Command{
		Short: "Launch a wardle API server",
		Long:  "Launch a wardle API server",
		RunE: func(c *cobra.Command, args []string) error {
			if err := o.Complete(); err != nil {
				return err
			}
			if err := o.Validate(args); err != nil {
				return err
			}
			if err := o.RunWardleServer(stopCh); err != nil {
				return err
			}
			return nil
		},
	}

	flags := cmd.Flags()
	o.RecommendedOptions.AddFlags(flags)

	return cmd
}

func (o WardleServerOptions) Validate(args []string) error {
	return nil
}

func (o *WardleServerOptions) Complete() error {
	return nil
}

func (o WardleServerOptions) Config() (*apiserver.Config, error) {
	// TODO have a "real" external address
	if err := o.RecommendedOptions.SecureServing.MaybeDefaultWithSelfSignedCerts("localhost"); err != nil {
		return nil, fmt.Errorf("error creating self-signed certificates: %v", err)
	}

	serverConfig := genericapiserver.NewConfig().WithSerializer(apiserver.Codecs)
	if err := o.RecommendedOptions.ApplyTo(serverConfig); err != nil {
		return nil, err
	}

	config := &apiserver.Config{
		GenericConfig: serverConfig,
	}
	return config, nil
}

func (o WardleServerOptions) RunWardleServer(stopCh <-chan struct{}) error {
	config, err := o.Config()
	if err != nil {
		return err
	}

	server, err := config.Complete().New()
	if err != nil {
		return err
	}
	server.GenericAPIServer.PrepareRun().Run(stopCh)

	return nil
}
<commit_msg>tweak defaults for recommended apiserver options<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package server

import (
	"fmt"
	"io"

	"github.com/spf13/cobra"

	genericapiserver "k8s.io/apiserver/pkg/server"
	genericoptions "k8s.io/apiserver/pkg/server/options"
	"k8s.io/sample-apiserver/pkg/apis/wardle/v1alpha1"
	"k8s.io/sample-apiserver/pkg/apiserver"
)

const defaultEtcdPathPrefix = "/registry/wardle.kubernetes.io"

type WardleServerOptions struct {
	RecommendedOptions *genericoptions.RecommendedOptions

	StdOut io.Writer
	StdErr io.Writer
}

func NewWardleServerOptions(out, errOut io.Writer) *WardleServerOptions {
	o := &WardleServerOptions{
		RecommendedOptions: genericoptions.NewRecommendedOptions(defaultEtcdPathPrefix, apiserver.Scheme, apiserver.Codecs.LegacyCodec(v1alpha1.SchemeGroupVersion)),

		StdOut: out,
		StdErr: errOut,
	}

	return o
}

// NewCommandStartMaster provides a CLI handler for 'start master' command
func NewCommandStartWardleServer(out, errOut io.Writer, stopCh <-chan struct{}) *cobra.Command {
	o := NewWardleServerOptions(out, errOut)

	cmd := &cobra.Command{
		Short: "Launch a wardle API server",
		Long:  "Launch a wardle API server",
		RunE: func(c *cobra.Command, args []string) error {
			if err := o.Complete(); err != nil {
				return err
			}
			if err := o.Validate(args); err != nil {
				return err
			}
			if err := o.RunWardleServer(stopCh); err != nil {
				return err
			}
			return nil
		},
	}

	flags := cmd.Flags()
	o.RecommendedOptions.AddFlags(flags)

	return cmd
}

func (o WardleServerOptions) Validate(args []string) error {
	return nil
}

func (o *WardleServerOptions) Complete() error {
	return nil
}

func (o WardleServerOptions) Config() (*apiserver.Config, error) {
	// TODO have a "real" external address
	if err := o.RecommendedOptions.SecureServing.MaybeDefaultWithSelfSignedCerts("localhost"); err != nil {
		return nil, fmt.Errorf("error creating self-signed certificates: %v", err)
	}

	serverConfig := genericapiserver.NewConfig().WithSerializer(apiserver.Codecs)
	if err := o.RecommendedOptions.ApplyTo(serverConfig); err != nil {
		return nil, err
	}

	config := &apiserver.Config{
		GenericConfig: serverConfig,
	}
	return config, nil
}

func (o WardleServerOptions) RunWardleServer(stopCh <-chan struct{}) error {
	config, err := o.Config()
	if err != nil {
		return err
	}

	server, err := config.Complete().New()
	if err != nil {
		return err
	}
	server.GenericAPIServer.PrepareRun().Run(stopCh)

	return nil
}
<|endoftext|>


<commit_before>package msg_test

import (
	"testing"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/event"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/trade"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/mux/msg"
	"github.com/stretchr/testify/assert"
)

func TestIsEvent(t *testing.T) {
	cases := map[string]struct {
		pld      []byte
		expected bool
	}{
		"event type": {
			pld:      []byte(`{}`),
			expected: true,
		},
		"not event type": {
			pld:      []byte(`[]`),
			expected: false,
		},
	}

	for k, v := range cases {
		t.Run(k, func(t *testing.T) {
			m := msg.Msg{
				Data: v.pld,
			}

			got := m.IsEvent()
			assert.Equal(t, v.expected, got)
		})
	}
}

func TestIsRaw(t *testing.T) {
	cases := map[string]struct {
		pld      []byte
		expected bool
	}{
		"raw type": {
			pld:      []byte(`[]`),
			expected: true,
		},
		"not raw type": {
			pld:      []byte(`{}`),
			expected: false,
		},
	}

	for k, v := range cases {
		t.Run(k, func(t *testing.T) {
			m := msg.Msg{
				Data: v.pld,
			}

			got := m.IsRaw()
			assert.Equal(t, v.expected, got)
		})
	}
}

func TestProcessEvent(t *testing.T) {
	m := msg.Msg{
		Data: []byte(`{"event":"info","version":2,"serverId":"dbea77ee-4740-4a82-84f3-c6bc1b5abb9a","platform":{"status":1}}`),
	}

	expected := event.Info{
		Subscribe: event.Subscribe{
			Event: "info",
		},
		Version:  2,
		ServerID: "dbea77ee-4740-4a82-84f3-c6bc1b5abb9a",
		Platform: event.Platform{Status: 1},
	}

	got, err := m.ProcessEvent()
	assert.NoError(t, err)
	assert.Equal(t, expected, got)
}

func TestProcessRaw(t *testing.T) {
	cases := map[string]struct {
		pld      []byte
		expected interface{}
		inf      map[int64]event.Info
	}{
		"raw info event": {
			pld: []byte(`[123, "hb"]`),
			inf: map[int64]event.Info{123: {}},
			expected: event.Info{
				ChanID:    123,
				Subscribe: event.Subscribe{Event: "hb"},
			},
		},
		"raw trades snapshot": {
			pld: []byte(`[111,[[559273857,1609665708633,-0.0048,34113]]]`),
			inf: map[int64]event.Info{
				111: {
					Subscribe: event.Subscribe{
						Channel: "trades",
						Symbol:  "tBTCUST",
					},
				},
			},
			expected: &trade.Snapshot{
				Snapshot: []*trade.Trade{
					{
						Pair:   "tBTCUST",
						ID:     559273857,
						MTS:    1609665708633,
						Amount: -0.0048,
						Price:  34113,
					},
				},
			},
		},
		"raw trade": {
			pld: []byte(`[111,[559273857,1609665708633,-0.0048,34113]]`),
			inf: map[int64]event.Info{
				111: {
					Subscribe: event.Subscribe{
						Channel: "trades",
						Symbol:  "tBTCUST",
					},
				},
			},
			expected: &trade.Trade{
				Pair:   "tBTCUST",
				ID:     559273857,
				MTS:    1609665708633,
				Amount: -0.0048,
				Price:  34113,
			},
		},
	}

	for k, v := range cases {
		t.Run(k, func(t *testing.T) {
			m := msg.Msg{Data: v.pld}
			got, err := m.ProcessRaw(v.inf)
			assert.NoError(t, err)
			assert.Equal(t, v.expected, got)
		})
	}
}
<commit_msg>ticker, candle and candle snapshot "decode" test coverage<commit_after>package msg_test

import (
	"testing"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/candle"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/event"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/ticker"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/trade"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/mux/msg"
	"github.com/stretchr/testify/assert"
)

func TestIsEvent(t *testing.T) {
	cases := map[string]struct {
		pld      []byte
		expected bool
	}{
		"event type": {
			pld:      []byte(`{}`),
			expected: true,
		},
		"not event type": {
			pld:      []byte(`[]`),
			expected: false,
		},
	}

	for k, v := range cases {
		t.Run(k, func(t *testing.T) {
			m := msg.Msg{
				Data: v.pld,
			}

			got := m.IsEvent()
			assert.Equal(t, v.expected, got)
		})
	}
}

func TestIsRaw(t *testing.T) {
	cases := map[string]struct {
		pld      []byte
		expected bool
	}{
		"raw msg type": {
			pld:      []byte(`[]`),
			expected: true,
		},
		"raw info type": {
			pld:      []byte(`{}`),
			expected: false,
		},
	}

	for k, v := range cases {
		t.Run(k, func(t *testing.T) {
			m := msg.Msg{
				Data: v.pld,
			}

			got := m.IsRaw()
			assert.Equal(t, v.expected, got)
		})
	}
}

func TestProcessEvent(t *testing.T) {
	m := msg.Msg{
		Data: []byte(`{
			"event": "info",
			"version": 2,
			"serverId": "dbea77ee-4740-4a82-84f3-c6bc1b5abb9a",
			"platform": {
				"status":1
			}
		}`),
	}

	expected := event.Info{
		Subscribe: event.Subscribe{
			Event: "info",
		},
		Version:  2,
		ServerID: "dbea77ee-4740-4a82-84f3-c6bc1b5abb9a",
		Platform: event.Platform{Status: 1},
	}

	got, err := m.ProcessEvent()
	assert.NoError(t, err)
	assert.Equal(t, expected, got)
}

func TestProcessRaw(t *testing.T) {
	cases := map[string]struct {
		pld      []byte
		expected interface{}
		inf      map[int64]event.Info
	}{
		"info event": {
			pld: []byte(`[123, "hb"]`),
			inf: map[int64]event.Info{123: {}},
			expected: event.Info{
				ChanID:    123,
				Subscribe: event.Subscribe{Event: "hb"},
			},
		},
		"trades snapshot": {
			pld: []byte(`[111,[[559273857,1609665708633,-0.0048,34113]]]`),
			inf: map[int64]event.Info{
				111: {
					Subscribe: event.Subscribe{
						Channel: "trades",
						Symbol:  "tBTCUST",
					},
				},
			},
			expected: &trade.Snapshot{
				Snapshot: []*trade.Trade{
					{
						Pair:   "tBTCUST",
						ID:     559273857,
						MTS:    1609665708633,
						Amount: -0.0048,
						Price:  34113,
					},
				},
			},
		},
		"trade": {
			pld: []byte(`[111,[559273857,1609665708633,-0.0048,34113]]`),
			inf: map[int64]event.Info{
				111: {
					Subscribe: event.Subscribe{
						Channel: "trades",
						Symbol:  "tBTCUST",
					},
				},
			},
			expected: &trade.Trade{
				Pair:   "tBTCUST",
				ID:     559273857,
				MTS:    1609665708633,
				Amount: -0.0048,
				Price:  34113,
			},
		},
		"ticker": {
			pld: []byte(`[
				111,
				[
					34072,0.019999999999999997,34080,6.69793272,4350,
					0.1464,34062,4047.85335915,34758,29490
				]
			]`),
			inf: map[int64]event.Info{
				111: {
					Subscribe: event.Subscribe{
						Channel: "ticker",
						Symbol:  "tBTCUST",
					},
				},
			},
			expected: &ticker.Ticker{
				Symbol:          "tBTCUST",
				Bid:             34072,
				BidSize:         0.019999999999999997,
				Ask:             34080,
				AskSize:         6.69793272,
				DailyChange:     4350,
				DailyChangePerc: 0.1464,
				LastPrice:       34062,
				Volume:          4047.85335915,
				High:            34758,
				Low:             29490,
			},
		},
		"candles snapshot": {
			pld: []byte(`[111,[[1609668540000,828.01,827.67,828.42,827.67,2.32080241]]]`),
			inf: map[int64]event.Info{
				111: {
					Subscribe: event.Subscribe{
						Channel: "candles",
						Key:     "trade:1m:tETHUST",
					},
				},
			},
			expected: &candle.Snapshot{
				Snapshot: []*candle.Candle{
					{
						Symbol:     "tETHUST",
						Resolution: "1m",
						MTS:        1609668540000,
						Open:       828.01,
						Close:      827.67,
						High:       828.42,
						Low:        827.67,
						Volume:     2.32080241,
					},
				},
			},
		},
		"candle": {
			pld: []byte(`[111,[1609668540000,828.01,827.67,828.42,827.67,2.32080241]]`),
			inf: map[int64]event.Info{
				111: {
					Subscribe: event.Subscribe{
						Channel: "candles",
						Key:     "trade:1m:tETHUST",
					},
				},
			},
			expected: &candle.Candle{
				Symbol:     "tETHUST",
				Resolution: "1m",
				MTS:        1609668540000,
				Open:       828.01,
				Close:      827.67,
				High:       828.42,
				Low:        827.67,
				Volume:     2.32080241,
			},
		},
	}

	for k, v := range cases {
		t.Run(k, func(t *testing.T) {
			m := msg.Msg{Data: v.pld}
			got, err := m.ProcessRaw(v.inf)
			assert.NoError(t, err)
			assert.Equal(t, v.expected, got)
		})
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/apigateway"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSAPIGatewayRestApi_basic(t *testing.T) {
	var conf apigateway.RestApi

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAPIGatewayRestAPIDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSAPIGatewayRestAPIConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPINameAttribute(&conf, "bar"),
					testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttribute(&conf, 0),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "name", "bar"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "description", ""),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "minimum_compression_size", "0"),
					resource.TestCheckResourceAttrSet("aws_api_gateway_rest_api.test", "created_date"),
					resource.TestCheckNoResourceAttr("aws_api_gateway_rest_api.test", "binary_media_types"),
				),
			},

			{
				Config: testAccAWSAPIGatewayRestAPIUpdateConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPINameAttribute(&conf, "test"),
					testAccCheckAWSAPIGatewayRestAPIDescriptionAttribute(&conf, "test"),
					testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttribute(&conf, 10485760),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "name", "test"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "description", "test"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "minimum_compression_size", "10485760"),
					resource.TestCheckResourceAttrSet("aws_api_gateway_rest_api.test", "created_date"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "binary_media_types.#", "1"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "binary_media_types.0", "application/octet-stream"),
				),
			},

			{
				Config: testAccAWSAPIGatewayRestAPIDisableCompressionConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttributeIsNil(&conf),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "minimum_compression_size", "-1"),
				),
			},
		},
	})
}

func TestAccAWSAPIGatewayRestApi_openapi(t *testing.T) {
	var conf apigateway.RestApi

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAPIGatewayRestAPIDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSAPIGatewayRestAPIConfigOpenAPI,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPINameAttribute(&conf, "test"),
					testAccCheckAWSAPIGatewayRestAPIRoutes(&conf, []string{"/", "/test"}),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "name", "test"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "description", ""),
					resource.TestCheckResourceAttrSet("aws_api_gateway_rest_api.test", "created_date"),
					resource.TestCheckNoResourceAttr("aws_api_gateway_rest_api.test", "binary_media_types"),
				),
			},

			{
				Config: testAccAWSAPIGatewayRestAPIUpdateConfigOpenAPI,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPINameAttribute(&conf, "test"),
					testAccCheckAWSAPIGatewayRestAPIRoutes(&conf, []string{"/", "/update"}),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "name", "test"),
					resource.TestCheckResourceAttrSet("aws_api_gateway_rest_api.test", "created_date"),
				),
			},
		},
	})
}

func testAccCheckAWSAPIGatewayRestAPINameAttribute(conf *apigateway.RestApi, name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if *conf.Name != name {
			return fmt.Errorf("Wrong Name: %q", *conf.Name)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIDescriptionAttribute(conf *apigateway.RestApi, description string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if *conf.Description != description {
			return fmt.Errorf("Wrong Description: %q", *conf.Description)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttribute(conf *apigateway.RestApi, minimumCompressionSize int64) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if conf.MinimumCompressionSize == nil {
			return fmt.Errorf("MinimumCompressionSize should not be nil")
		}
		if *conf.MinimumCompressionSize != minimumCompressionSize {
			return fmt.Errorf("Wrong MinimumCompressionSize: %d", *conf.MinimumCompressionSize)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttributeIsNil(conf *apigateway.RestApi) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if conf.MinimumCompressionSize != nil {
			return fmt.Errorf("MinimumCompressionSize should be nil: %d", *conf.MinimumCompressionSize)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIRoutes(conf *apigateway.RestApi, routes []string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.Meta().(*AWSClient).apigateway

		resp, err := conn.GetResources(&apigateway.GetResourcesInput{
			RestApiId: conf.Id,
		})
		if err != nil {
			return err
		}

		actualRoutePaths := map[string]bool{}
		for _, resource := range resp.Items {
			actualRoutePaths[*resource.Path] = true
		}

		for _, route := range routes {
			if _, ok := actualRoutePaths[route]; !ok {
				return fmt.Errorf("Expected path %v but did not find it in %v", route, actualRoutePaths)
			}
			delete(actualRoutePaths, route)
		}

		if len(actualRoutePaths) > 0 {
			return fmt.Errorf("Found unexpected paths %v", actualRoutePaths)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIExists(n string, res *apigateway.RestApi) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No API Gateway ID is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).apigateway

		req := &apigateway.GetRestApiInput{
			RestApiId: aws.String(rs.Primary.ID),
		}
		describe, err := conn.GetRestApi(req)
		if err != nil {
			return err
		}

		if *describe.Id != rs.Primary.ID {
			return fmt.Errorf("APIGateway not found")
		}

		*res = *describe

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).apigateway

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_api_gateway_rest_api" {
			continue
		}

		req := &apigateway.GetRestApisInput{}
		describe, err := conn.GetRestApis(req)

		if err == nil {
			if len(describe.Items) != 0 &&
				*describe.Items[0].Id == rs.Primary.ID {
				return fmt.Errorf("API Gateway still exists")
			}
		}

		return err
	}

	return nil
}

const testAccAWSAPIGatewayRestAPIConfig = `
resource "aws_api_gateway_rest_api" "test" {
  name = "bar"
  minimum_compression_size = 0
}
`

const testAccAWSAPIGatewayRestAPIUpdateConfig = `
resource "aws_api_gateway_rest_api" "test" {
  name = "test"
  description = "test"
  binary_media_types = ["application/octet-stream"]
  minimum_compression_size = 10485760
}
`

const testAccAWSAPIGatewayRestAPIDisableCompressionConfig = `
resource "aws_api_gateway_rest_api" "test" {
  name = "test"
  description = "test"
  binary_media_types = ["application/octet-stream"]
  minimum_compression_size = -1
}
`

const testAccAWSAPIGatewayRestAPIConfigOpenAPI = `
resource "aws_api_gateway_rest_api" "test" {
  name = "test"
  body = <<EOF
{
  "swagger": "2.0",
  "info": {
    "title": "test",
    "version": "2017-04-20T04:08:08Z"
  },
  "schemes": [
    "https"
  ],
  "paths": {
    "/test": {
      "get": {
        "responses": {
          "200": {
            "description": "200 response"
          }
        },
        "x-amazon-apigateway-integration": {
          "type": "HTTP",
          "uri": "https://www.google.de",
          "httpMethod": "GET",
          "responses": {
            "default": {
              "statusCode": 200
            }
          }
        }
      }
    }
  }
}
EOF
}
`

const testAccAWSAPIGatewayRestAPIUpdateConfigOpenAPI = `
resource "aws_api_gateway_rest_api" "test" {
  name = "test"
  body = <<EOF
{
  "swagger": "2.0",
  "info": {
    "title": "test",
    "version": "2017-04-20T04:08:08Z"
  },
  "schemes": [
    "https"
  ],
  "paths": {
    "/update": {
      "get": {
        "responses": {
          "200": {
            "description": "200 response"
          }
        },
        "x-amazon-apigateway-integration": {
          "type": "HTTP",
          "uri": "https://www.google.de",
          "httpMethod": "GET",
          "responses": {
            "default": {
              "statusCode": 200
            }
          }
        }
      }
    }
  }
}
EOF
}
`
<commit_msg>tests/resource/aws_api_gateway_rest_api: Add sweeper<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/apigateway"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func init() {
	resource.AddTestSweepers("aws_api_gateway_rest_api", &resource.Sweeper{
		Name: "aws_api_gateway_rest_api",
		F:    testSweepAPIGatewayRestApis,
	})
}

func testSweepAPIGatewayRestApis(region string) error {
	client, err := sharedClientForRegion(region)
	if err != nil {
		return fmt.Errorf("error getting client: %s", err)
	}
	conn := client.(*AWSClient).apigateway

	// https://github.com/terraform-providers/terraform-provider-aws/issues/3808
	prefixes := []string{
		"test",
		"tf_acc_",
		"tf-acc-",
	}

	err = conn.GetRestApisPages(&apigateway.GetRestApisInput{}, func(page *apigateway.GetRestApisOutput, lastPage bool) bool {
		for _, item := range page.Items {
			skip := true
			for _, prefix := range prefixes {
				if strings.HasPrefix(*item.Name, prefix) {
					skip = false
					break
				}
			}
			if skip {
				log.Printf("[INFO] Skipping API Gateway REST API: %s", *item.Name)
				continue
			}

			input := &apigateway.DeleteRestApiInput{
				RestApiId: item.Id,
			}
			log.Printf("[INFO] Deleting API Gateway REST API: %s", input)
			// TooManyRequestsException: Too Many Requests can take over a minute to resolve itself
			err := resource.Retry(2*time.Minute, func() *resource.RetryError {
				_, err := conn.DeleteRestApi(input)
				if err != nil {
					if isAWSErr(err, apigateway.ErrCodeTooManyRequestsException, "") {
						return resource.RetryableError(err)
					}
					return resource.NonRetryableError(err)
				}
				return nil
			})
			if err != nil {
				log.Printf("[ERROR] Failed to delete API Gateway REST API %s: %s", *item.Name, err)
				continue
			}
		}
		return !lastPage
	})
	if err != nil {
		return fmt.Errorf("Error retrieving API Gateway REST APIs: %s", err)
	}

	return nil
}

func TestAccAWSAPIGatewayRestApi_basic(t *testing.T) {
	var conf apigateway.RestApi

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAPIGatewayRestAPIDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSAPIGatewayRestAPIConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPINameAttribute(&conf, "bar"),
					testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttribute(&conf, 0),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "name", "bar"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "description", ""),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "minimum_compression_size", "0"),
					resource.TestCheckResourceAttrSet("aws_api_gateway_rest_api.test", "created_date"),
					resource.TestCheckNoResourceAttr("aws_api_gateway_rest_api.test", "binary_media_types"),
				),
			},

			{
				Config: testAccAWSAPIGatewayRestAPIUpdateConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPINameAttribute(&conf, "test"),
					testAccCheckAWSAPIGatewayRestAPIDescriptionAttribute(&conf, "test"),
					testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttribute(&conf, 10485760),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "name", "test"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "description", "test"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "minimum_compression_size", "10485760"),
					resource.TestCheckResourceAttrSet("aws_api_gateway_rest_api.test", "created_date"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "binary_media_types.#", "1"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "binary_media_types.0", "application/octet-stream"),
				),
			},

			{
				Config: testAccAWSAPIGatewayRestAPIDisableCompressionConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttributeIsNil(&conf),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "minimum_compression_size", "-1"),
				),
			},
		},
	})
}

func TestAccAWSAPIGatewayRestApi_openapi(t *testing.T) {
	var conf apigateway.RestApi

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAPIGatewayRestAPIDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSAPIGatewayRestAPIConfigOpenAPI,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPINameAttribute(&conf, "test"),
					testAccCheckAWSAPIGatewayRestAPIRoutes(&conf, []string{"/", "/test"}),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "name", "test"),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "description", ""),
					resource.TestCheckResourceAttrSet("aws_api_gateway_rest_api.test", "created_date"),
					resource.TestCheckNoResourceAttr("aws_api_gateway_rest_api.test", "binary_media_types"),
				),
			},

			{
				Config: testAccAWSAPIGatewayRestAPIUpdateConfigOpenAPI,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAPIGatewayRestAPIExists("aws_api_gateway_rest_api.test", &conf),
					testAccCheckAWSAPIGatewayRestAPINameAttribute(&conf, "test"),
					testAccCheckAWSAPIGatewayRestAPIRoutes(&conf, []string{"/", "/update"}),
					resource.TestCheckResourceAttr("aws_api_gateway_rest_api.test", "name", "test"),
					resource.TestCheckResourceAttrSet("aws_api_gateway_rest_api.test", "created_date"),
				),
			},
		},
	})
}

func testAccCheckAWSAPIGatewayRestAPINameAttribute(conf *apigateway.RestApi, name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if *conf.Name != name {
			return fmt.Errorf("Wrong Name: %q", *conf.Name)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIDescriptionAttribute(conf *apigateway.RestApi, description string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if *conf.Description != description {
			return fmt.Errorf("Wrong Description: %q", *conf.Description)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttribute(conf *apigateway.RestApi, minimumCompressionSize int64) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if conf.MinimumCompressionSize == nil {
			return fmt.Errorf("MinimumCompressionSize should not be nil")
		}
		if *conf.MinimumCompressionSize != minimumCompressionSize {
			return fmt.Errorf("Wrong MinimumCompressionSize: %d", *conf.MinimumCompressionSize)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIMinimumCompressionSizeAttributeIsNil(conf *apigateway.RestApi) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if conf.MinimumCompressionSize != nil {
			return fmt.Errorf("MinimumCompressionSize should be nil: %d", *conf.MinimumCompressionSize)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIRoutes(conf *apigateway.RestApi, routes []string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.Meta().(*AWSClient).apigateway

		resp, err := conn.GetResources(&apigateway.GetResourcesInput{
			RestApiId: conf.Id,
		})
		if err != nil {
			return err
		}

		actualRoutePaths := map[string]bool{}
		for _, resource := range resp.Items {
			actualRoutePaths[*resource.Path] = true
		}

		for _, route := range routes {
			if _, ok := actualRoutePaths[route]; !ok {
				return fmt.Errorf("Expected path %v but did not find it in %v", route, actualRoutePaths)
			}
			delete(actualRoutePaths, route)
		}

		if len(actualRoutePaths) > 0 {
			return fmt.Errorf("Found unexpected paths %v", actualRoutePaths)
		}

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIExists(n string, res *apigateway.RestApi) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No API Gateway ID is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).apigateway

		req := &apigateway.GetRestApiInput{
			RestApiId: aws.String(rs.Primary.ID),
		}
		describe, err := conn.GetRestApi(req)
		if err != nil {
			return err
		}

		if *describe.Id != rs.Primary.ID {
			return fmt.Errorf("APIGateway not found")
		}

		*res = *describe

		return nil
	}
}

func testAccCheckAWSAPIGatewayRestAPIDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).apigateway

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_api_gateway_rest_api" {
			continue
		}

		req := &apigateway.GetRestApisInput{}
		describe, err := conn.GetRestApis(req)

		if err == nil {
			if len(describe.Items) != 0 &&
				*describe.Items[0].Id == rs.Primary.ID {
				return fmt.Errorf("API Gateway still exists")
			}
		}

		return err
	}

	return nil
}

const testAccAWSAPIGatewayRestAPIConfig = `
resource "aws_api_gateway_rest_api" "test" {
  name = "bar"
  minimum_compression_size = 0
}
`

const testAccAWSAPIGatewayRestAPIUpdateConfig = `
resource "aws_api_gateway_rest_api" "test" {
  name = "test"
  description = "test"
  binary_media_types = ["application/octet-stream"]
  minimum_compression_size = 10485760
}
`

const testAccAWSAPIGatewayRestAPIDisableCompressionConfig = `
resource "aws_api_gateway_rest_api" "test" {
  name = "test"
  description = "test"
  binary_media_types = ["application/octet-stream"]
  minimum_compression_size = -1
}
`

const testAccAWSAPIGatewayRestAPIConfigOpenAPI = `
resource "aws_api_gateway_rest_api" "test" {
  name = "test"
  body = <<EOF
{
  "swagger": "2.0",
  "info": {
    "title": "test",
    "version": "2017-04-20T04:08:08Z"
  },
  "schemes": [
    "https"
  ],
  "paths": {
    "/test": {
      "get": {
        "responses": {
          "200": {
            "description": "200 response"
          }
        },
        "x-amazon-apigateway-integration": {
          "type": "HTTP",
          "uri": "https://www.google.de",
          "httpMethod": "GET",
          "responses": {
            "default": {
              "statusCode": 200
            }
          }
        }
      }
    }
  }
}
EOF
}
`

const testAccAWSAPIGatewayRestAPIUpdateConfigOpenAPI = `
resource "aws_api_gateway_rest_api" "test" {
  name = "test"
  body = <<EOF
{
  "swagger": "2.0",
  "info": {
    "title": "test",
    "version": "2017-04-20T04:08:08Z"
  },
  "schemes": [
    "https"
  ],
  "paths": {
    "/update": {
      "get": {
        "responses": {
          "200": {
            "description": "200 response"
          }
        },
        "x-amazon-apigateway-integration": {
          "type": "HTTP",
          "uri": "https://www.google.de",
          "httpMethod": "GET",
          "responses": {
            "default": {
              "statusCode": 200
            }
          }
        }
      }
    }
  }
}
EOF
}
`
<|endoftext|>


<commit_before>/*
Copyright 2012 Google Inc.

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

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

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

package main

import (
	"bytes"
	"compress/zlib"
	"flag"
	"fmt"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

const (
	maxUncompressed = 50 << 10 // 50KB
	// Threshold ratio for compression.
	// Files which don't compress at least as well are kept uncompressed.
	zRatio = 0.5
)

func main() {
	flag.Parse()
	dir := "."
	switch flag.NArg() {
	case 0:
	case 1:
		dir = flag.Arg(0)
		if err := os.Chdir(dir); err != nil {
			log.Fatalf("chdir(%q) = %v", dir, err)
		}
	default:
		fmt.Fprintf(os.Stderr, "usage: genfileembed [<dir>]\n")
		os.Exit(2)
	}

	pkgName, filePattern, err := parseFileEmbed()
	if err != nil {
		log.Fatalf("Error parsing %s/fileembed.go: %v", dir, err)
	}

	for _, fileName := range matchingFiles(filePattern) {
		embedName := "zembed_" + fileName + ".go"
		fi, err := os.Stat(fileName)
		if err != nil {
			log.Fatal(err)
		}
		efi, err := os.Stat(embedName)
		if err == nil && !efi.ModTime().Before(fi.ModTime()) {
			continue
		}
		log.Printf("Updating %s (package %s)", filepath.Join(dir, embedName), pkgName)

		zb, fileSize := compressFile(fileName)
		ratio := float64(len(zb)) / float64(fileSize)
		var bs []byte
		byteStreamType := ""
		if fileSize < maxUncompressed || ratio > zRatio {
			byteStreamType = "fileembed.String"
			bs, err = ioutil.ReadFile(fileName)
			if err != nil {
				log.Fatal(err)
			}
		} else {
			byteStreamType = "fileembed.ZlibCompressed"
			bs = zb
		}
		qb := quote(bs) // quoted bytes

		var b bytes.Buffer
		fmt.Fprintf(&b, "// THIS FILE IS AUTO-GENERATED FROM %s\n", fileName)
		fmt.Fprintf(&b, "// DO NOT EDIT.\n\n")
		fmt.Fprintf(&b, "package %s\n\n", pkgName)
		fmt.Fprintf(&b, "import \"time\"\n\n")
		fmt.Fprintf(&b, "import \"camlistore.org/pkg/fileembed\"\n\n")
		fmt.Fprintf(&b, "func init() {\n\tFiles.Add(%q, %d, %s(%s), time.Unix(0, %d));\n}\n",
			fileName, fileSize, byteStreamType, qb, fi.ModTime().UnixNano())

		// gofmt it
		fset := token.NewFileSet()
		ast, err := parser.ParseFile(fset, "", b.Bytes(), parser.ParseComments)
		if err != nil {
			log.Fatal(err)
		}

		var clean bytes.Buffer
		config := &printer.Config{
			Mode:     printer.TabIndent | printer.UseSpaces,
			Tabwidth: 8,
		}
		err = config.Fprint(&clean, fset, ast)
		if err != nil {
			log.Fatal(err)
		}

		if err := ioutil.WriteFile(embedName, clean.Bytes(), 0644); err != nil {
			log.Fatal(err)
		}
	}
}

func compressFile(fileName string) ([]byte, int64) {
	var zb bytes.Buffer
	f, err := os.Open(fileName)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	w := zlib.NewWriter(&zb)
	n, err := io.Copy(w, f)
	if err != nil {
		log.Fatal(err)
	}
	w.Close()
	return zb.Bytes(), n
}

func quote(bs []byte) []byte {
	var qb bytes.Buffer
	qb.WriteByte('"')
	run := 0
	for _, b := range bs {
		if b == '\n' {
			qb.WriteString(`\n`)
		}
		if b == '\n' || run > 80 {
			qb.WriteString("\" +\n\t\"")
			run = 0
		}
		if b == '\n' {
			continue
		}
		run++
		if b == '\\' {
			qb.WriteString(`\\`)
			continue
		}
		if b == '"' {
			qb.WriteString(`\"`)
			continue
		}
		if (b >= 32 && b <= 126) || b == '\t' {
			qb.WriteByte(b)
			continue
		}
		fmt.Fprintf(&qb, "\\x%02x", b)
	}
	qb.WriteByte('"')
	return qb.Bytes()
}

func matchingFiles(p *regexp.Regexp) []string {
	var f []string
	d, err := os.Open(".")
	if err != nil {
		log.Fatal(err)
	}
	defer d.Close()
	names, err := d.Readdirnames(-1)
	if err != nil {
		log.Fatal(err)
	}
	for _, n := range names {
		if strings.HasPrefix(n, "zembed_") {
			continue
		}
		if p.MatchString(n) {
			f = append(f, n)
		}
	}
	return f
}

func parseFileEmbed() (pkgName string, filePattern *regexp.Regexp, err error) {
	fe, err := os.Open("fileembed.go")
	if err != nil {
		return
	}
	defer fe.Close()

	fs := token.NewFileSet()
	astf, err := parser.ParseFile(fs, "fileembed.go", fe, parser.PackageClauseOnly|parser.ParseComments)
	if err != nil {
		return
	}
	pkgName = astf.Name.Name

	if astf.Doc == nil {
		err = fmt.Errorf("no package comment before the %q line", "package "+pkgName)
		return
	}

	pkgComment := astf.Doc.Text()
	findPattern := regexp.MustCompile(`(?m)^#fileembed\s+pattern\s+(\S+)\s*$`)
	m := findPattern.FindStringSubmatch(pkgComment)
	if m == nil {
		err = fmt.Errorf("package comment lacks line of form: #fileembed pattern <pattern>")
		return
	}
	pattern := m[1]
	filePattern, err = regexp.Compile(pattern)
	if err != nil {
		err = fmt.Errorf("bad regexp %q: %v", pattern, err)
		return
	}
	return
}
<commit_msg>Add -fileembed-package and -all flags<commit_after>/*
Copyright 2012 Google Inc.

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

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

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

package main

import (
	"bytes"
	"compress/zlib"
	"flag"
	"fmt"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

const (
	maxUncompressed = 50 << 10 // 50KB
	// Threshold ratio for compression.
	// Files which don't compress at least as well are kept uncompressed.
	zRatio = 0.5
)

func usage() {
	fmt.Fprintf(os.Stderr, "usage: genfileembed [flags] [<dir>]\n")
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	var processAll bool
	flag.BoolVar(&processAll, "all", false, "process all files (if false, only process modified files)")
	var fileEmbedPkgPath string
	flag.StringVar(&fileEmbedPkgPath, "fileembed-package", "camlistore.org/pkg/fileembed", "the Go package name for fileembed. If you have vendored fileembed (e.g. with goven), you can use this flag to ensure that generated code imports the vendored package.")
	flag.Usage = usage
	flag.Parse()

	dir := "."
	switch flag.NArg() {
	case 0:
	case 1:
		dir = flag.Arg(0)
		if err := os.Chdir(dir); err != nil {
			log.Fatalf("chdir(%q) = %v", dir, err)
		}
	default:
		flag.Usage()
	}

	pkgName, filePattern, err := parseFileEmbed()
	if err != nil {
		log.Fatalf("Error parsing %s/fileembed.go: %v", dir, err)
	}

	for _, fileName := range matchingFiles(filePattern) {
		embedName := "zembed_" + fileName + ".go"
		fi, err := os.Stat(fileName)
		if err != nil {
			log.Fatal(err)
		}
		efi, err := os.Stat(embedName)
		if err == nil && !efi.ModTime().Before(fi.ModTime()) && !processAll {
			continue
		}
		log.Printf("Updating %s (package %s)", filepath.Join(dir, embedName), pkgName)

		zb, fileSize := compressFile(fileName)
		ratio := float64(len(zb)) / float64(fileSize)
		var bs []byte
		byteStreamType := ""
		if fileSize < maxUncompressed || ratio > zRatio {
			byteStreamType = "fileembed.String"
			bs, err = ioutil.ReadFile(fileName)
			if err != nil {
				log.Fatal(err)
			}
		} else {
			byteStreamType = "fileembed.ZlibCompressed"
			bs = zb
		}
		qb := quote(bs) // quoted bytes

		var b bytes.Buffer
		fmt.Fprintf(&b, "// THIS FILE IS AUTO-GENERATED FROM %s\n", fileName)
		fmt.Fprintf(&b, "// DO NOT EDIT.\n\n")
		fmt.Fprintf(&b, "package %s\n\n", pkgName)
		fmt.Fprintf(&b, "import \"time\"\n\n")
		fmt.Fprintf(&b, "import \""+fileEmbedPkgPath+"\"\n\n")
		fmt.Fprintf(&b, "func init() {\n\tFiles.Add(%q, %d, %s(%s), time.Unix(0, %d));\n}\n",
			fileName, fileSize, byteStreamType, qb, fi.ModTime().UnixNano())

		// gofmt it
		fset := token.NewFileSet()
		ast, err := parser.ParseFile(fset, "", b.Bytes(), parser.ParseComments)
		if err != nil {
			log.Fatal(err)
		}

		var clean bytes.Buffer
		config := &printer.Config{
			Mode:     printer.TabIndent | printer.UseSpaces,
			Tabwidth: 8,
		}
		err = config.Fprint(&clean, fset, ast)
		if err != nil {
			log.Fatal(err)
		}

		if err := ioutil.WriteFile(embedName, clean.Bytes(), 0644); err != nil {
			log.Fatal(err)
		}
	}
}

func compressFile(fileName string) ([]byte, int64) {
	var zb bytes.Buffer
	f, err := os.Open(fileName)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	w := zlib.NewWriter(&zb)
	n, err := io.Copy(w, f)
	if err != nil {
		log.Fatal(err)
	}
	w.Close()
	return zb.Bytes(), n
}

func quote(bs []byte) []byte {
	var qb bytes.Buffer
	qb.WriteByte('"')
	run := 0
	for _, b := range bs {
		if b == '\n' {
			qb.WriteString(`\n`)
		}
		if b == '\n' || run > 80 {
			qb.WriteString("\" +\n\t\"")
			run = 0
		}
		if b == '\n' {
			continue
		}
		run++
		if b == '\\' {
			qb.WriteString(`\\`)
			continue
		}
		if b == '"' {
			qb.WriteString(`\"`)
			continue
		}
		if (b >= 32 && b <= 126) || b == '\t' {
			qb.WriteByte(b)
			continue
		}
		fmt.Fprintf(&qb, "\\x%02x", b)
	}
	qb.WriteByte('"')
	return qb.Bytes()
}

func matchingFiles(p *regexp.Regexp) []string {
	var f []string
	d, err := os.Open(".")
	if err != nil {
		log.Fatal(err)
	}
	defer d.Close()
	names, err := d.Readdirnames(-1)
	if err != nil {
		log.Fatal(err)
	}
	for _, n := range names {
		if strings.HasPrefix(n, "zembed_") {
			continue
		}
		if p.MatchString(n) {
			f = append(f, n)
		}
	}
	return f
}

func parseFileEmbed() (pkgName string, filePattern *regexp.Regexp, err error) {
	fe, err := os.Open("fileembed.go")
	if err != nil {
		return
	}
	defer fe.Close()

	fs := token.NewFileSet()
	astf, err := parser.ParseFile(fs, "fileembed.go", fe, parser.PackageClauseOnly|parser.ParseComments)
	if err != nil {
		return
	}
	pkgName = astf.Name.Name

	if astf.Doc == nil {
		err = fmt.Errorf("no package comment before the %q line", "package "+pkgName)
		return
	}

	pkgComment := astf.Doc.Text()
	findPattern := regexp.MustCompile(`(?m)^#fileembed\s+pattern\s+(\S+)\s*$`)
	m := findPattern.FindStringSubmatch(pkgComment)
	if m == nil {
		err = fmt.Errorf("package comment lacks line of form: #fileembed pattern <pattern>")
		return
	}
	pattern := m[1]
	filePattern, err = regexp.Compile(pattern)
	if err != nil {
		err = fmt.Errorf("bad regexp %q: %v", pattern, err)
		return
	}
	return
}
<|endoftext|>


<commit_before><commit_msg>Add "json:,string" annotation to Volume.Size<commit_after><|endoftext|>


<commit_before><commit_msg>retry<commit_after><|endoftext|>


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

package mixnet

import (
	"container/list"
	"errors"
	"math/rand"
	"net"

	"github.com/golang/glog"
)

type sendQueueError struct {
	id uint64 // Serial identifier of sender to which error pertains.
	error
}

// The SendQueue structure maps a serial identifier corresponding to a sender
// (in the router context) to a destination. It also maintains a message buffer
// for each sender. Once there messages ready on enough buffers, a batch of
// messages are transmitted simultaneously.
type SendQueue struct {
	batchSize int // Number of messages to transmit in a roud.
	ct        int // Current number of buffers with messages ready.

	network string // Network protocol, e.g. "tcp".

	nextAddr   map[uint64]string     // Address of destination.
	nextConn   map[uint64]net.Conn   // Connection to destination.
	sendBuffer map[uint64]*list.List // Message buffer of sender.

	queue chan Queueable      // Channel for queueing messages/directives.
	err   chan sendQueueError // Channel for handling errors.
}

// NewSendQueue creates a new SendQueue structure.
func NewSendQueue(network string, batchSize int) (sq *SendQueue) {
	sq = new(SendQueue)
	sq.batchSize = batchSize
	sq.network = network

	sq.nextAddr = make(map[uint64]string)
	sq.nextConn = make(map[uint64]net.Conn)
	sq.sendBuffer = make(map[uint64]*list.List)

	sq.queue = make(chan Queueable)
	sq.err = make(chan sendQueueError)
	return sq
}

// Enqueue adds a queueable object, such as a message or directive, to the
// send queue.
func (sq *SendQueue) Enqueue(q *Queueable) {
	sq.queue <- *q
}

// DoSendQueue adds messages to a queue and transmits messages in batches.
// Typically a message is a cell, but when the calling router is an exit point,
// the message length is arbitrary. A batch is transmitted when there are
// messages on batchSize distinct sender channels.
func (sq *SendQueue) DoSendQueue(kill <-chan bool) {
	for {
		select {
		case <-kill:
			for _, c := range sq.nextConn {
				c.Close()
			}
			return

		case q := <-sq.queue:
			// Set the next-hop address. We don't allow the destination address
			// or connection to be overwritten in order to avoid accumulating
			// stale connections on routers.
			if _, def := sq.nextAddr[*q.Id]; !def && q.Addr != nil {
				sq.nextAddr[*q.Id] = *q.Addr
			}

			if q.Dir != nil {
				sq.err <- sendQueueError{*q.Id,
					errors.New("directives not implemented")}

			} else if _, def := sq.nextAddr[*q.Id]; !def && q.Msg != nil {
				sq.err <- sendQueueError{*q.Id,
					errors.New("request to send message without a destination")}

			} else {

				// Create a send buffer for the sender ID if it doesn't exist.
				if _, def := sq.sendBuffer[*q.Id]; !def {
					sq.sendBuffer[*q.Id] = list.New()
				}
				buf := sq.sendBuffer[*q.Id]

				// The buffer was empty but now has a message ready; increment
				// the counter.
				if buf.Len() == 0 {
					sq.ct++
				}

				// Add message to send buffer.
				buf.PushBack(q.Msg)
			}

			// Transmit the message batch if it is full.
			if sq.ct == sq.batchSize {
				sq.dequeue()
			}
		}
	}
}

// DoSendQueueErrorHandler handles errors produced by DoSendQueue. When this
// is fully fleshed out, it will enqueue into the response queue a Directive
// containing an error message. For now, just print out the error.
func (sq *SendQueue) DoSendQueueErrorHandler(kill <-chan bool) {
	for {
		select {
		case <-kill:
			return
		case err := <-sq.err:
			glog.Errorf("send queue (%d): %s\n", err.id, err.Error())
		}
	}
}

// dequeue sends one message from each send buffer for each serial ID in a
// random order. This is called by DoSendQueue and is not safe to call directly
// elsewhere.
func (sq *SendQueue) dequeue() {

	// Shuffle the serial IDs.
	order := rand.Perm(int(sq.ct)) // TODO(cjpatton) Use tao.GetRandomBytes().
	ids := make([]uint64, sq.ct)
	i := 0
	for id := range sq.sendBuffer {
		ids[order[i]] = id
		i++
	}

	// Issue a sendWorker thread for each message to be sent.
	ch := make(chan senderResult)
	for _, id := range ids {
		addr := sq.nextAddr[id]
		msg := sq.sendBuffer[id].Front().Value.([]byte)
		c, def := sq.nextConn[id]
		go senderWorker(network, addr, id, msg, c, def, ch, sq.err)
	}

	// Wait for workers to finish.
	for _ = range ids {
		res := <-ch
		if res.c != nil {
			// Save the connection.
			sq.nextConn[res.id] = res.c

			// Pop the message from the buffer and decrement the counter
			// if the buffer is empty.
			buf := sq.sendBuffer[res.id]
			buf.Remove(buf.Front())
			if buf.Len() == 0 {
				sq.ct--
			}
		}
	}
}

type senderResult struct {
	c  net.Conn
	id uint64
}

func senderWorker(network, addr string, id uint64, msg []byte, c net.Conn, def bool,
	res chan<- senderResult, err chan<- sendQueueError) {
	var e error

	// Wait to connect until the queue is dequeued in order to prevent
	// an observer from correlating an incoming cell with the handshake
	// with the destination server.
	if !def {
		c, e = net.Dial(network, addr) // TODO(cjpatton) timeout.
		if e != nil {
			err <- sendQueueError{id, e}
			res <- senderResult{nil, id}
			return
		}
	}

	// Send the message.
	if _, e := c.Write(msg); e != nil {
		err <- sendQueueError{id, e}
		res <- senderResult{nil, id}
		return
	}

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

package mixnet

import (
	"container/list"
	"errors"
	"math/rand"
	"net"

	"github.com/golang/glog"
)

type sendQueueError struct {
	id uint64 // Serial identifier of sender to which error pertains.
	error
}

// The SendQueue structure maps a serial identifier corresponding to a sender
// (in the router context) to a destination. It also maintains a message buffer
// for each sender. Once there messages ready on enough buffers, a batch of
// messages are transmitted simultaneously.
type SendQueue struct {
	batchSize int // Number of messages to transmit in a roud.
	ct        int // Current number of buffers with messages ready.

	network string // Network protocol, e.g. "tcp".

	nextAddr   map[uint64]string     // Address of destination.
	nextConn   map[uint64]net.Conn   // Connection to destination.
	sendBuffer map[uint64]*list.List // Message buffer of sender.

	queue chan Queueable      // Channel for queueing messages/directives.
	err   chan sendQueueError // Channel for handling errors.
}

// NewSendQueue creates a new SendQueue structure.
func NewSendQueue(network string, batchSize int) (sq *SendQueue) {
	sq = new(SendQueue)
	sq.batchSize = batchSize
	sq.network = network

	sq.nextAddr = make(map[uint64]string)
	sq.nextConn = make(map[uint64]net.Conn)
	sq.sendBuffer = make(map[uint64]*list.List)

	sq.queue = make(chan Queueable)
	sq.err = make(chan sendQueueError)
	return sq
}

// Enqueue adds a queueable object, such as a message or directive, to the
// send queue.
func (sq *SendQueue) Enqueue(q *Queueable) {
	sq.queue <- *q
}

// DoSendQueue adds messages to a queue and transmits messages in batches.
// Typically a message is a cell, but when the calling router is an exit point,
// the message length is arbitrary. A batch is transmitted when there are
// messages on batchSize distinct sender channels.
func (sq *SendQueue) DoSendQueue(kill <-chan bool) {
	for {
		select {
		case <-kill:
			for _, c := range sq.nextConn {
				c.Close()
			}
			return

		case q := <-sq.queue:
			// Set the next-hop address. We don't allow the destination address
			// or connection to be overwritten in order to avoid accumulating
			// stale connections on routers.
			if _, def := sq.nextAddr[*q.Id]; !def && q.Addr != nil {
				sq.nextAddr[*q.Id] = *q.Addr
			}

			if q.Dir != nil {
				sq.err <- sendQueueError{*q.Id,
					errors.New("directives not implemented")}

			} else if _, def := sq.nextAddr[*q.Id]; !def && q.Msg != nil {
				sq.err <- sendQueueError{*q.Id,
					errors.New("request to send message without a destination")}

			} else {

				// Create a send buffer for the sender ID if it doesn't exist.
				if _, def := sq.sendBuffer[*q.Id]; !def {
					sq.sendBuffer[*q.Id] = list.New()
				}
				buf := sq.sendBuffer[*q.Id]

				// The buffer was empty but now has a message ready; increment
				// the counter.
				if buf.Len() == 0 {
					sq.ct++
				}

				// Add message to send buffer.
				buf.PushBack(q.Msg)
			}

			// Transmit the message batch if it is full.
			if sq.ct == sq.batchSize {
				sq.dequeue()
			}
		}
	}
}

// DoSendQueueErrorHandler handles errors produced by DoSendQueue. When this
// is fully fleshed out, it will enqueue into the response queue a Directive
// containing an error message. For now, just print out the error.
func (sq *SendQueue) DoSendQueueErrorHandler(kill <-chan bool) {
	for {
		select {
		case <-kill:
			return
		case err := <-sq.err:
			glog.Errorf("send queue (%d): %s\n", err.id, err.Error())
		}
	}
}

// dequeue sends one message from each send buffer for each serial ID in a
// random order. This is called by DoSendQueue and is not safe to call directly
// elsewhere.
func (sq *SendQueue) dequeue() {

	// Shuffle the serial IDs.
	order := rand.Perm(int(sq.ct)) // TODO(cjpatton) Use tao.GetRandomBytes().
	ids := make([]uint64, sq.ct)
	i := 0
	for id := range sq.sendBuffer {
		ids[order[i]] = id
		i++
	}

	// Issue a sendWorker thread for each message to be sent.
	ch := make(chan senderResult)
	for _, id := range ids {
		addr := sq.nextAddr[id]
		msg := sq.sendBuffer[id].Front().Value.([]byte)
		c, def := sq.nextConn[id]
		go senderWorker(sq.network, addr, id, msg, c, def, ch, sq.err)
	}

	// Wait for workers to finish.
	for _ = range ids {
		res := <-ch
		if res.c != nil {
			// Save the connection.
			sq.nextConn[res.id] = res.c

			// Pop the message from the buffer and decrement the counter
			// if the buffer is empty.
			buf := sq.sendBuffer[res.id]
			buf.Remove(buf.Front())
			if buf.Len() == 0 {
				sq.ct--
			}
		}
	}
}

type senderResult struct {
	c  net.Conn
	id uint64
}

func senderWorker(network, addr string, id uint64, msg []byte, c net.Conn, def bool,
	res chan<- senderResult, err chan<- sendQueueError) {
	var e error

	// Wait to connect until the queue is dequeued in order to prevent
	// an observer from correlating an incoming cell with the handshake
	// with the destination server.
	if !def {
		c, e = net.Dial(network, addr) // TODO(cjpatton) timeout.
		if e != nil {
			err <- sendQueueError{id, e}
			res <- senderResult{nil, id}
			return
		}
	}

	// Send the message.
	if _, e := c.Write(msg); e != nil {
		err <- sendQueueError{id, e}
		res <- senderResult{nil, id}
		return
	}

	res <- senderResult{c, id}
}
<|endoftext|>


<commit_before>package alita

import (
	"reflect"
	"testing"
)

type JustfyTest struct {
	justfy Justify
	width  int
	src    string
	dst    string
}

var indexTestsJustfy = []JustfyTest{
	// Case odd
	{JustLeft, 5, "abc", "abc  "},
	{JustRight, 5, "abc", "  abc"},
	{JustCenter, 5, "abc", " abc "},

	// Case odd && equal
	{JustLeft, 5, "abcde", "abcde"},
	{JustRight, 5, "abcde", "abcde"},
	{JustCenter, 5, "abcde", "abcde"},

	// Case even
	{JustLeft, 6, "abc", "abc   "},
	{JustRight, 6, "abc", "   abc"},
	{JustCenter, 6, "abc", " abc  "},

	// Case even && equal
	{JustLeft, 6, "abcdef", "abcdef"},
	{JustRight, 6, "abcdef", "abcdef"},
	{JustCenter, 6, "abcdef", "abcdef"},

	// Case min
	{JustLeft, 2, "a", "a "},
	{JustRight, 2, "a", " a"},
	{JustCenter, 2, "a", "a "},

	// Case over
	{JustLeft, 4, "abcdefg", "abcdefg"},
	{JustRight, 4, "abcdefg", "abcdefg"},
	{JustCenter, 4, "abcdefg", "abcdefg"},
}

func TestJustfy(t *testing.T) {
	for _, test := range indexTestsJustfy {
		actual := test.justfy.Just(test.width, test.src)
		expect := test.dst
		if actual != expect {
			kind := "lrc"[test.justfy]
			t.Errorf("%q.Justfy(%q) = %q; want %q",
				kind, test.src, actual, expect)
		}
	}
}

type UpdateWidthTest struct {
	a      []string
	before []int
	after  []int
}

var indexTestsUpdateWidth = []UpdateWidthTest{
	// Through update
	{[]string{"a"},
		[]int{}, []int{}},
	{[]string{"aaaa"},
		[]int{}, []int{}},
	{[]string{"aaaa"},
		[]int{3, 5}, []int{3, 5}},

	// Update
	{[]string{"aaa", "aaa"},
		[]int{1, 2}, []int{3, 3}},
	{[]string{"aaa", "aaa", "aaaa"},
		[]int{1, 2, 3}, []int{3, 3, 4}},

	// Update with stretch
	{[]string{"a", "aa"},
		[]int{}, []int{1, 2}},
	{[]string{"a", "aa", "aaa"},
		[]int{1, 2}, []int{1, 2, 3}},

	// No update
	{[]string{"a", "a"},
		[]int{1, 1}, []int{1, 1}},
	{[]string{"aaaaa", "aa"},
		[]int{3, 3}, []int{5, 3}},
	{[]string{"aaa", "aaa", "aaa"},
		[]int{4, 4, 4}, []int{4, 4, 4}},

	// Correspondence double-width character
	{[]string{"日本語", "日本語日本語"},
		[]int{1, 1}, []int{6, 12}},
	{[]string{"「おはよう」", "『こんにちは』"},
		[]int{1, 1}, []int{12, 14}},
}

func TestUpdateWidth(t *testing.T) {
	p := NewPadding()
	for _, test := range indexTestsUpdateWidth {
		p.width = test.before

		p.UpdateWidth(test.a)
		actual := p.width
		expect := test.after
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("%v -> UpdateWidth(%v) got %v; want %v",
				test.before, test.a, actual, expect)
		}
	}
}
<commit_msg>Add case for test of justfy<commit_after>package alita

import (
	"reflect"
	"testing"
)

type JustfyTest struct {
	justfy Justify
	width  int
	src    string
	dst    string
}

var indexTestsJustfy = []JustfyTest{
	// Case normal
	{JustLeft, 5, "abc", "abc  "},
	{JustRight, 5, "abc", "  abc"},
	{JustCenter, 5, "abc", " abc "},
	{JustLeft, 6, "abc", "abc   "},
	{JustRight, 6, "abc", "   abc"},
	{JustCenter, 6, "abc", " abc  "},

	// Case equal
	{JustLeft, 5, "abcde", "abcde"},
	{JustRight, 5, "abcde", "abcde"},
	{JustCenter, 5, "abcde", "abcde"},
	{JustLeft, 6, "abcdef", "abcdef"},
	{JustRight, 6, "abcdef", "abcdef"},
	{JustCenter, 6, "abcdef", "abcdef"},

	// Case min
	{JustLeft, 2, "a", "a "},
	{JustRight, 2, "a", " a"},
	{JustCenter, 2, "a", "a "},

	// Case over
	{JustLeft, 4, "abcdefg", "abcdefg"},
	{JustRight, 4, "abcdefg", "abcdefg"},
	{JustCenter, 4, "abcdefg", "abcdefg"},

	// Case minus
	{JustLeft, -5, "abc", "abc"},
	{JustRight, -5, "abc", "abc"},
	{JustCenter, -5, "abc", "abc"},

	// Case double-width character
	{JustLeft, 8, "日本語", "日本語  "},
	{JustRight, 8, "日本語", "  日本語"},
	{JustCenter, 8, "日本語", " 日本語 "},
	{JustLeft, 9, "日本語", "日本語   "},
	{JustRight, 9, "日本語", "   日本語"},
	{JustCenter, 9, "日本語", " 日本語  "},
}

func TestJustfy(t *testing.T) {
	for _, test := range indexTestsJustfy {
		actual := test.justfy.Just(test.width, test.src)
		expect := test.dst
		if actual != expect {
			kind := "lrc"[test.justfy]
			t.Errorf("%q.Justfy(%q) = %q; want %q",
				kind, test.src, actual, expect)
		}
	}
}

type UpdateWidthTest struct {
	a      []string
	before []int
	after  []int
}

var indexTestsUpdateWidth = []UpdateWidthTest{
	// Through update
	{[]string{"a"},
		[]int{}, []int{}},
	{[]string{"aaaa"},
		[]int{}, []int{}},
	{[]string{"aaaa"},
		[]int{3, 5}, []int{3, 5}},

	// Update
	{[]string{"aaa", "aaa"},
		[]int{1, 2}, []int{3, 3}},
	{[]string{"aaa", "aaa", "aaaa"},
		[]int{1, 2, 3}, []int{3, 3, 4}},

	// Update with stretch
	{[]string{"a", "aa"},
		[]int{}, []int{1, 2}},
	{[]string{"a", "aa", "aaa"},
		[]int{1, 2}, []int{1, 2, 3}},

	// No update
	{[]string{"a", "a"},
		[]int{1, 1}, []int{1, 1}},
	{[]string{"aaaaa", "aa"},
		[]int{3, 3}, []int{5, 3}},
	{[]string{"aaa", "aaa", "aaa"},
		[]int{4, 4, 4}, []int{4, 4, 4}},

	// Correspondence double-width character
	{[]string{"日本語", "日本語日本語"},
		[]int{1, 1}, []int{6, 12}},
	{[]string{"「おはよう」", "『こんにちは』"},
		[]int{1, 1}, []int{12, 14}},
}

func TestUpdateWidth(t *testing.T) {
	p := NewPadding()
	for _, test := range indexTestsUpdateWidth {
		p.width = test.before

		p.UpdateWidth(test.a)
		actual := p.width
		expect := test.after
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("%v -> UpdateWidth(%v) got %v; want %v",
				test.before, test.a, actual, expect)
		}
	}
}
<|endoftext|>


<commit_before>package lib

import (
	"context"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/stats"
	"gopkg.in/guregu/null.v3"
	"strconv"
	"sync"
	"time"
)

const TickRate = 1 * time.Millisecond

var (
	MetricVUs    = &stats.Metric{Name: "vus", Type: stats.Gauge}
	MetricVUsMax = &stats.Metric{Name: "vus_max", Type: stats.Gauge}
	MetricErrors = &stats.Metric{Name: "errors", Type: stats.Counter}

	ErrTooManyVUs = errors.New("More VUs than the maximum requested")
	ErrMaxTooLow  = errors.New("Can't lower max below current VU count")

	// Special error used to taint a test, without printing an error.
	ErrVUWantsTaint = errors.New("[ErrVUWantsTaint is never logged]")
)

type vuEntry struct {
	VU        VU
	Buffer    []stats.Sample
	ExtBuffer stats.Buffer
	Cancel    context.CancelFunc
}

type Engine struct {
	Runner    Runner
	Status    Status
	Stages    []Stage
	Collector stats.Collector
	Quit      bool
	Pause     sync.WaitGroup

	Metrics map[*stats.Metric]stats.Sink

	ctx    context.Context
	vus    []*vuEntry
	nextID int64

	vuMutex sync.Mutex
}

func NewEngine(r Runner) (*Engine, error) {
	e := &Engine{
		Runner: r,
		Status: Status{
			Running: null.BoolFrom(false),
			Tainted: null.BoolFrom(false),
			VUs:     null.IntFrom(0),
			VUsMax:  null.IntFrom(0),
			AtTime:  null.IntFrom(0),
		},
		Metrics: make(map[*stats.Metric]stats.Sink),
	}

	e.Status.Running = null.BoolFrom(false)
	e.Pause.Add(1)

	e.Status.VUs = null.IntFrom(0)
	e.Status.VUsMax = null.IntFrom(0)

	return e, nil
}

func (e *Engine) Run(ctx context.Context) error {
	if len(e.Stages) == 0 {
		return errors.New("Engine has no stages")
	}

	e.ctx = ctx
	e.nextID = 1

	if e.Collector != nil {
		go e.Collector.Run(ctx)
	} else {
		log.Debug("Engine: No Collector")
	}

	e.consumeEngineStats()

	ticker := time.NewTicker(TickRate)
	lastTick := time.Now()

loop:
	for {
		select {
		case now := <-ticker.C:
			timeDelta := now.Sub(lastTick)
			e.Status.AtTime.Int64 += int64(timeDelta)
			lastTick = now

			stage, left, ok := StageAt(e.Stages, time.Duration(e.Status.AtTime.Int64))
			if stage.StartVUs.Valid && stage.EndVUs.Valid {
				progress := (float64(stage.Duration.Int64-int64(left)) / float64(stage.Duration.Int64))
				vus := Lerp(stage.StartVUs.Int64, stage.EndVUs.Int64, progress)
				e.SetVUs(vus)
			}

			if !ok {
				e.SetRunning(false)

				if e.Quit {
					break loop
				} else {
					log.Info("Test finished, press Ctrl+C to exit")
				}
			}

			e.consumeEngineStats()
		case <-ctx.Done():
			break loop
		}
	}

	e.vus = nil

	e.Status.Running = null.BoolFrom(false)
	e.Status.VUs = null.IntFrom(0)
	e.Status.VUsMax = null.IntFrom(0)
	e.consumeEngineStats()

	return nil
}

func (e *Engine) SetRunning(running bool) {
	if running && !e.Status.Running.Bool {
		e.Pause.Done()
		log.Debug("Engine Unpaused")
	} else if !running && e.Status.Running.Bool {
		e.Pause.Add(1)
		log.Debug("Engine Paused")
	}
	e.Status.Running.Bool = running
}

func (e *Engine) SetVUs(v int64) error {
	if e.Status.VUs.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUs.Int64, "to": v}).Debug("Setting VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v > e.Status.VUsMax.Int64 {
		return ErrTooManyVUs
	}

	current := e.Status.VUs.Int64
	for i := current; i < v; i++ {
		entry := e.vus[i]
		if entry.Cancel != nil {
			panic(errors.New("ATTEMPTED TO RESCHEDULE RUNNING VU"))
		}

		ctx, cancel := context.WithCancel(e.ctx)
		entry.Cancel = cancel

		if err := entry.VU.Reconfigure(e.nextID); err != nil {
			return err
		}
		go e.runVU(ctx, e.nextID, entry)
		e.nextID++
	}
	for i := current - 1; i >= v; i-- {
		entry := e.vus[i]
		entry.Cancel()
		entry.Cancel = nil
	}

	e.Status.VUs.Int64 = v
	return nil
}

func (e *Engine) SetMaxVUs(v int64) error {
	if e.Status.VUsMax.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUsMax.Int64, "to": v}).Debug("Setting Max VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v < e.Status.VUs.Int64 {
		return ErrMaxTooLow
	}

	current := e.Status.VUsMax.Int64
	if v > current {
		vus := e.vus
		for i := current; i < v; i++ {
			vu, err := e.Runner.NewVU()
			if err != nil {
				return err
			}
			entry := &vuEntry{VU: vu}
			if e.Collector != nil {
				entry.ExtBuffer = e.Collector.Buffer()
			}
			vus = append(vus, entry)
		}
		e.vus = vus
	} else if v < current {
		e.vus = e.vus[:v]
	}

	e.Status.VUsMax.Int64 = v
	return nil
}

func (e *Engine) runVU(ctx context.Context, id int64, vu *vuEntry) {
	idString := strconv.FormatInt(id, 10)

waitForPause:
	e.Pause.Wait()

	for {
		samples, err := vu.VU.RunOnce(ctx)

		// If the context is cancelled, the iteration is likely to produce erroneous output
		// due to cancelled HTTP requests and whatnot. Discard output from such runs.
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err != nil {
			e.Status.Tainted.Bool = true

			if err != ErrVUWantsTaint {
				if s, ok := err.(fmt.Stringer); ok {
					log.Error(s.String())
				} else {
					log.WithError(err).Error("Runtime Error")
				}

				samples = append(samples, stats.Sample{
					Metric: MetricErrors,
					Time:   time.Now(),
					Tags:   map[string]string{"vu": idString, "error": err.Error()},
					Value:  float64(1),
				})
			}
		}

		vu.Buffer = append(vu.Buffer, samples...)
		if vu.ExtBuffer != nil {
			vu.ExtBuffer.Add(samples...)
		}

		if !e.Status.Running.Bool {
			goto waitForPause
		}
	}
}

// Returns a value sink for a metric, created from the type if unavailable.
func (e *Engine) getSink(m *stats.Metric) stats.Sink {
	s, ok := e.Metrics[m]
	if !ok {
		switch m.Type {
		case stats.Counter:
			s = &stats.CounterSink{}
		case stats.Gauge:
			s = &stats.GaugeSink{}
		case stats.Trend:
			s = &stats.TrendSink{}
		}
		e.Metrics[m] = s
	}
	return s
}

func (e *Engine) consumeBuffer(buffer []stats.Sample) {
	for _, sample := range buffer {
		e.getSink(sample.Metric).Add(sample)
	}
}

func (e *Engine) consumeEngineStats() {
	t := time.Now()
	e.consumeBuffer([]stats.Sample{
		stats.Sample{Metric: MetricVUs, Time: t, Value: float64(e.Status.VUs.Int64)},
		stats.Sample{Metric: MetricVUsMax, Time: t, Value: float64(e.Status.VUsMax.Int64)},
	})
}
<commit_msg>[fix] Whoops I simplified away external outputs<commit_after>package lib

import (
	"context"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/stats"
	"gopkg.in/guregu/null.v3"
	"strconv"
	"sync"
	"time"
)

const TickRate = 1 * time.Millisecond

var (
	MetricVUs    = &stats.Metric{Name: "vus", Type: stats.Gauge}
	MetricVUsMax = &stats.Metric{Name: "vus_max", Type: stats.Gauge}
	MetricErrors = &stats.Metric{Name: "errors", Type: stats.Counter}

	ErrTooManyVUs = errors.New("More VUs than the maximum requested")
	ErrMaxTooLow  = errors.New("Can't lower max below current VU count")

	// Special error used to taint a test, without printing an error.
	ErrVUWantsTaint = errors.New("[ErrVUWantsTaint is never logged]")
)

type vuEntry struct {
	VU        VU
	Buffer    []stats.Sample
	ExtBuffer stats.Buffer
	Cancel    context.CancelFunc
}

type Engine struct {
	Runner    Runner
	Status    Status
	Stages    []Stage
	Collector stats.Collector
	Quit      bool
	Pause     sync.WaitGroup

	Metrics map[*stats.Metric]stats.Sink

	ctx    context.Context
	vus    []*vuEntry
	nextID int64

	vuMutex sync.Mutex
}

func NewEngine(r Runner) (*Engine, error) {
	e := &Engine{
		Runner: r,
		Status: Status{
			Running: null.BoolFrom(false),
			Tainted: null.BoolFrom(false),
			VUs:     null.IntFrom(0),
			VUsMax:  null.IntFrom(0),
			AtTime:  null.IntFrom(0),
		},
		Metrics: make(map[*stats.Metric]stats.Sink),
	}

	e.Status.Running = null.BoolFrom(false)
	e.Pause.Add(1)

	e.Status.VUs = null.IntFrom(0)
	e.Status.VUsMax = null.IntFrom(0)

	return e, nil
}

func (e *Engine) Run(ctx context.Context) error {
	if len(e.Stages) == 0 {
		return errors.New("Engine has no stages")
	}

	e.ctx = ctx
	e.nextID = 1

	if e.Collector != nil {
		go e.Collector.Run(ctx)
	} else {
		log.Debug("Engine: No Collector")
	}

	e.consumeEngineStats()

	ticker := time.NewTicker(TickRate)
	lastTick := time.Now()

loop:
	for {
		select {
		case now := <-ticker.C:
			timeDelta := now.Sub(lastTick)
			e.Status.AtTime.Int64 += int64(timeDelta)
			lastTick = now

			stage, left, ok := StageAt(e.Stages, time.Duration(e.Status.AtTime.Int64))
			if stage.StartVUs.Valid && stage.EndVUs.Valid {
				progress := (float64(stage.Duration.Int64-int64(left)) / float64(stage.Duration.Int64))
				vus := Lerp(stage.StartVUs.Int64, stage.EndVUs.Int64, progress)
				e.SetVUs(vus)
			}

			for _, vu := range e.vus {
				e.consumeBuffer(vu.Buffer)
			}

			if !ok {
				e.SetRunning(false)

				if e.Quit {
					break loop
				} else {
					log.Info("Test finished, press Ctrl+C to exit")
				}
			}

			e.consumeEngineStats()
		case <-ctx.Done():
			break loop
		}
	}

	e.vus = nil

	e.Status.Running = null.BoolFrom(false)
	e.Status.VUs = null.IntFrom(0)
	e.Status.VUsMax = null.IntFrom(0)
	e.consumeEngineStats()

	return nil
}

func (e *Engine) SetRunning(running bool) {
	if running && !e.Status.Running.Bool {
		e.Pause.Done()
		log.Debug("Engine Unpaused")
	} else if !running && e.Status.Running.Bool {
		e.Pause.Add(1)
		log.Debug("Engine Paused")
	}
	e.Status.Running.Bool = running
}

func (e *Engine) SetVUs(v int64) error {
	if e.Status.VUs.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUs.Int64, "to": v}).Debug("Setting VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v > e.Status.VUsMax.Int64 {
		return ErrTooManyVUs
	}

	current := e.Status.VUs.Int64
	for i := current; i < v; i++ {
		entry := e.vus[i]
		if entry.Cancel != nil {
			panic(errors.New("ATTEMPTED TO RESCHEDULE RUNNING VU"))
		}

		ctx, cancel := context.WithCancel(e.ctx)
		entry.Cancel = cancel

		if err := entry.VU.Reconfigure(e.nextID); err != nil {
			return err
		}
		go e.runVU(ctx, e.nextID, entry)
		e.nextID++
	}
	for i := current - 1; i >= v; i-- {
		entry := e.vus[i]
		entry.Cancel()
		entry.Cancel = nil
	}

	e.Status.VUs.Int64 = v
	return nil
}

func (e *Engine) SetMaxVUs(v int64) error {
	if e.Status.VUsMax.Int64 == v {
		return nil
	}

	log.WithFields(log.Fields{"from": e.Status.VUsMax.Int64, "to": v}).Debug("Setting Max VUs")

	e.vuMutex.Lock()
	defer e.vuMutex.Unlock()

	if v < e.Status.VUs.Int64 {
		return ErrMaxTooLow
	}

	current := e.Status.VUsMax.Int64
	if v > current {
		vus := e.vus
		for i := current; i < v; i++ {
			vu, err := e.Runner.NewVU()
			if err != nil {
				return err
			}
			entry := &vuEntry{VU: vu}
			if e.Collector != nil {
				entry.ExtBuffer = e.Collector.Buffer()
			}
			vus = append(vus, entry)
		}
		e.vus = vus
	} else if v < current {
		e.vus = e.vus[:v]
	}

	e.Status.VUsMax.Int64 = v
	return nil
}

func (e *Engine) runVU(ctx context.Context, id int64, vu *vuEntry) {
	idString := strconv.FormatInt(id, 10)

waitForPause:
	e.Pause.Wait()

	for {
		samples, err := vu.VU.RunOnce(ctx)

		// If the context is cancelled, the iteration is likely to produce erroneous output
		// due to cancelled HTTP requests and whatnot. Discard output from such runs.
		select {
		case <-ctx.Done():
			return
		default:
		}

		if err != nil {
			e.Status.Tainted.Bool = true

			if err != ErrVUWantsTaint {
				if s, ok := err.(fmt.Stringer); ok {
					log.Error(s.String())
				} else {
					log.WithError(err).Error("Runtime Error")
				}

				samples = append(samples, stats.Sample{
					Metric: MetricErrors,
					Time:   time.Now(),
					Tags:   map[string]string{"vu": idString, "error": err.Error()},
					Value:  float64(1),
				})
			}
		}

		vu.Buffer = append(vu.Buffer, samples...)
		if vu.ExtBuffer != nil {
			vu.ExtBuffer.Add(samples...)
		}

		if !e.Status.Running.Bool {
			goto waitForPause
		}
	}
}

// Returns a value sink for a metric, created from the type if unavailable.
func (e *Engine) getSink(m *stats.Metric) stats.Sink {
	s, ok := e.Metrics[m]
	if !ok {
		switch m.Type {
		case stats.Counter:
			s = &stats.CounterSink{}
		case stats.Gauge:
			s = &stats.GaugeSink{}
		case stats.Trend:
			s = &stats.TrendSink{}
		}
		e.Metrics[m] = s
	}
	return s
}

func (e *Engine) consumeBuffer(buffer []stats.Sample) {
	for _, sample := range buffer {
		e.getSink(sample.Metric).Add(sample)
	}
}

func (e *Engine) consumeEngineStats() {
	t := time.Now()
	e.consumeBuffer([]stats.Sample{
		stats.Sample{Metric: MetricVUs, Time: t, Value: float64(e.Status.VUs.Int64)},
		stats.Sample{Metric: MetricVUsMax, Time: t, Value: float64(e.Status.VUsMax.Int64)},
	})
}
<|endoftext|>


<commit_before>// Timezoneless date/datetime types

package types

import (
	"database/sql/driver"
	"time"
)

// Truncate timezone
//
//   clickhouse.Date(time.Date(2017, 1, 1, 0, 0, 0, 0, time.Local)) -> time.Date(2017, 1, 1, 0, 0, 0, 0, time.UTC)
type Date time.Time

func (date Date) Value() (driver.Value, error) {
	return date.convert(), nil
}

func (date Date) convert() time.Time {
	return time.Date(time.Time(date).Year(), time.Time(date).Month(), time.Time(date).Day(), 0, 0, 0, 0, time.UTC)
}

// Truncate timezone
//
//   clickhouse.DateTime(time.Date(2017, 1, 1, 0, 0, 0, 0, time.Local)) -> time.Date(2017, 1, 1, 0, 0, 0, 0, time.UTC)
type DateTime time.Time

func (datetime DateTime) Value() (driver.Value, error) {
	return datetime.convert(), nil
}

func (datetime DateTime) convert() time.Time {
	return time.Date(
		time.Time(datetime).Year(),
		time.Time(datetime).Month(),
		time.Time(datetime).Day(),
		time.Time(datetime).Hour(),
		time.Time(datetime).Minute(),
		time.Time(datetime).Second(),
		0,
		time.UTC,
	)
}

var (
	_ driver.Valuer = Date{}
	_ driver.Valuer = DateTime{}
)
<commit_msg>#150 DateTime converter: add nanosecond<commit_after>// Timezoneless date/datetime types

package types

import (
	"database/sql/driver"
	"time"
)

// Truncate timezone
//
//   clickhouse.Date(time.Date(2017, 1, 1, 0, 0, 0, 0, time.Local)) -> time.Date(2017, 1, 1, 0, 0, 0, 0, time.UTC)
type Date time.Time

func (date Date) Value() (driver.Value, error) {
	return date.convert(), nil
}

func (date Date) convert() time.Time {
	return time.Date(time.Time(date).Year(), time.Time(date).Month(), time.Time(date).Day(), 0, 0, 0, 0, time.UTC)
}

// Truncate timezone
//
//   clickhouse.DateTime(time.Date(2017, 1, 1, 0, 0, 0, 0, time.Local)) -> time.Date(2017, 1, 1, 0, 0, 0, 0, time.UTC)
type DateTime time.Time

func (datetime DateTime) Value() (driver.Value, error) {
	return datetime.convert(), nil
}

func (datetime DateTime) convert() time.Time {
	return time.Date(
		time.Time(datetime).Year(),
		time.Time(datetime).Month(),
		time.Time(datetime).Day(),
		time.Time(datetime).Hour(),
		time.Time(datetime).Minute(),
		time.Time(datetime).Second(),
		1,
		time.UTC,
	)
}

var (
	_ driver.Valuer = Date{}
	_ driver.Valuer = DateTime{}
)
<|endoftext|>


<commit_before>package formatter

import (
  "github.com/phemmer/sawmill/event"
  "golang.org/x/crypto/ssh/terminal"
  "fmt"
  "reflect"
  "sort"
)

func IsTerminal(stream interface{Fd() uintptr}) bool {
  return terminal.IsTerminal(int(stream.Fd()))
}
var colors = terminal.EscapeCodes{
	Black:   []byte{27, '[', '3', '0', 'm'},
	Red:     []byte{27, '[', '3', '1', 'm'},
	Green:   []byte{27, '[', '3', '2', 'm'},
	Yellow:  []byte{27, '[', '3', '3', 'm'},
	Blue:    []byte{27, '[', '3', '4', 'm'},
	Magenta: []byte{27, '[', '3', '5', 'm'},
	Cyan:    []byte{27, '[', '3', '6', 'm'},
	White:   []byte{27, '[', '3', '7', 'm'},

	Reset:   []byte{27, '[', '0', 'm'},
}

type TextFormatter struct {
  UseColor bool
	TimeFormat string
}
func NewTextFormatter(color bool) *TextFormatter {
  return &TextFormatter{
		UseColor: color,
		TimeFormat: "2006-01-02_15:04:05.000",
	}
}
func (formatter *TextFormatter) Format(logEvent *event.Event) ([]byte) {
  fields := flatten(logEvent.Fields)

	timestamp := formatter.FormatTimestamp(logEvent)
	levelName := formatter.FormatLevel(logEvent)
  buf := []byte(fmt.Sprintf("%s %s> %s ", timestamp, levelName, logEvent.Message))

  flatFields := flatten(fields)
  keys := make([]string, len(flatFields))
  i := 0
  for k := range flatFields {
    keys[i] = k
    i += 1
  }
  sort.Strings(keys)
  for _, k := range keys {
    v := flatFields[k]
    buf = append(buf, []byte(fmt.Sprintf("%s=%v ", k, v))...)
  }
  return buf
}

func (formatter *TextFormatter) FormatTimestamp(logEvent *event.Event) string {
	return logEvent.Timestamp.Format(formatter.TimeFormat)
}
func (formatter *TextFormatter) FormatLevel(logEvent *event.Event) string {
	var levelName string
	if formatter.UseColor {
		var levelColor []byte
		if logEvent.Level <= event.Error {
			levelColor = colors.Red
		} else if logEvent.Level == event.Warning {
			levelColor = colors.Yellow
		} else {
			levelColor = colors.Blue
		}
		levelName = fmt.Sprintf("%s%s%s", levelColor, logEvent.LevelName(), colors.Reset)
	} else {
		levelName = logEvent.LevelName()
	}

	return levelName
}

func flatten(fields interface{}) (map[string]interface{}) {
  flat := make(map[string]interface{})

  value := reflect.ValueOf(fields)
  for value.Kind() == reflect.Ptr { // shouldn't ever happen since deStruct() also does this
    value = value.Elem()
  }

  //fmt.Printf("flattening: %#v\n", fields)
  if value.Kind() == reflect.Map {
    for _, kV := range value.MapKeys() {
      vV := value.MapIndex(kV)
      k := kV.Interface()
      flattenValue(flat, k, vV)
    }
  } else if value.Kind() == reflect.Array || value.Kind() == reflect.Slice {
    for k, v := range value.Interface().([]interface{}) {
      flattenValue(flat, k, reflect.ValueOf(v))
    }
  } else {
    if value.IsValid() {
      flat[""] = value.Interface()
    } else {
      flat[""] = nil
    }
  }

  //fmt.Printf("Flatten result: %#v\n", flat)
  return flat
}
func flattenValue(flattened map[string]interface{}, parentKey interface{}, value reflect.Value) {
  kind := value.Kind()
  for kind == reflect.Ptr || kind == reflect.Interface {
    kind = value.Elem().Kind()
  }

  if kind == reflect.Struct || kind == reflect.Map || kind == reflect.Array || kind == reflect.Slice {
    for vk,vv := range flatten(value.Interface()) {
      flat_k := fmt.Sprintf("%s.%s", parentKey, vk)
      flattened[flat_k] = vv
    }
  } else {
    flat_k := fmt.Sprintf("%s", parentKey)
    flattened[flat_k] = value.Interface()
  }
}
<commit_msg>split message & fields formatters to their own functions<commit_after>package formatter

import (
  "github.com/phemmer/sawmill/event"
  "golang.org/x/crypto/ssh/terminal"
  "fmt"
  "reflect"
  "sort"
)

func IsTerminal(stream interface{Fd() uintptr}) bool {
  return terminal.IsTerminal(int(stream.Fd()))
}
var colors = terminal.EscapeCodes{
	Black:   []byte{27, '[', '3', '0', 'm'},
	Red:     []byte{27, '[', '3', '1', 'm'},
	Green:   []byte{27, '[', '3', '2', 'm'},
	Yellow:  []byte{27, '[', '3', '3', 'm'},
	Blue:    []byte{27, '[', '3', '4', 'm'},
	Magenta: []byte{27, '[', '3', '5', 'm'},
	Cyan:    []byte{27, '[', '3', '6', 'm'},
	White:   []byte{27, '[', '3', '7', 'm'},

	Reset:   []byte{27, '[', '0', 'm'},
}

type TextFormatter struct {
  UseColor bool
	TimeFormat string
}
func NewTextFormatter(color bool) *TextFormatter {
  return &TextFormatter{
		UseColor: color,
		TimeFormat: "2006-01-02_15:04:05.000",
	}
}
func (formatter *TextFormatter) Format(logEvent *event.Event) ([]byte) {
	timestamp := formatter.FormatTimestamp(logEvent)
	level := formatter.FormatLevel(logEvent)
	message := formatter.FormatMessage(logEvent)
	fields := formatter.FormatFields(logEvent)

  buf := []byte(fmt.Sprintf("%s %s> %s %s", timestamp, level, message, fields))

	return buf
}

func (formatter *TextFormatter) FormatTimestamp(logEvent *event.Event) string {
	return logEvent.Timestamp.Format(formatter.TimeFormat)
}
func (formatter *TextFormatter) FormatLevel(logEvent *event.Event) string {
	var levelName string
	if formatter.UseColor {
		var levelColor []byte
		if logEvent.Level <= event.Error {
			levelColor = colors.Red
		} else if logEvent.Level == event.Warning {
			levelColor = colors.Yellow
		} else {
			levelColor = colors.Blue
		}
		levelName = fmt.Sprintf("%s%s%s", levelColor, logEvent.LevelName(), colors.Reset)
	} else {
		levelName = logEvent.LevelName()
	}

	return levelName
}
func (formatter *TextFormatter) FormatMessage(logEvent *event.Event) string {
	return logEvent.Message
}
func (formatter *TextFormatter) FormatFields(logEvent *event.Event) string {
  flatFields := flatten(logEvent.Fields)

  keys := make([]string, len(flatFields))
  i := 0
  for k := range flatFields {
    keys[i] = k
    i += 1
  }
  sort.Strings(keys)

	buf := []byte{}

  for _, k := range keys {
		if len(buf) > 0 {
			buf = append(buf, ' ')
		}
    v := flatFields[k]
    buf = append(buf, []byte(fmt.Sprintf("%s=%v", k, v))...)
  }

  return string(buf)
}

func flatten(fields interface{}) (map[string]interface{}) {
  flat := make(map[string]interface{})

  value := reflect.ValueOf(fields)
  for value.Kind() == reflect.Ptr { // shouldn't ever happen since deStruct() also does this
    value = value.Elem()
  }

  //fmt.Printf("flattening: %#v\n", fields)
  if value.Kind() == reflect.Map {
    for _, kV := range value.MapKeys() {
      vV := value.MapIndex(kV)
      k := kV.Interface()
      flattenValue(flat, k, vV)
    }
  } else if value.Kind() == reflect.Array || value.Kind() == reflect.Slice {
    for k, v := range value.Interface().([]interface{}) {
      flattenValue(flat, k, reflect.ValueOf(v))
    }
  } else {
    if value.IsValid() {
      flat[""] = value.Interface()
    } else {
      flat[""] = nil
    }
  }

  //fmt.Printf("Flatten result: %#v\n", flat)
  return flat
}
func flattenValue(flattened map[string]interface{}, parentKey interface{}, value reflect.Value) {
  kind := value.Kind()
  for kind == reflect.Ptr || kind == reflect.Interface {
    kind = value.Elem().Kind()
  }

  if kind == reflect.Struct || kind == reflect.Map || kind == reflect.Array || kind == reflect.Slice {
    for vk,vv := range flatten(value.Interface()) {
      flat_k := fmt.Sprintf("%s.%s", parentKey, vk)
      flattened[flat_k] = vv
    }
  } else {
    flat_k := fmt.Sprintf("%s", parentKey)
    flattened[flat_k] = value.Interface()
  }
}
<|endoftext|>