<commit_before>// Copyright 2014 The Prometheus 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 prometheus

import (
	"math"
	"sync/atomic"
	"time"

	dto "github.com/prometheus/client_model/go"
)

// Gauge is a Metric that represents a single numerical value that can
// arbitrarily go up and down.
//
// A Gauge is typically used for measured values like temperatures or current
// memory usage, but also "counts" that can go up and down, like the number of
// running goroutines.
//
// To create Gauge instances, use NewGauge.
type Gauge interface {
	Metric
	Collector

	// Set sets the Gauge to an arbitrary value.
	Set(float64)
	// Inc increments the Gauge by 1. Use Add to increment it by arbitrary
	// values.
	Inc()
	// Dec decrements the Gauge by 1. Use Sub to decrement it by arbitrary
	// values.
	Dec()
	// Add adds the given value to the Gauge. (The value can be negative,
	// resulting in a decrease of the Gauge.)
	Add(float64)
	// Sub subtracts the given value from the Gauge. (The value can be
	// negative, resulting in an increase of the Gauge.)
	Sub(float64)

	// SetToCurrentTime sets the Gauge to the current Unix time in seconds.
	SetToCurrentTime()
}

// GaugeOpts is an alias for Opts. See there for doc comments.
type GaugeOpts Opts

// NewGauge creates a new Gauge based on the provided GaugeOpts.
//
// The returned implementation is optimized for a fast Set method. If you have a
// choice for managing the value of a Gauge via Set vs. Inc/Dec/Add/Sub, pick
// the former. For example, the Inc method of the returned Gauge is slower than
// the Inc method of a Counter returned by NewCounter. This matches the typical
// scenarios for Gauges and Counters, where the former tends to be Set-heavy and
// the latter Inc-heavy.
func NewGauge(opts GaugeOpts) Gauge {
	desc := NewDesc(
		BuildFQName(opts.Namespace, opts.Subsystem, opts.Name),
		opts.Help,
		nil,
		opts.ConstLabels,
	)
	result := &gauge{desc: desc, labelPairs: desc.constLabelPairs}
	result.init(result) // Init self-collection.
	return result
}

type gauge struct {
	// valBits contains the bits of the represented float64 value. It has
	// to go first in the struct to guarantee alignment for atomic
	// operations.  http://golang.org/pkg/sync/atomic/#pkg-note-BUG
	valBits uint64

	selfCollector

	desc       *Desc
	labelPairs []*dto.LabelPair
}

func (g *gauge) Desc() *Desc {
	return g.desc
}

func (g *gauge) Set(val float64) {
	atomic.StoreUint64(&g.valBits, math.Float64bits(val))
}

func (g *gauge) SetToCurrentTime() {
	g.Set(float64(time.Now().UnixNano()) / 1e9)
}

func (g *gauge) Inc() {
	g.Add(1)
}

func (g *gauge) Dec() {
	g.Add(-1)
}

func (g *gauge) Add(val float64) {
	for {
		oldBits := atomic.LoadUint64(&g.valBits)
		newBits := math.Float64bits(math.Float64frombits(oldBits) + val)
		if atomic.CompareAndSwapUint64(&g.valBits, oldBits, newBits) {
			return
		}
	}
}

func (g *gauge) Sub(val float64) {
	g.Add(val * -1)
}

func (g *gauge) Write(out *dto.Metric) error {
	val := math.Float64frombits(atomic.LoadUint64(&g.valBits))
	return populateMetric(GaugeValue, val, g.labelPairs, out)
}

// GaugeVec is a Collector that bundles a set of Gauges that all share the same
// Desc, but have different values for their variable labels. This is used if
// you want to count the same thing partitioned by various dimensions
// (e.g. number of operations queued, partitioned by user and operation
// type). Create instances with NewGaugeVec.
type GaugeVec struct {
	*metricVec
}

// NewGaugeVec creates a new GaugeVec based on the provided GaugeOpts and
// partitioned by the given label names.
func NewGaugeVec(opts GaugeOpts, labelNames []string) *GaugeVec {
	desc := NewDesc(
		BuildFQName(opts.Namespace, opts.Subsystem, opts.Name),
		opts.Help,
		labelNames,
		opts.ConstLabels,
	)
	return &GaugeVec{
		metricVec: newMetricVec(desc, func(lvs ...string) Metric {
			if len(lvs) != len(desc.variableLabels) {
				panic(makeInconsistentCardinalityError(desc.fqName, desc.variableLabels, lvs))
			}
			result := &gauge{desc: desc, labelPairs: makeLabelPairs(desc, lvs)}
			result.init(result) // Init self-collection.
			return result
		}),
	}
}

// GetMetricWithLabelValues returns the Gauge for the given slice of label
// values (same order as the VariableLabels in Desc). If that combination of
// label values is accessed for the first time, a new Gauge is created.
//
// It is possible to call this method without using the returned Gauge to only
// create the new Gauge but leave it at its starting value 0. See also the
// SummaryVec example.
//
// Keeping the Gauge for later use is possible (and should be considered if
// performance is critical), but keep in mind that Reset, DeleteLabelValues and
// Delete can be used to delete the Gauge from the GaugeVec. In that case, the
// Gauge will still exist, but it will not be exported anymore, even if a
// Gauge with the same label values is created later. See also the CounterVec
// example.
//
// An error is returned if the number of label values is not the same as the
// number of VariableLabels in Desc (minus any curried labels).
//
// Note that for more than one label value, this method is prone to mistakes
// caused by an incorrect order of arguments. Consider GetMetricWith(Labels) as
// an alternative to avoid that type of mistake. For higher label numbers, the
// latter has a much more readable (albeit more verbose) syntax, but it comes
// with a performance overhead (for creating and processing the Labels map).
func (v *GaugeVec) GetMetricWithLabelValues(lvs ...string) (Gauge, error) {
	metric, err := v.metricVec.getMetricWithLabelValues(lvs...)
	if metric != nil {
		return metric.(Gauge), err
	}
	return nil, err
}

// GetMetricWith returns the Gauge for the given Labels map (the label names
// must match those of the VariableLabels in Desc). If that label map is
// accessed for the first time, a new Gauge is created. Implications of
// creating a Gauge without using it and keeping the Gauge for later use are
// the same as for GetMetricWithLabelValues.
//
// An error is returned if the number and names of the Labels are inconsistent
// with those of the VariableLabels in Desc (minus any curried labels).
//
// This method is used for the same purpose as
// GetMetricWithLabelValues(...string). See there for pros and cons of the two
// methods.
func (v *GaugeVec) GetMetricWith(labels Labels) (Gauge, error) {
	metric, err := v.metricVec.getMetricWith(labels)
	if metric != nil {
		return metric.(Gauge), err
	}
	return nil, err
}

// WithLabelValues works as GetMetricWithLabelValues, but panics where
// GetMetricWithLabelValues would have returned an error. Not returning an
// error allows shortcuts like
//     myVec.WithLabelValues("404", "GET").Add(42)
func (v *GaugeVec) WithLabelValues(lvs ...string) Gauge {
	g, err := v.GetMetricWithLabelValues(lvs...)
	if err != nil {
		panic(err)
	}
	return g
}

// With works as GetMetricWith, but panics where GetMetricWithLabels would have
// returned an error. Not returning an error allows shortcuts like
//     myVec.With(prometheus.Labels{"code": "404", "method": "GET"}).Add(42)
func (v *GaugeVec) With(labels Labels) Gauge {
	g, err := v.GetMetricWith(labels)
	if err != nil {
		panic(err)
	}
	return g
}

// CurryWith returns a vector curried with the provided labels, i.e. the
// returned vector has those labels pre-set for all labeled operations performed
// on it. The cardinality of the curried vector is reduced accordingly. The
// order of the remaining labels stays the same (just with the curried labels
// taken out of the sequence – which is relevant for the
// (GetMetric)WithLabelValues methods). It is possible to curry a curried
// vector, but only with labels not yet used for currying before.
//
// The metrics contained in the GaugeVec are shared between the curried and
// uncurried vectors. They are just accessed differently. Curried and uncurried
// vectors behave identically in terms of collection. Only one must be
// registered with a given registry (usually the uncurried version). The Reset
// method deletes all metrics, even if called on a curried vector.
func (v *GaugeVec) CurryWith(labels Labels) (*GaugeVec, error) {
	vec, err := v.curryWith(labels)
	if vec != nil {
		return &GaugeVec{vec}, err
	}
	return nil, err
}

// MustCurryWith works as CurryWith but panics where CurryWith would have
// returned an error.
func (v *GaugeVec) MustCurryWith(labels Labels) *GaugeVec {
	vec, err := v.CurryWith(labels)
	if err != nil {
		panic(err)
	}
	return vec
}

// GaugeFunc is a Gauge whose value is determined at collect time by calling a
// provided function.
//
// To create GaugeFunc instances, use NewGaugeFunc.
type GaugeFunc interface {
	Metric
	Collector
}

// NewGaugeFunc creates a new GaugeFunc based on the provided GaugeOpts. The
// value reported is determined by calling the given function from within the
// Write method. Take into account that metric collection may happen
// concurrently. If that results in concurrent calls to Write, like in the case
// where a GaugeFunc is directly registered with Prometheus, the provided
// function must be concurrency-safe.
func NewGaugeFunc(opts GaugeOpts, function func() float64) GaugeFunc {
	return newValueFunc(NewDesc(
		BuildFQName(opts.Namespace, opts.Subsystem, opts.Name),
		opts.Help,
		nil,
		opts.ConstLabels,
	), GaugeValue, function)
}
<commit_msg>Improve doc comment for NewGaugeFunc<commit_after>// Copyright 2014 The Prometheus 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 prometheus

import (
	"math"
	"sync/atomic"
	"time"

	dto "github.com/prometheus/client_model/go"
)

// Gauge is a Metric that represents a single numerical value that can
// arbitrarily go up and down.
//
// A Gauge is typically used for measured values like temperatures or current
// memory usage, but also "counts" that can go up and down, like the number of
// running goroutines.
//
// To create Gauge instances, use NewGauge.
type Gauge interface {
	Metric
	Collector

	// Set sets the Gauge to an arbitrary value.
	Set(float64)
	// Inc increments the Gauge by 1. Use Add to increment it by arbitrary
	// values.
	Inc()
	// Dec decrements the Gauge by 1. Use Sub to decrement it by arbitrary
	// values.
	Dec()
	// Add adds the given value to the Gauge. (The value can be negative,
	// resulting in a decrease of the Gauge.)
	Add(float64)
	// Sub subtracts the given value from the Gauge. (The value can be
	// negative, resulting in an increase of the Gauge.)
	Sub(float64)

	// SetToCurrentTime sets the Gauge to the current Unix time in seconds.
	SetToCurrentTime()
}

// GaugeOpts is an alias for Opts. See there for doc comments.
type GaugeOpts Opts

// NewGauge creates a new Gauge based on the provided GaugeOpts.
//
// The returned implementation is optimized for a fast Set method. If you have a
// choice for managing the value of a Gauge via Set vs. Inc/Dec/Add/Sub, pick
// the former. For example, the Inc method of the returned Gauge is slower than
// the Inc method of a Counter returned by NewCounter. This matches the typical
// scenarios for Gauges and Counters, where the former tends to be Set-heavy and
// the latter Inc-heavy.
func NewGauge(opts GaugeOpts) Gauge {
	desc := NewDesc(
		BuildFQName(opts.Namespace, opts.Subsystem, opts.Name),
		opts.Help,
		nil,
		opts.ConstLabels,
	)
	result := &gauge{desc: desc, labelPairs: desc.constLabelPairs}
	result.init(result) // Init self-collection.
	return result
}

type gauge struct {
	// valBits contains the bits of the represented float64 value. It has
	// to go first in the struct to guarantee alignment for atomic
	// operations.  http://golang.org/pkg/sync/atomic/#pkg-note-BUG
	valBits uint64

	selfCollector

	desc       *Desc
	labelPairs []*dto.LabelPair
}

func (g *gauge) Desc() *Desc {
	return g.desc
}

func (g *gauge) Set(val float64) {
	atomic.StoreUint64(&g.valBits, math.Float64bits(val))
}

func (g *gauge) SetToCurrentTime() {
	g.Set(float64(time.Now().UnixNano()) / 1e9)
}

func (g *gauge) Inc() {
	g.Add(1)
}

func (g *gauge) Dec() {
	g.Add(-1)
}

func (g *gauge) Add(val float64) {
	for {
		oldBits := atomic.LoadUint64(&g.valBits)
		newBits := math.Float64bits(math.Float64frombits(oldBits) + val)
		if atomic.CompareAndSwapUint64(&g.valBits, oldBits, newBits) {
			return
		}
	}
}

func (g *gauge) Sub(val float64) {
	g.Add(val * -1)
}

func (g *gauge) Write(out *dto.Metric) error {
	val := math.Float64frombits(atomic.LoadUint64(&g.valBits))
	return populateMetric(GaugeValue, val, g.labelPairs, out)
}

// GaugeVec is a Collector that bundles a set of Gauges that all share the same
// Desc, but have different values for their variable labels. This is used if
// you want to count the same thing partitioned by various dimensions
// (e.g. number of operations queued, partitioned by user and operation
// type). Create instances with NewGaugeVec.
type GaugeVec struct {
	*metricVec
}

// NewGaugeVec creates a new GaugeVec based on the provided GaugeOpts and
// partitioned by the given label names.
func NewGaugeVec(opts GaugeOpts, labelNames []string) *GaugeVec {
	desc := NewDesc(
		BuildFQName(opts.Namespace, opts.Subsystem, opts.Name),
		opts.Help,
		labelNames,
		opts.ConstLabels,
	)
	return &GaugeVec{
		metricVec: newMetricVec(desc, func(lvs ...string) Metric {
			if len(lvs) != len(desc.variableLabels) {
				panic(makeInconsistentCardinalityError(desc.fqName, desc.variableLabels, lvs))
			}
			result := &gauge{desc: desc, labelPairs: makeLabelPairs(desc, lvs)}
			result.init(result) // Init self-collection.
			return result
		}),
	}
}

// GetMetricWithLabelValues returns the Gauge for the given slice of label
// values (same order as the VariableLabels in Desc). If that combination of
// label values is accessed for the first time, a new Gauge is created.
//
// It is possible to call this method without using the returned Gauge to only
// create the new Gauge but leave it at its starting value 0. See also the
// SummaryVec example.
//
// Keeping the Gauge for later use is possible (and should be considered if
// performance is critical), but keep in mind that Reset, DeleteLabelValues and
// Delete can be used to delete the Gauge from the GaugeVec. In that case, the
// Gauge will still exist, but it will not be exported anymore, even if a
// Gauge with the same label values is created later. See also the CounterVec
// example.
//
// An error is returned if the number of label values is not the same as the
// number of VariableLabels in Desc (minus any curried labels).
//
// Note that for more than one label value, this method is prone to mistakes
// caused by an incorrect order of arguments. Consider GetMetricWith(Labels) as
// an alternative to avoid that type of mistake. For higher label numbers, the
// latter has a much more readable (albeit more verbose) syntax, but it comes
// with a performance overhead (for creating and processing the Labels map).
func (v *GaugeVec) GetMetricWithLabelValues(lvs ...string) (Gauge, error) {
	metric, err := v.metricVec.getMetricWithLabelValues(lvs...)
	if metric != nil {
		return metric.(Gauge), err
	}
	return nil, err
}

// GetMetricWith returns the Gauge for the given Labels map (the label names
// must match those of the VariableLabels in Desc). If that label map is
// accessed for the first time, a new Gauge is created. Implications of
// creating a Gauge without using it and keeping the Gauge for later use are
// the same as for GetMetricWithLabelValues.
//
// An error is returned if the number and names of the Labels are inconsistent
// with those of the VariableLabels in Desc (minus any curried labels).
//
// This method is used for the same purpose as
// GetMetricWithLabelValues(...string). See there for pros and cons of the two
// methods.
func (v *GaugeVec) GetMetricWith(labels Labels) (Gauge, error) {
	metric, err := v.metricVec.getMetricWith(labels)
	if metric != nil {
		return metric.(Gauge), err
	}
	return nil, err
}

// WithLabelValues works as GetMetricWithLabelValues, but panics where
// GetMetricWithLabelValues would have returned an error. Not returning an
// error allows shortcuts like
//     myVec.WithLabelValues("404", "GET").Add(42)
func (v *GaugeVec) WithLabelValues(lvs ...string) Gauge {
	g, err := v.GetMetricWithLabelValues(lvs...)
	if err != nil {
		panic(err)
	}
	return g
}

// With works as GetMetricWith, but panics where GetMetricWithLabels would have
// returned an error. Not returning an error allows shortcuts like
//     myVec.With(prometheus.Labels{"code": "404", "method": "GET"}).Add(42)
func (v *GaugeVec) With(labels Labels) Gauge {
	g, err := v.GetMetricWith(labels)
	if err != nil {
		panic(err)
	}
	return g
}

// CurryWith returns a vector curried with the provided labels, i.e. the
// returned vector has those labels pre-set for all labeled operations performed
// on it. The cardinality of the curried vector is reduced accordingly. The
// order of the remaining labels stays the same (just with the curried labels
// taken out of the sequence – which is relevant for the
// (GetMetric)WithLabelValues methods). It is possible to curry a curried
// vector, but only with labels not yet used for currying before.
//
// The metrics contained in the GaugeVec are shared between the curried and
// uncurried vectors. They are just accessed differently. Curried and uncurried
// vectors behave identically in terms of collection. Only one must be
// registered with a given registry (usually the uncurried version). The Reset
// method deletes all metrics, even if called on a curried vector.
func (v *GaugeVec) CurryWith(labels Labels) (*GaugeVec, error) {
	vec, err := v.curryWith(labels)
	if vec != nil {
		return &GaugeVec{vec}, err
	}
	return nil, err
}

// MustCurryWith works as CurryWith but panics where CurryWith would have
// returned an error.
func (v *GaugeVec) MustCurryWith(labels Labels) *GaugeVec {
	vec, err := v.CurryWith(labels)
	if err != nil {
		panic(err)
	}
	return vec
}

// GaugeFunc is a Gauge whose value is determined at collect time by calling a
// provided function.
//
// To create GaugeFunc instances, use NewGaugeFunc.
type GaugeFunc interface {
	Metric
	Collector
}

// NewGaugeFunc creates a new GaugeFunc based on the provided GaugeOpts. The
// value reported is determined by calling the given function from within the
// Write method. Take into account that metric collection may happen
// concurrently. Therefore, it must be safe to call the provided function
// concurrently.
//
// NewGaugeFunc is a good way to create an “info” style metric with a constant
// value of 1. Example:
// https://github.com/prometheus/common/blob/8558a5b7db3c84fa38b4766966059a7bd5bfa2ee/version/info.go#L36-L56
func NewGaugeFunc(opts GaugeOpts, function func() float64) GaugeFunc {
	return newValueFunc(NewDesc(
		BuildFQName(opts.Namespace, opts.Subsystem, opts.Name),
		opts.Help,
		nil,
		opts.ConstLabels,
	), GaugeValue, function)
}
<|endoftext|>


<commit_before>package proto3_test

import (
	"fmt"
	"testing"

	. "github.com/muxinc/protogen/proto3"
)

func TestScalarField_Validate(t *testing.T) {
	type fields struct {
		Name    NameType
		Tag     TagType
		Rule    FieldRule
		Comment string
		Typing  FieldType
	}
	tests := []struct {
		name    string
		fields  fields
		wantErr bool
	}{
		{
			name:    "Valid Scalar field",
			fields:  fields{Name: "MyMap", Tag: 1, Typing: StringType},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		s := &ScalarField{
			Name:    tt.fields.Name,
			Tag:     tt.fields.Tag,
			Rule:    tt.fields.Rule,
			Comment: tt.fields.Comment,
			Typing:  tt.fields.Typing,
		}
		if err := s.Validate(); (err != nil) != tt.wantErr {
			t.Errorf("%q. ScalarField.Validate() error = %v, wantErr %v", tt.name, err, tt.wantErr)
		}
	}
}

func TestSpec_Write(t *testing.T) {
	type fields struct {
		Package  string
		Imports  []ImportType
		Messages []Message
	}
	tests := []struct {
		name    string
		fields  fields
		want    string
		wantErr bool
	}{
		{
			name:    "Invalid spec with zero messages",
			fields:  fields{},
			wantErr: true,
		},
		{
			name: "Spec with message using reserved tags",
			fields: fields{
				Package: "foo",
				Messages: []Message{
					Message{
						Name:    "Beacon",
						Comment: "Beacon Message containing event information",
						Messages: []Message{
							Message{
								Name: "Event",
								ReservedValues: []Reserved{
									ReservedTagValue{Tag: 1},
									ReservedTagValue{Tag: 2},
									ReservedTagValue{Tag: 3},
									ReservedTagRange{LowerTag: 6, UpperTag: 9},
								},
								Fields: []Field{
									CustomField{Name: "Habitat", Typing: "string", Tag: 10, Rule: Repeated, Comment: "What am I?"},
									ScalarField{Name: "Continent", Typing: StringType, Tag: 11, Comment: "Where am I?"},
									MapField{Name: "LanguageMap", KeyTyping: StringType, ValueTyping: StringType, Tag: 12, Comment: "Super essential"},
								},
							},
						},
						ReservedValues: []Reserved{
							ReservedTagValue{Tag: 1},
							ReservedTagValue{Tag: 2},
							ReservedTagValue{Tag: 3},
							ReservedTagRange{LowerTag: 6, UpperTag: 9},
						},
						Fields: []Field{
							CustomField{Name: "Habitat", Typing: "string", Tag: 20, Comment: "What am I?"},
							ScalarField{Name: "Continent", Typing: StringType, Tag: 21, Rule: Repeated, Comment: "Where am I?"},
							MapField{Name: "LanguageMap", KeyTyping: StringType, ValueTyping: StringType, Tag: 22, Comment: "Super essential"},
							CustomMapField{Name: "CustomMap", KeyTyping: StringType, ValueTyping: "Event", Tag: 23},
						},
						OneOfs: []OneOf{
							OneOf{
								Name:    "test_oneof",
								Comment: "Can have a name or sub-message, but not both",
								Fields: []Field{
									ScalarField{Name: "name", Typing: StringType, Tag: 24, Comment: "Name"},
									CustomField{Name: "sub_message", Typing: "Event", Tag: 25, Comment: "Sub-Message"},
								},
							},
						},
						Enums: []Enum{
							Enum{
								Name: "Country",
								Values: []EnumValue{
									EnumValue{Name: "US", Tag: 0},
									EnumValue{Name: "CA", Tag: 1, Comment: "Canada"},
									EnumValue{Name: "GB", Tag: 2, Comment: "Great Britain"},
									EnumValue{Name: "MX", Tag: 3, Comment: "Mexico"},
								},
							},
							Enum{
								Name:       "PlaybackState",
								AllowAlias: true,
								Values: []EnumValue{
									EnumValue{Name: "Waiting", Tag: 0},
									EnumValue{Name: "Playing", Tag: 1},
									EnumValue{Name: "Started", Tag: 1},
									EnumValue{Name: "Stopped", Tag: 2},
								},
							},
						},
					},
				},
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		s := &Spec{
			Package:  tt.fields.Package,
			Imports:  tt.fields.Imports,
			Messages: tt.fields.Messages,
		}
		got, err := s.Write()
		if (err != nil) != tt.wantErr {
			t.Errorf("%q. Spec.Write() error = %v, wantErr %v", tt.name, err, tt.wantErr)
			continue
		}
		if tt.wantErr == false && got == "" {
			t.Error("String representing Protobuf spec was unexpectedly empty")
			continue
		}
	}
}

func ExampleSpec_Write() {
	spec := &Spec{
		Package: "foo",
		Messages: []Message{
			Message{
				Name:    "Beacon",
				Comment: "Beacon Message containing event information",
				Messages: []Message{
					Message{
						Name: "Event",
						ReservedValues: []Reserved{
							ReservedTagValue{Tag: 1},
							ReservedTagValue{Tag: 2},
							ReservedTagValue{Tag: 3},
							ReservedTagRange{LowerTag: 6, UpperTag: 9},
						},
						Fields: []Field{
							CustomField{Name: "Habitat", Typing: "string", Tag: 10, Rule: Repeated, Comment: "What am I?"},
							ScalarField{Name: "Continent", Typing: StringType, Tag: 11, Comment: "Where am I?"},
							MapField{Name: "LanguageMap", KeyTyping: StringType, ValueTyping: StringType, Tag: 12, Comment: "Super essential"},
						},
					},
				},
				ReservedValues: []Reserved{
					ReservedTagValue{Tag: 1},
					ReservedTagValue{Tag: 2},
					ReservedTagValue{Tag: 3},
					ReservedTagRange{LowerTag: 6, UpperTag: 9},
				},
				Fields: []Field{
					CustomField{Name: "Habitat", Typing: "string", Tag: 20, Comment: "What am I?"},
					ScalarField{Name: "Continent", Typing: StringType, Tag: 21, Rule: Repeated, Comment: "Where am I?"},
					MapField{Name: "LanguageMap", KeyTyping: StringType, ValueTyping: StringType, Tag: 22, Comment: "Super essential"},
					CustomMapField{Name: "CustomMap", KeyTyping: StringType, ValueTyping: "Event", Tag: 23},
				},
				OneOfs: []OneOf{
					OneOf{
						Name:    "test_oneof",
						Comment: "Can have a name or sub-message, but not both",
						Fields: []Field{
							ScalarField{Name: "name", Typing: StringType, Tag: 24, Comment: "Name"},
							CustomField{Name: "sub_message", Typing: "Event", Tag: 25, Comment: "Sub-Message"},
						},
					},
				},
				Enums: []Enum{
					Enum{
						Name:    "Country",
						Comment: "Country code",
						Values: []EnumValue{
							EnumValue{Name: "CA", Tag: 1, Comment: "Canada"},
							EnumValue{Name: "US", Tag: 0},
							EnumValue{Name: "MX", Tag: 3, Comment: "Mexico"},
							EnumValue{Name: "GB", Tag: 2, Comment: "Great Britain"},
						},
					},
					Enum{
						Name:       "PlaybackState",
						AllowAlias: true,
						Values: []EnumValue{
							EnumValue{Name: "Playing", Tag: 1},
							EnumValue{Name: "Waiting", Tag: 0},
							EnumValue{Name: "Stopped", Tag: 2},
							EnumValue{Name: "Started", Tag: 1},
						},
					},
				},
			},
		},
	}

	s, err := spec.Write()
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println(s)
	}

	// Output:
	// syntax = "proto3";
	// package foo;
	//
	// // Beacon Message containing event information
	// message Beacon {
	//   message Event {
	//     reserved 1;
	//     reserved 2;
	//     reserved 3;
	//     reserved 6 to 9;
	//
	//     repeated string Habitat = 10;   // What am I?
	//     string Continent = 11;   // Where am I?
	//     map<string, string> LanguageMap = 12;   // Super essential
	//
	//   }
	//
	//   // Country code
	//   enum Country {
	//     US = 0;
	//     CA = 1;   // Canada
	//     GB = 2;   // Great Britain
	//     MX = 3;   // Mexico
	//   }
	//
	//   enum PlaybackState {
	//     option allow_alias = true;
	//     Waiting = 0;
	//     Playing = 1;
	//     Started = 1;
	//     Stopped = 2;
	//   }
	//
	//   reserved 1;
	//   reserved 2;
	//   reserved 3;
	//   reserved 6 to 9;
	//
	//   string Habitat = 20;   // What am I?
	//   repeated string Continent = 21;   // Where am I?
	//   map<string, string> LanguageMap = 22;   // Super essential
	//   map<string, Event> CustomMap = 23;
	//
	//   // Can have a name or sub-message, but not both
	//   oneof test_oneof {
	//     string name = 24;   // Name
	//     Event sub_message = 25;   // Sub-Message
	//   }
	// }
}
<commit_msg>Format test source file<commit_after>package proto3_test

import (
	"fmt"
	"testing"

	. "github.com/muxinc/protogen/proto3"
)

func TestScalarField_Validate(t *testing.T) {
	type fields struct {
		Name    NameType
		Tag     TagType
		Rule    FieldRule
		Comment string
		Typing  FieldType
	}
	tests := []struct {
		name    string
		fields  fields
		wantErr bool
	}{
		{
			name:    "Valid Scalar field",
			fields:  fields{Name: "MyMap", Tag: 1, Typing: StringType},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		s := &ScalarField{
			Name:    tt.fields.Name,
			Tag:     tt.fields.Tag,
			Rule:    tt.fields.Rule,
			Comment: tt.fields.Comment,
			Typing:  tt.fields.Typing,
		}
		if err := s.Validate(); (err != nil) != tt.wantErr {
			t.Errorf("%q. ScalarField.Validate() error = %v, wantErr %v", tt.name, err, tt.wantErr)
		}
	}
}

func TestSpec_Write(t *testing.T) {
	type fields struct {
		Package  string
		Imports  []ImportType
		Messages []Message
	}
	tests := []struct {
		name    string
		fields  fields
		want    string
		wantErr bool
	}{
		{
			name:    "Invalid spec with zero messages",
			fields:  fields{},
			wantErr: true,
		},
		{
			name: "Spec with message using reserved tags",
			fields: fields{
				Package: "foo",
				Messages: []Message{
					{
						Name:    "Beacon",
						Comment: "Beacon Message containing event information",
						Messages: []Message{
							{
								Name: "Event",
								ReservedValues: []Reserved{
									ReservedTagValue{Tag: 1},
									ReservedTagValue{Tag: 2},
									ReservedTagValue{Tag: 3},
									ReservedTagRange{LowerTag: 6, UpperTag: 9},
								},
								Fields: []Field{
									CustomField{Name: "Habitat", Typing: "string", Tag: 10, Rule: Repeated, Comment: "What am I?"},
									ScalarField{Name: "Continent", Typing: StringType, Tag: 11, Comment: "Where am I?"},
									MapField{Name: "LanguageMap", KeyTyping: StringType, ValueTyping: StringType, Tag: 12, Comment: "Super essential"},
								},
							},
						},
						ReservedValues: []Reserved{
							ReservedTagValue{Tag: 1},
							ReservedTagValue{Tag: 2},
							ReservedTagValue{Tag: 3},
							ReservedTagRange{LowerTag: 6, UpperTag: 9},
						},
						Fields: []Field{
							CustomField{Name: "Habitat", Typing: "string", Tag: 20, Comment: "What am I?"},
							ScalarField{Name: "Continent", Typing: StringType, Tag: 21, Rule: Repeated, Comment: "Where am I?"},
							MapField{Name: "LanguageMap", KeyTyping: StringType, ValueTyping: StringType, Tag: 22, Comment: "Super essential"},
							CustomMapField{Name: "CustomMap", KeyTyping: StringType, ValueTyping: "Event", Tag: 23},
						},
						OneOfs: []OneOf{
							{
								Name:    "test_oneof",
								Comment: "Can have a name or sub-message, but not both",
								Fields: []Field{
									ScalarField{Name: "name", Typing: StringType, Tag: 24, Comment: "Name"},
									CustomField{Name: "sub_message", Typing: "Event", Tag: 25, Comment: "Sub-Message"},
								},
							},
						},
						Enums: []Enum{
							{
								Name: "Country",
								Values: []EnumValue{
									{Name: "US", Tag: 0},
									{Name: "CA", Tag: 1, Comment: "Canada"},
									{Name: "GB", Tag: 2, Comment: "Great Britain"},
									{Name: "MX", Tag: 3, Comment: "Mexico"},
								},
							},
							{
								Name:       "PlaybackState",
								AllowAlias: true,
								Values: []EnumValue{
									{Name: "Waiting", Tag: 0},
									{Name: "Playing", Tag: 1},
									{Name: "Started", Tag: 1},
									{Name: "Stopped", Tag: 2},
								},
							},
						},
					},
				},
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		s := &Spec{
			Package:  tt.fields.Package,
			Imports:  tt.fields.Imports,
			Messages: tt.fields.Messages,
		}
		got, err := s.Write()
		if (err != nil) != tt.wantErr {
			t.Errorf("%q. Spec.Write() error = %v, wantErr %v", tt.name, err, tt.wantErr)
			continue
		}
		if tt.wantErr == false && got == "" {
			t.Error("String representing Protobuf spec was unexpectedly empty")
			continue
		}
	}
}

func ExampleSpec_Write() {
	spec := &Spec{
		Package: "foo",
		Messages: []Message{
			{
				Name:    "Beacon",
				Comment: "Beacon Message containing event information",
				Messages: []Message{
					{
						Name: "Event",
						ReservedValues: []Reserved{
							ReservedTagValue{Tag: 1},
							ReservedTagValue{Tag: 2},
							ReservedTagValue{Tag: 3},
							ReservedTagRange{LowerTag: 6, UpperTag: 9},
						},
						Fields: []Field{
							CustomField{Name: "Habitat", Typing: "string", Tag: 10, Rule: Repeated, Comment: "What am I?"},
							ScalarField{Name: "Continent", Typing: StringType, Tag: 11, Comment: "Where am I?"},
							MapField{Name: "LanguageMap", KeyTyping: StringType, ValueTyping: StringType, Tag: 12, Comment: "Super essential"},
						},
					},
				},
				ReservedValues: []Reserved{
					ReservedTagValue{Tag: 1},
					ReservedTagValue{Tag: 2},
					ReservedTagValue{Tag: 3},
					ReservedTagRange{LowerTag: 6, UpperTag: 9},
				},
				Fields: []Field{
					CustomField{Name: "Habitat", Typing: "string", Tag: 20, Comment: "What am I?"},
					ScalarField{Name: "Continent", Typing: StringType, Tag: 21, Rule: Repeated, Comment: "Where am I?"},
					MapField{Name: "LanguageMap", KeyTyping: StringType, ValueTyping: StringType, Tag: 22, Comment: "Super essential"},
					CustomMapField{Name: "CustomMap", KeyTyping: StringType, ValueTyping: "Event", Tag: 23},
				},
				OneOfs: []OneOf{
					{
						Name:    "test_oneof",
						Comment: "Can have a name or sub-message, but not both",
						Fields: []Field{
							ScalarField{Name: "name", Typing: StringType, Tag: 24, Comment: "Name"},
							CustomField{Name: "sub_message", Typing: "Event", Tag: 25, Comment: "Sub-Message"},
						},
					},
				},
				Enums: []Enum{
					{
						Name:    "Country",
						Comment: "Country code",
						Values: []EnumValue{
							{Name: "CA", Tag: 1, Comment: "Canada"},
							{Name: "US", Tag: 0},
							{Name: "MX", Tag: 3, Comment: "Mexico"},
							{Name: "GB", Tag: 2, Comment: "Great Britain"},
						},
					},
					{
						Name:       "PlaybackState",
						AllowAlias: true,
						Values: []EnumValue{
							{Name: "Playing", Tag: 1},
							{Name: "Waiting", Tag: 0},
							{Name: "Stopped", Tag: 2},
							{Name: "Started", Tag: 1},
						},
					},
				},
			},
		},
	}

	s, err := spec.Write()
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println(s)
	}

	// Output:
	// syntax = "proto3";
	// package foo;
	//
	// // Beacon Message containing event information
	// message Beacon {
	//   message Event {
	//     reserved 1;
	//     reserved 2;
	//     reserved 3;
	//     reserved 6 to 9;
	//
	//     repeated string Habitat = 10;   // What am I?
	//     string Continent = 11;   // Where am I?
	//     map<string, string> LanguageMap = 12;   // Super essential
	//
	//   }
	//
	//   // Country code
	//   enum Country {
	//     US = 0;
	//     CA = 1;   // Canada
	//     GB = 2;   // Great Britain
	//     MX = 3;   // Mexico
	//   }
	//
	//   enum PlaybackState {
	//     option allow_alias = true;
	//     Waiting = 0;
	//     Playing = 1;
	//     Started = 1;
	//     Stopped = 2;
	//   }
	//
	//   reserved 1;
	//   reserved 2;
	//   reserved 3;
	//   reserved 6 to 9;
	//
	//   string Habitat = 20;   // What am I?
	//   repeated string Continent = 21;   // Where am I?
	//   map<string, string> LanguageMap = 22;   // Super essential
	//   map<string, Event> CustomMap = 23;
	//
	//   // Can have a name or sub-message, but not both
	//   oneof test_oneof {
	//     string name = 24;   // Name
	//     Event sub_message = 25;   // Sub-Message
	//   }
	// }
}
<|endoftext|>


<commit_before>package provider

import (
	"errors"
	"fmt"
	"strings"

	"github.com/kubernetes-incubator/external-dns/endpoint"
	"github.com/kubernetes-incubator/external-dns/plan"
	log "github.com/sirupsen/logrus"
	"github.com/transip/gotransip"
	transip "github.com/transip/gotransip/domain"
)

const (
	// 60 seconds is the current minimal TTL for TransIP and will replace unconfigured
	// TTL's for Endpoints
	transipMinimalValidTTL = 60
)

// TransIPProvider is an implementation of Provider for TransIP.
type TransIPProvider struct {
	client       gotransip.SOAPClient
	domainFilter DomainFilter
	dryRun       bool
}

// NewTransIPProvider initializes a new TransIP Provider.
func NewTransIPProvider(accountName, privateKeyFile string, domainFilter DomainFilter, dryRun bool) (*TransIPProvider, error) {
	// check given arguments
	if accountName == "" {
		return nil, errors.New("required --transip-account not set")
	}

	if privateKeyFile == "" {
		return nil, errors.New("required --transip-keyfile not set")
	}

	var apiMode gotransip.APIMode
	if dryRun {
		apiMode = gotransip.APIModeReadOnly
	} else {
		apiMode = gotransip.APIModeReadWrite
	}

	// create new TransIP API client
	c, err := gotransip.NewSOAPClient(gotransip.ClientConfig{
		AccountName:    accountName,
		PrivateKeyPath: privateKeyFile,
		Mode:           apiMode,
	})
	if err != nil {
		return nil, fmt.Errorf("could not setup TransIP API client: %s", err.Error())
	}

	// return tipCloud struct
	return &TransIPProvider{
		client:       c,
		domainFilter: domainFilter,
		dryRun:       dryRun,
	}, nil
}

// ApplyChanges applies a given set of changes in a given zone.
func (p *TransIPProvider) ApplyChanges(changes *plan.Changes) error {
	// build zonefinder with all our zones so we can use FindZone
	// and a mapping of zones and their domain name
	zones, err := p.fetchZones()
	if err != nil {
		return err
	}

	zoneNameMapper := zoneIDName{}
	zonesByName := make(map[string]transip.Domain)
	updatedZones := make(map[string]bool)
	for _, zone := range zones {
		// TransIP API doesn't expose a unique identifier for zones, other than than
		// the domain name itself
		zoneNameMapper.Add(zone.Name, zone.Name)
		zonesByName[zone.Name] = zone
	}

	zoneForZoneName := func(name string, m zoneIDName, z map[string]transip.Domain) (zone transip.Domain, err error) {
		_, zoneName := m.FindZone(name)
		if zoneName == "" {
			err = fmt.Errorf("could not find zoneName for %s", name)
			return
		}

		var ok bool
		zone, ok = zonesByName[zoneName]
		if !ok {
			err = fmt.Errorf("could not find zone for %s", zoneName)
		}

		return
	}

	// first see if we need to delete anything
	for _, ep := range changes.Delete {
		log.WithFields(log.Fields{"record": ep.DNSName, "type": ep.RecordType}).Info("endpoint has to go")

		zone, err := zoneForZoneName(ep.DNSName, zoneNameMapper, zonesByName)
		if err != nil {
			log.Errorf("could not find zone for %s: %s", ep.DNSName, err.Error())
			continue
		}

		log.Debugf("removing records for %s", zone.Name)

		// remove current records from DNS entry set
		entries := p.removeEndpointFromEntries(ep, zone)

		// update zone in zone map
		zone.DNSEntries = entries
		zonesByName[zone.Name] = zone
		// flag zone for updating
		updatedZones[zone.Name] = true
	}

	for _, ep := range changes.Create {
		log.WithFields(log.Fields{"record": ep.DNSName, "type": ep.RecordType}).Info("endpoint is missing")

		zone, err := zoneForZoneName(ep.DNSName, zoneNameMapper, zonesByName)
		if err != nil {
			log.Errorf("could not find zone for %s: %s", ep.DNSName, err.Error())
			continue
		}

		log.Debugf("creating records for %s", zone.Name)

		// add new entries to set
		zone.DNSEntries = p.addEndpointToEntries(ep, zone, zone.DNSEntries)

		// update zone in zone map
		zonesByName[zone.Name] = zone
		// flag zone for updating
		updatedZones[zone.Name] = true
		log.WithFields(log.Fields{"zone": zone.Name}).Debug("flagging for update")
	}

	for _, ep := range changes.UpdateNew {
		log.WithFields(log.Fields{"record": ep.DNSName, "type": ep.RecordType}).Debug("needs updating")

		zone, err := zoneForZoneName(ep.DNSName, zoneNameMapper, zonesByName)
		if err != nil {
			log.WithFields(log.Fields{"record": ep.DNSName}).Warn(err.Error())
			continue
		}

		// updating the records is basically finding all matching records according
		// to the name and the type, removing them from the set and add the new
		// records
		log.WithFields(log.Fields{
			"zone":       zone.Name,
			"dnsname":    ep.DNSName,
			"recordtype": ep.RecordType,
		}).Debug("removing matching entries")

		// remove current records from DNS entry set
		entries := p.removeEndpointFromEntries(ep, zone)

		// add new entries to set
		entries = p.addEndpointToEntries(ep, zone, entries)

		// check to see if actually anything changed in the DNSEntry set
		if p.dnsEntriesAreEqual(entries, zone.DNSEntries) {
			log.WithFields(log.Fields{"zone": zone.Name}).Debug("not updating identical entries")
			continue
		}

		// update zone in zone map
		zone.DNSEntries = entries
		zonesByName[zone.Name] = zone
		// flag zone for updating
		updatedZones[zone.Name] = true

		log.WithFields(log.Fields{"zone": zone.Name}).Debug("flagging for update")
	}

	// go over all updated zones and set new DNSEntry set
	for uz := range updatedZones {
		zone, ok := zonesByName[uz]
		if !ok {
			log.WithFields(log.Fields{"zone": uz}).Debug("updated zone no longer found")
			continue
		}

		if p.dryRun {
			log.WithFields(log.Fields{"zone": zone.Name}).Info("not updating in dry-run mode")
			continue
		}

		log.WithFields(log.Fields{"zone": zone.Name}).Info("updating DNS entries")
		if err := transip.SetDNSEntries(p.client, zone.Name, zone.DNSEntries); err != nil {
			log.WithFields(log.Fields{"zone": zone.Name, "error": err.Error()}).Warn("failed to update")
		}
	}

	return nil
}

// fetchZones returns a list of all domains within the account
func (p *TransIPProvider) fetchZones() ([]transip.Domain, error) {
	domainNames, err := transip.GetDomainNames(p.client)
	if err != nil {
		return nil, err
	}

	domains, err := transip.BatchGetInfo(p.client, domainNames)
	if err != nil {
		return nil, err
	}

	var zones []transip.Domain
	for _, d := range domains {
		if !p.domainFilter.Match(d.Name) {
			continue
		}

		zones = append(zones, d)
	}

	return zones, nil
}

// Zones returns the list of hosted zones.
func (p *TransIPProvider) Zones() ([]transip.Domain, error) {
	zones, err := p.fetchZones()
	if err != nil {
		return nil, err
	}

	return zones, nil
}

// Records returns the list of records in a given zone.
func (p *TransIPProvider) Records() ([]*endpoint.Endpoint, error) {
	zones, err := p.Zones()
	if err != nil {
		return nil, err
	}

	var endpoints []*endpoint.Endpoint
	var name string
	// go over all zones and their DNS entries and create endpoints for them
	for _, zone := range zones {
		for _, r := range zone.DNSEntries {
			if !supportedRecordType(string(r.Type)) {
				continue
			}

			name = p.endpointNameForRecord(r, zone)
			endpoints = append(endpoints, endpoint.NewEndpointWithTTL(name, string(r.Type), endpoint.TTL(r.TTL), r.Content))
		}
	}

	return endpoints, nil
}

// endpointNameForRecord returns "www.example.org" for DNSEntry with Name "www" and
// Doman with Name "example.org"
func (p *TransIPProvider) endpointNameForRecord(r transip.DNSEntry, d transip.Domain) string {
	// root name is identified by "@" and should be translated to domain name for
	// the endpoint entry.
	if r.Name == "@" {
		return d.Name
	}

	return fmt.Sprintf("%s.%s", r.Name, d.Name)
}

// recordNameForEndpoint returns "www" for Endpoint with DNSName "www.example.org"
// and Domain with Name "example.org"
func (p *TransIPProvider) recordNameForEndpoint(ep *endpoint.Endpoint, d transip.Domain) string {
	// root name is identified by "@" and should be translated to domain name for
	// the endpoint entry.
	if ep.DNSName == d.Name {
		return "@"
	}

	return strings.TrimSuffix(ep.DNSName, "."+d.Name)
}

// getMinimalValidTTL returns max between given Endpoint's RecordTTL and
// transipMinimalValidTTL
func (p *TransIPProvider) getMinimalValidTTL(ep *endpoint.Endpoint) int64 {
	// TTL cannot be lower than transipMinimalValidTTL
	if ep.RecordTTL < transipMinimalValidTTL {
		return transipMinimalValidTTL
	}

	return int64(ep.RecordTTL)
}

// dnsEntriesAreEqual compares the entries in 2 sets and returns true if the
// content of the entries is equal
func (p *TransIPProvider) dnsEntriesAreEqual(a, b transip.DNSEntries) bool {
	if len(a) != len(b) {
		return false
	}

	match := 0
	for _, aa := range a {
		for _, bb := range b {
			if aa.Content != bb.Content {
				continue
			}

			if aa.Name != bb.Name {
				continue
			}

			if aa.TTL != bb.TTL {
				continue
			}

			if aa.Type != bb.Type {
				continue
			}

			match += 1
		}
	}

	return (len(a) == match)
}

// removeEndpointFromEntries removes DNS entries from zone's set that match the
// type and name from given endpoint and returns the resulting DNS entry set
func (p *TransIPProvider) removeEndpointFromEntries(ep *endpoint.Endpoint, zone transip.Domain) transip.DNSEntries {
	// create new entry set
	entries := transip.DNSEntries{}
	// go over each DNS entry to see if it is a match
	for _, e := range zone.DNSEntries {
		// if we have match, don't copy it to the new entry set
		if p.endpointNameForRecord(e, zone) == ep.DNSName && string(e.Type) == ep.RecordType {
			log.WithFields(log.Fields{
				"name":    e.Name,
				"content": e.Content,
				"type":    e.Type,
			}).Debug("found match")
			continue
		}

		entries = append(entries, e)
	}

	return entries
}

// addEndpointToEntries creates DNS entries for given endpoint and returns
// resulting DNS entry set
func (p *TransIPProvider) addEndpointToEntries(ep *endpoint.Endpoint, zone transip.Domain, entries transip.DNSEntries) transip.DNSEntries {
	ttl := p.getMinimalValidTTL(ep)
	for _, target := range ep.Targets {
		log.WithFields(log.Fields{
			"zone":       zone.Name,
			"dnsname":    ep.DNSName,
			"recordtype": ep.RecordType,
			"ttl":        ttl,
			"target":     target,
		}).Debugf("adding new record")
		entries = append(entries, transip.DNSEntry{
			Name:    p.recordNameForEndpoint(ep, zone),
			TTL:     ttl,
			Type:    transip.DNSEntryType(ep.RecordType),
			Content: target,
		})
	}

	return entries
}
<commit_msg>moved zoneForZoneName function into TransIPProvider struct<commit_after>package provider

import (
	"errors"
	"fmt"
	"strings"

	"github.com/kubernetes-incubator/external-dns/endpoint"
	"github.com/kubernetes-incubator/external-dns/plan"
	log "github.com/sirupsen/logrus"
	"github.com/transip/gotransip"
	transip "github.com/transip/gotransip/domain"
)

const (
	// 60 seconds is the current minimal TTL for TransIP and will replace unconfigured
	// TTL's for Endpoints
	transipMinimalValidTTL = 60
)

// TransIPProvider is an implementation of Provider for TransIP.
type TransIPProvider struct {
	client       gotransip.SOAPClient
	domainFilter DomainFilter
	dryRun       bool
}

// NewTransIPProvider initializes a new TransIP Provider.
func NewTransIPProvider(accountName, privateKeyFile string, domainFilter DomainFilter, dryRun bool) (*TransIPProvider, error) {
	// check given arguments
	if accountName == "" {
		return nil, errors.New("required --transip-account not set")
	}

	if privateKeyFile == "" {
		return nil, errors.New("required --transip-keyfile not set")
	}

	var apiMode gotransip.APIMode
	if dryRun {
		apiMode = gotransip.APIModeReadOnly
	} else {
		apiMode = gotransip.APIModeReadWrite
	}

	// create new TransIP API client
	c, err := gotransip.NewSOAPClient(gotransip.ClientConfig{
		AccountName:    accountName,
		PrivateKeyPath: privateKeyFile,
		Mode:           apiMode,
	})
	if err != nil {
		return nil, fmt.Errorf("could not setup TransIP API client: %s", err.Error())
	}

	// return tipCloud struct
	return &TransIPProvider{
		client:       c,
		domainFilter: domainFilter,
		dryRun:       dryRun,
	}, nil
}

// ApplyChanges applies a given set of changes in a given zone.
func (p *TransIPProvider) ApplyChanges(changes *plan.Changes) error {
	// build zonefinder with all our zones so we can use FindZone
	// and a mapping of zones and their domain name
	zones, err := p.fetchZones()
	if err != nil {
		return err
	}

	zoneNameMapper := zoneIDName{}
	zonesByName := make(map[string]transip.Domain)
	updatedZones := make(map[string]bool)
	for _, zone := range zones {
		// TransIP API doesn't expose a unique identifier for zones, other than than
		// the domain name itself
		zoneNameMapper.Add(zone.Name, zone.Name)
		zonesByName[zone.Name] = zone
	}

	// first see if we need to delete anything
	for _, ep := range changes.Delete {
		log.WithFields(log.Fields{"record": ep.DNSName, "type": ep.RecordType}).Info("endpoint has to go")

		zone, err := p.zoneForZoneName(ep.DNSName, zoneNameMapper, zonesByName)
		if err != nil {
			log.Errorf("could not find zone for %s: %s", ep.DNSName, err.Error())
			continue
		}

		log.Debugf("removing records for %s", zone.Name)

		// remove current records from DNS entry set
		entries := p.removeEndpointFromEntries(ep, zone)

		// update zone in zone map
		zone.DNSEntries = entries
		zonesByName[zone.Name] = zone
		// flag zone for updating
		updatedZones[zone.Name] = true
	}

	for _, ep := range changes.Create {
		log.WithFields(log.Fields{"record": ep.DNSName, "type": ep.RecordType}).Info("endpoint is missing")

		zone, err := p.zoneForZoneName(ep.DNSName, zoneNameMapper, zonesByName)
		if err != nil {
			log.Errorf("could not find zone for %s: %s", ep.DNSName, err.Error())
			continue
		}

		log.Debugf("creating records for %s", zone.Name)

		// add new entries to set
		zone.DNSEntries = p.addEndpointToEntries(ep, zone, zone.DNSEntries)

		// update zone in zone map
		zonesByName[zone.Name] = zone
		// flag zone for updating
		updatedZones[zone.Name] = true
		log.WithFields(log.Fields{"zone": zone.Name}).Debug("flagging for update")
	}

	for _, ep := range changes.UpdateNew {
		log.WithFields(log.Fields{"record": ep.DNSName, "type": ep.RecordType}).Debug("needs updating")

		zone, err := p.zoneForZoneName(ep.DNSName, zoneNameMapper, zonesByName)
		if err != nil {
			log.WithFields(log.Fields{"record": ep.DNSName}).Warn(err.Error())
			continue
		}

		// updating the records is basically finding all matching records according
		// to the name and the type, removing them from the set and add the new
		// records
		log.WithFields(log.Fields{
			"zone":       zone.Name,
			"dnsname":    ep.DNSName,
			"recordtype": ep.RecordType,
		}).Debug("removing matching entries")

		// remove current records from DNS entry set
		entries := p.removeEndpointFromEntries(ep, zone)

		// add new entries to set
		entries = p.addEndpointToEntries(ep, zone, entries)

		// check to see if actually anything changed in the DNSEntry set
		if p.dnsEntriesAreEqual(entries, zone.DNSEntries) {
			log.WithFields(log.Fields{"zone": zone.Name}).Debug("not updating identical entries")
			continue
		}

		// update zone in zone map
		zone.DNSEntries = entries
		zonesByName[zone.Name] = zone
		// flag zone for updating
		updatedZones[zone.Name] = true

		log.WithFields(log.Fields{"zone": zone.Name}).Debug("flagging for update")
	}

	// go over all updated zones and set new DNSEntry set
	for uz := range updatedZones {
		zone, ok := zonesByName[uz]
		if !ok {
			log.WithFields(log.Fields{"zone": uz}).Debug("updated zone no longer found")
			continue
		}

		if p.dryRun {
			log.WithFields(log.Fields{"zone": zone.Name}).Info("not updating in dry-run mode")
			continue
		}

		log.WithFields(log.Fields{"zone": zone.Name}).Info("updating DNS entries")
		if err := transip.SetDNSEntries(p.client, zone.Name, zone.DNSEntries); err != nil {
			log.WithFields(log.Fields{"zone": zone.Name, "error": err.Error()}).Warn("failed to update")
		}
	}

	return nil
}

// fetchZones returns a list of all domains within the account
func (p *TransIPProvider) fetchZones() ([]transip.Domain, error) {
	domainNames, err := transip.GetDomainNames(p.client)
	if err != nil {
		return nil, err
	}

	domains, err := transip.BatchGetInfo(p.client, domainNames)
	if err != nil {
		return nil, err
	}

	var zones []transip.Domain
	for _, d := range domains {
		if !p.domainFilter.Match(d.Name) {
			continue
		}

		zones = append(zones, d)
	}

	return zones, nil
}

// Zones returns the list of hosted zones.
func (p *TransIPProvider) Zones() ([]transip.Domain, error) {
	zones, err := p.fetchZones()
	if err != nil {
		return nil, err
	}

	return zones, nil
}

// Records returns the list of records in a given zone.
func (p *TransIPProvider) Records() ([]*endpoint.Endpoint, error) {
	zones, err := p.Zones()
	if err != nil {
		return nil, err
	}

	var endpoints []*endpoint.Endpoint
	var name string
	// go over all zones and their DNS entries and create endpoints for them
	for _, zone := range zones {
		for _, r := range zone.DNSEntries {
			if !supportedRecordType(string(r.Type)) {
				continue
			}

			name = p.endpointNameForRecord(r, zone)
			endpoints = append(endpoints, endpoint.NewEndpointWithTTL(name, string(r.Type), endpoint.TTL(r.TTL), r.Content))
		}
	}

	return endpoints, nil
}

// endpointNameForRecord returns "www.example.org" for DNSEntry with Name "www" and
// Doman with Name "example.org"
func (p *TransIPProvider) endpointNameForRecord(r transip.DNSEntry, d transip.Domain) string {
	// root name is identified by "@" and should be translated to domain name for
	// the endpoint entry.
	if r.Name == "@" {
		return d.Name
	}

	return fmt.Sprintf("%s.%s", r.Name, d.Name)
}

// recordNameForEndpoint returns "www" for Endpoint with DNSName "www.example.org"
// and Domain with Name "example.org"
func (p *TransIPProvider) recordNameForEndpoint(ep *endpoint.Endpoint, d transip.Domain) string {
	// root name is identified by "@" and should be translated to domain name for
	// the endpoint entry.
	if ep.DNSName == d.Name {
		return "@"
	}

	return strings.TrimSuffix(ep.DNSName, "."+d.Name)
}

// getMinimalValidTTL returns max between given Endpoint's RecordTTL and
// transipMinimalValidTTL
func (p *TransIPProvider) getMinimalValidTTL(ep *endpoint.Endpoint) int64 {
	// TTL cannot be lower than transipMinimalValidTTL
	if ep.RecordTTL < transipMinimalValidTTL {
		return transipMinimalValidTTL
	}

	return int64(ep.RecordTTL)
}

// dnsEntriesAreEqual compares the entries in 2 sets and returns true if the
// content of the entries is equal
func (p *TransIPProvider) dnsEntriesAreEqual(a, b transip.DNSEntries) bool {
	if len(a) != len(b) {
		return false
	}

	match := 0
	for _, aa := range a {
		for _, bb := range b {
			if aa.Content != bb.Content {
				continue
			}

			if aa.Name != bb.Name {
				continue
			}

			if aa.TTL != bb.TTL {
				continue
			}

			if aa.Type != bb.Type {
				continue
			}

			match += 1
		}
	}

	return (len(a) == match)
}

// removeEndpointFromEntries removes DNS entries from zone's set that match the
// type and name from given endpoint and returns the resulting DNS entry set
func (p *TransIPProvider) removeEndpointFromEntries(ep *endpoint.Endpoint, zone transip.Domain) transip.DNSEntries {
	// create new entry set
	entries := transip.DNSEntries{}
	// go over each DNS entry to see if it is a match
	for _, e := range zone.DNSEntries {
		// if we have match, don't copy it to the new entry set
		if p.endpointNameForRecord(e, zone) == ep.DNSName && string(e.Type) == ep.RecordType {
			log.WithFields(log.Fields{
				"name":    e.Name,
				"content": e.Content,
				"type":    e.Type,
			}).Debug("found match")
			continue
		}

		entries = append(entries, e)
	}

	return entries
}

// addEndpointToEntries creates DNS entries for given endpoint and returns
// resulting DNS entry set
func (p *TransIPProvider) addEndpointToEntries(ep *endpoint.Endpoint, zone transip.Domain, entries transip.DNSEntries) transip.DNSEntries {
	ttl := p.getMinimalValidTTL(ep)
	for _, target := range ep.Targets {
		log.WithFields(log.Fields{
			"zone":       zone.Name,
			"dnsname":    ep.DNSName,
			"recordtype": ep.RecordType,
			"ttl":        ttl,
			"target":     target,
		}).Debugf("adding new record")
		entries = append(entries, transip.DNSEntry{
			Name:    p.recordNameForEndpoint(ep, zone),
			TTL:     ttl,
			Type:    transip.DNSEntryType(ep.RecordType),
			Content: target,
		})
	}

	return entries
}

// zoneForZoneName returns the zone mapped to given name or error if zone could
// not be found
func (p *TransIPProvider) zoneForZoneName(name string, m zoneIDName, z map[string]transip.Domain) (transip.Domain, error) {
	_, zoneName := m.FindZone(name)
	if zoneName == "" {
		return transip.Domain{}, fmt.Errorf("could not find zoneName for %s", name)
	}

	zone, ok := z[zoneName]
	if !ok {
		return zone, fmt.Errorf("could not find zone for %s", zoneName)
	}

	return zone, nil
}
<|endoftext|>


<commit_before>package fads

import (
	"math"
	"reflect"

	"github.com/go-hep/fmom"
	"github.com/go-hep/fwk"
)

type ParticlePropagator struct {
	fwk.TaskBase

	radius  float64
	radius2 float64
	halflen float64
	bz      float64

	input  string
	output string

	hadrons string
	eles    string
	muons   string
}

func (tsk *ParticlePropagator) Configure(ctx fwk.Context) fwk.Error {
	var err fwk.Error

	tsk.radius2 = tsk.radius * tsk.radius
	if tsk.radius < 1.0e-2 {
		return fwk.Errorf("%s: too small radius value (%v)", tsk.Name(), tsk.radius)
	}

	if tsk.halflen < 1.0e-2 {
		return fwk.Errorf("%s: too small 1/2-length value (%v)", tsk.Name(), tsk.halflen)
	}

	err = tsk.DeclInPort(tsk.input, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	err = tsk.DeclOutPort(tsk.output, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	err = tsk.DeclOutPort(tsk.hadrons, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	err = tsk.DeclOutPort(tsk.eles, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	err = tsk.DeclOutPort(tsk.muons, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	return err
}

func (tsk *ParticlePropagator) StartTask(ctx fwk.Context) fwk.Error {
	var err fwk.Error

	return err
}

func (tsk *ParticlePropagator) StopTask(ctx fwk.Context) fwk.Error {
	var err fwk.Error

	return err
}

func (tsk *ParticlePropagator) Process(ctx fwk.Context) fwk.Error {
	var err fwk.Error
	store := ctx.Store()
	msg := ctx.Msg()

	v, err := store.Get(tsk.input)
	if err != nil {
		return err
	}

	input := v.([]Candidate)
	msg.Debugf(">>> candidates: %v\n", len(input))

	output := make([]Candidate, 0)
	defer func() {
		err = store.Put(tsk.output, output)
	}()

	hadrons := make([]Candidate, 0)
	eles := make([]Candidate, 0)
	muons := make([]Candidate, 0)
	defer func() {
		err = store.Put(tsk.hadrons, hadrons)
		if err != nil {
			return
		}
		err = store.Put(tsk.eles, eles)
		if err != nil {
			return
		}
		err = store.Put(tsk.muons, muons)
		if err != nil {
			return
		}
	}()

	const (
		cLight  = 2.99792458E8
		cLight2 = cLight * cLight
	)

	for i := range input {
		cand := &input[i]
		x := cand.Pos.X() * 1e-3
		y := cand.Pos.Y() * 1e-3
		z := cand.Pos.Z() * 1e-3

		// is particle inside cylinder ?
		if math.Hypot(x, y) > tsk.radius || math.Abs(z) > tsk.halflen {
			continue
		}

		px := cand.Mom.Px()
		py := cand.Mom.Py()
		pt2 := px*px + py*py
		if pt2 < 1e-9 {
			continue
		}

		q := float64(cand.Charge())
		if math.Abs(q) < 1e-9 || math.Abs(tsk.bz) < 1e-9 {
			// solve pt2*t^2 + 2(px.x + py.y)*t + (radius2 - x.x - y.y) = 0
			v := px*y - py*x
			discr2 := pt2*tsk.radius2 - v*v
			if discr2 < 0 {
				// no solution
				continue
			}

			v = px*x + py*y
			discr := math.Sqrt(discr2)
			t1 := (-v + discr) / pt2
			t2 := (-v - discr) / pt2
			t := 0.0
			switch t1 < 0 {
			case true:
				t = t2
			case false:
				t = t1
			}

			pz := cand.Mom.Pz()
			e := cand.Mom.E()

			xt := x + px*t
			yt := y + py*t
			zt := z + pz*t

			mother := cand
			c := cand.Clone()
			c.Pos = fmom.NewPxPyPzE(xt*1e3, yt*1e3, zt*1e3, cand.Pos.T()+t*e*1e3)
			c.Add(mother)

			output = append(output, *c)
			if math.Abs(q) > 1e-9 {
				pid := c.Pid
				if pid < 0 {
					pid = 0
				}
				switch pid {
				case 11:
					eles = append(eles, *c)
				case 13:
					muons = append(muons, *c)
				default:
					hadrons = append(hadrons, *c)
				}
			}
		} else {
			// 1.  initial transverse momentum p_{T0} : Part->pt
			//     initial transverse momentum direction \phi_0 = -atan(p_X0/p_Y0)
			//     relativistic gamma : gamma = E/mc² ; gammam = gamma \times m
			//     giration frequency \omega = q/(gamma m) fBz
			//     helix radius r = p_T0 / (omega gamma m)

			e := cand.Mom.E()
			pt := cand.Mom.Pt()
			pz := cand.Mom.Pz()

			gammam := e * 1e9 / cLight2  // in eV/c^2
			omega := q * tsk.bz / gammam // omega is in [89875518 / s]
			r := pt / (q * tsk.bz) * 1e9 / cLight

			phi0 := math.Atan2(py, px) // [rad] in [-pi,pi)

			// 2. helix axis coordinates
			xc := x + r*math.Sin(phi0)
			yc := y - r*math.Cos(phi0)
			rc := math.Hypot(xc, yc)
			phic := math.Atan2(yc, xc)
			phi := phic
			if xc < 0 {
				phi += math.Pi
			}

			// 3. time evaluation t = TMath::Min(t_r, t_z)
			//    t_r : time to exit from the sides
			//    t_z : time to exit from the front or the back
			tr := 0.0 // in [ns]
			signpz := 1.0
			if math.Signbit(pz) {
				signpz = -1.0
			}

			tz := 0.0
			if pz == 0 {
				tz = 1e99
			} else {
				tz = gammam / (pz * 1e9 / cLight) * (-z * tsk.halflen * signpz)
			}

			absr := math.Abs(r)
			t := 0.0
			if rc+absr < tsk.radius {
				// helix does not cross the cylinder sides
				t = tz
			} else {
				asinrho := math.Asin((tsk.radius*tsk.radius - rc*rc - r*r) * 0.5 / (absr * rc))
				delta := phi0 - phi
				if delta < -math.Pi {
					delta += 2.0 * math.Pi
				}
				if delta > math.Pi {
					delta -= 2.0 * math.Pi
				}

				t1 := (delta + asinrho) / omega
				t2 := (delta + math.Pi - asinrho) / omega
				t3 := (delta + math.Pi + asinrho) / omega
				t4 := (delta - asinrho) / omega
				t5 := (delta - math.Pi - asinrho) / omega
				t6 := (delta - math.Pi + asinrho) / omega

				if t1 < 0 {
					t1 = 1.0E99
				}
				if t2 < 0 {
					t2 = 1.0E99
				}
				if t3 < 0 {
					t3 = 1.0E99
				}
				if t4 < 0 {
					t4 = 1.0E99
				}
				if t5 < 0 {
					t5 = 1.0E99
				}
				if t6 < 0 {
					t6 = 1.0E99
				}

				tra := math.Min(t1, math.Min(t2, t3))
				trb := math.Min(t4, math.Min(t5, t6))
				tr = math.Min(tra, trb)
				t = math.Min(tr, tz)
			}

			// 4. position in terms of x(t), y(t), z(t)
			xt := xc + r*math.Sin(omega*t-phi0)
			yt := yc + r*math.Cos(omega*t-phi0)
			zt := z + pz*1.0E9/cLight/gammam*t
			rt := math.Hypot(xt, yt)

			if rt > 0.0 {
				mother := cand
				c := cand.Clone()
				c.Pos = fmom.NewPxPyPzE(xt*1e3, yt*1e3, zt*1e3, cand.Pos.T()+t*e*1e3)
				c.Add(mother)

				output = append(output, *c)
				if math.Abs(q) > 1e-9 {
					pid := c.Pid
					if pid < 0 {
						pid = 0
					}
					switch pid {
					case 11:
						eles = append(eles, *c)
					case 13:
						muons = append(muons, *c)
					default:
						hadrons = append(hadrons, *c)
					}
				}

			}

		}
	}

	msg.Debugf(">>> output:     %v\n", len(output))
	return err
}

func init() {
	fwk.Register(reflect.TypeOf(ParticlePropagator{}),
		func(typ, name string, mgr fwk.App) (fwk.Component, fwk.Error) {
			var err fwk.Error
			tsk := &ParticlePropagator{
				TaskBase: fwk.NewTask(typ, name, mgr),
			}
			tsk.radius = 1.0
			err = tsk.DeclProp("Radius", &tsk.radius)
			if err != nil {
				return nil, err
			}
			tsk.radius2 = tsk.radius * tsk.radius

			tsk.halflen = 3.0
			err = tsk.DeclProp("HalfLength", &tsk.halflen)
			if err != nil {
				return nil, err
			}

			tsk.bz = 0.0
			err = tsk.DeclProp("Bz", &tsk.bz)
			if err != nil {
				return nil, err
			}

			tsk.input = "/fads/StableParticles"
			err = tsk.DeclProp("Input", &tsk.input)
			if err != nil {
				return nil, err
			}

			tsk.output = "StableParticles"
			err = tsk.DeclProp("Output", &tsk.output)
			if err != nil {
				return nil, err
			}

			tsk.hadrons = "ChargedHadrons"
			err = tsk.DeclProp("ChargedHadrons", &tsk.hadrons)
			if err != nil {
				return nil, err
			}

			tsk.eles = "Electrons"
			err = tsk.DeclProp("Electrons", &tsk.eles)
			if err != nil {
				return nil, err
			}

			tsk.muons = "Muons"
			err = tsk.DeclProp("Muons", &tsk.muons)
			if err != nil {
				return nil, err
			}

			return tsk, err
		},
	)
}

// EOF
<commit_msg>propagator: fixed propagation along z-axis<commit_after>package fads

import (
	"math"
	"reflect"

	"github.com/go-hep/fmom"
	"github.com/go-hep/fwk"
)

type ParticlePropagator struct {
	fwk.TaskBase

	radius  float64
	radius2 float64
	halflen float64
	bz      float64

	input  string
	output string

	hadrons string
	eles    string
	muons   string
}

func (tsk *ParticlePropagator) Configure(ctx fwk.Context) fwk.Error {
	var err fwk.Error

	tsk.radius2 = tsk.radius * tsk.radius
	if tsk.radius < 1.0e-2 {
		return fwk.Errorf("%s: too small radius value (%v)", tsk.Name(), tsk.radius)
	}

	if tsk.halflen < 1.0e-2 {
		return fwk.Errorf("%s: too small 1/2-length value (%v)", tsk.Name(), tsk.halflen)
	}

	err = tsk.DeclInPort(tsk.input, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	err = tsk.DeclOutPort(tsk.output, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	err = tsk.DeclOutPort(tsk.hadrons, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	err = tsk.DeclOutPort(tsk.eles, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	err = tsk.DeclOutPort(tsk.muons, reflect.TypeOf([]Candidate{}))
	if err != nil {
		return err
	}

	return err
}

func (tsk *ParticlePropagator) StartTask(ctx fwk.Context) fwk.Error {
	var err fwk.Error

	return err
}

func (tsk *ParticlePropagator) StopTask(ctx fwk.Context) fwk.Error {
	var err fwk.Error

	return err
}

func (tsk *ParticlePropagator) Process(ctx fwk.Context) fwk.Error {
	var err fwk.Error
	store := ctx.Store()
	msg := ctx.Msg()

	v, err := store.Get(tsk.input)
	if err != nil {
		return err
	}

	input := v.([]Candidate)
	msg.Debugf(">>> candidates: %v\n", len(input))

	output := make([]Candidate, 0)
	defer func() {
		err = store.Put(tsk.output, output)
	}()

	hadrons := make([]Candidate, 0)
	eles := make([]Candidate, 0)
	muons := make([]Candidate, 0)
	defer func() {
		err = store.Put(tsk.hadrons, hadrons)
		if err != nil {
			return
		}
		err = store.Put(tsk.eles, eles)
		if err != nil {
			return
		}
		err = store.Put(tsk.muons, muons)
		if err != nil {
			return
		}
	}()

	const (
		cLight  = 2.99792458E8
		cLight2 = cLight * cLight
	)

	for i := range input {
		cand := &input[i]
		x := cand.Pos.X() * 1e-3
		y := cand.Pos.Y() * 1e-3
		z := cand.Pos.Z() * 1e-3

		// is particle inside cylinder ?
		if math.Hypot(x, y) > tsk.radius || math.Abs(z) > tsk.halflen {
			continue
		}

		px := cand.Mom.Px()
		py := cand.Mom.Py()
		pt2 := px*px + py*py
		if pt2 < 1e-9 {
			continue
		}

		q := float64(cand.Charge())
		if math.Abs(q) < 1e-9 || math.Abs(tsk.bz) < 1e-9 {
			// solve pt2*t^2 + 2(px.x + py.y)*t + (radius2 - x.x - y.y) = 0
			v := px*y - py*x
			discr2 := pt2*tsk.radius2 - v*v
			if discr2 < 0 {
				// no solution
				continue
			}

			v = px*x + py*y
			discr := math.Sqrt(discr2)
			t1 := (-v + discr) / pt2
			t2 := (-v - discr) / pt2
			t := 0.0
			switch t1 < 0 {
			case true:
				t = t2
			case false:
				t = t1
			}

			pz := cand.Mom.Pz()
			e := cand.Mom.E()

			zt := z + pz*t
			if math.Abs(zt) > tsk.halflen {
				invpz := 1.0 / pz
				t3 := (+tsk.halflen - z) * invpz
				t4 := (-tsk.halflen - z) * invpz
				if t3 < 0 {
					t = t4
				} else {
					t = t3
				}
			}

			xt := x + px*t
			yt := y + py*t
			zt = z + pz*t

			mother := cand
			c := cand.Clone()
			c.Pos = fmom.NewPxPyPzE(xt*1e3, yt*1e3, zt*1e3, cand.Pos.T()+t*e*1e3)
			c.Add(mother)

			output = append(output, *c)
			if math.Abs(q) > 1e-9 {
				pid := c.Pid
				if pid < 0 {
					pid = 0
				}
				switch pid {
				case 11:
					eles = append(eles, *c)
				case 13:
					muons = append(muons, *c)
				default:
					hadrons = append(hadrons, *c)
				}
			}
		} else {
			// 1.  initial transverse momentum p_{T0} : Part->pt
			//     initial transverse momentum direction \phi_0 = -atan(p_X0/p_Y0)
			//     relativistic gamma : gamma = E/mc² ; gammam = gamma \times m
			//     giration frequency \omega = q/(gamma m) fBz
			//     helix radius r = p_T0 / (omega gamma m)

			e := cand.Mom.E()
			pt := cand.Mom.Pt()
			pz := cand.Mom.Pz()

			gammam := e * 1e9 / cLight2  // in eV/c^2
			omega := q * tsk.bz / gammam // omega is in [89875518 / s]
			r := pt / (q * tsk.bz) * 1e9 / cLight

			phi0 := math.Atan2(py, px) // [rad] in [-pi,pi)

			// 2. helix axis coordinates
			xc := x + r*math.Sin(phi0)
			yc := y - r*math.Cos(phi0)
			rc := math.Hypot(xc, yc)
			phic := math.Atan2(yc, xc)
			phi := phic
			if xc < 0 {
				phi += math.Pi
			}

			// 3. time evaluation t = TMath::Min(t_r, t_z)
			//    t_r : time to exit from the sides
			//    t_z : time to exit from the front or the back
			tr := 0.0 // in [ns]
			signpz := 1.0
			if math.Signbit(pz) {
				signpz = -1.0
			}

			tz := 0.0
			if pz == 0 {
				tz = 1e99
			} else {
				tz = gammam / (pz * 1e9 / cLight) * (-z * tsk.halflen * signpz)
			}

			absr := math.Abs(r)
			t := 0.0
			if rc+absr < tsk.radius {
				// helix does not cross the cylinder sides
				t = tz
			} else {
				asinrho := math.Asin((tsk.radius*tsk.radius - rc*rc - r*r) * 0.5 / (absr * rc))
				delta := phi0 - phi
				if delta < -math.Pi {
					delta += 2.0 * math.Pi
				}
				if delta > math.Pi {
					delta -= 2.0 * math.Pi
				}

				t1 := (delta + asinrho) / omega
				t2 := (delta + math.Pi - asinrho) / omega
				t3 := (delta + math.Pi + asinrho) / omega
				t4 := (delta - asinrho) / omega
				t5 := (delta - math.Pi - asinrho) / omega
				t6 := (delta - math.Pi + asinrho) / omega

				if t1 < 0 {
					t1 = 1.0E99
				}
				if t2 < 0 {
					t2 = 1.0E99
				}
				if t3 < 0 {
					t3 = 1.0E99
				}
				if t4 < 0 {
					t4 = 1.0E99
				}
				if t5 < 0 {
					t5 = 1.0E99
				}
				if t6 < 0 {
					t6 = 1.0E99
				}

				tra := math.Min(t1, math.Min(t2, t3))
				trb := math.Min(t4, math.Min(t5, t6))
				tr = math.Min(tra, trb)
				t = math.Min(tr, tz)
			}

			// 4. position in terms of x(t), y(t), z(t)
			xt := xc + r*math.Sin(omega*t-phi0)
			yt := yc + r*math.Cos(omega*t-phi0)
			zt := z + pz*1.0E9/cLight/gammam*t
			rt := math.Hypot(xt, yt)

			if rt > 0.0 {
				mother := cand
				c := cand.Clone()
				c.Pos = fmom.NewPxPyPzE(xt*1e3, yt*1e3, zt*1e3, cand.Pos.T()+t*e*1e3)
				c.Add(mother)

				output = append(output, *c)
				if math.Abs(q) > 1e-9 {
					pid := c.Pid
					if pid < 0 {
						pid = 0
					}
					switch pid {
					case 11:
						eles = append(eles, *c)
					case 13:
						muons = append(muons, *c)
					default:
						hadrons = append(hadrons, *c)
					}
				}

			}

		}
	}

	msg.Debugf(">>> output:     %v\n", len(output))
	return err
}

func init() {
	fwk.Register(reflect.TypeOf(ParticlePropagator{}),
		func(typ, name string, mgr fwk.App) (fwk.Component, fwk.Error) {
			var err fwk.Error
			tsk := &ParticlePropagator{
				TaskBase: fwk.NewTask(typ, name, mgr),
			}
			tsk.radius = 1.0
			err = tsk.DeclProp("Radius", &tsk.radius)
			if err != nil {
				return nil, err
			}
			tsk.radius2 = tsk.radius * tsk.radius

			tsk.halflen = 3.0
			err = tsk.DeclProp("HalfLength", &tsk.halflen)
			if err != nil {
				return nil, err
			}

			tsk.bz = 0.0
			err = tsk.DeclProp("Bz", &tsk.bz)
			if err != nil {
				return nil, err
			}

			tsk.input = "/fads/StableParticles"
			err = tsk.DeclProp("Input", &tsk.input)
			if err != nil {
				return nil, err
			}

			tsk.output = "StableParticles"
			err = tsk.DeclProp("Output", &tsk.output)
			if err != nil {
				return nil, err
			}

			tsk.hadrons = "ChargedHadrons"
			err = tsk.DeclProp("ChargedHadrons", &tsk.hadrons)
			if err != nil {
				return nil, err
			}

			tsk.eles = "Electrons"
			err = tsk.DeclProp("Electrons", &tsk.eles)
			if err != nil {
				return nil, err
			}

			tsk.muons = "Muons"
			err = tsk.DeclProp("Muons", &tsk.muons)
			if err != nil {
				return nil, err
			}

			return tsk, err
		},
	)
}

// EOF
<|endoftext|>


<commit_before>package scripts_test

import (
	"net/http"
	"os"
	"os/exec"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"
)

type Response struct {
	Resources    []UserResource `json:"resources"`
	StartIndex   int            `json:"startIndex"`
	ItemsPerPage int            `json:"itemsPerPage"`
	TotalResults int            `json:"totalResults"`
	Schemas      []string       `json:"schemas"`
}

type UserResource struct {
	ID       string `json:"id"`
	Username string `json:"userName"`
	Active   bool   `json:"active"`
}

var _ = Describe("uaa_lock_user.rb", func() {
	var (
		args       []string
		session    *gexec.Session
		server     *ghttp.Server
		response   *Response
		statusCode int
	)

	BeforeEach(func() {
		args = []string{"jane.smith@gov.uk"}

		statusCode = http.StatusOK
		response = &Response{
			Resources: []UserResource{{
				ID:       "17742243-4250-4266-ba8b-5d76dc97d52e",
				Username: "jane.smith@gov.uk",
			}},
			StartIndex:   1,
			ItemsPerPage: 100,
			TotalResults: 1,
			Schemas:      []string{"urn:scim:schemas:core:1.0"},
		}

		server = ghttp.NewServer()
		server.RouteToHandler("GET", "/Users", ghttp.RespondWithJSONEncodedPtr(&statusCode, &response))
	})

	AfterEach(func() {
		server.Close()
	})

	JustBeforeEach(func() {
		os.Setenv("TARGET", server.URL())
		os.Setenv("TOKEN", "faketoken")

		args = append([]string{"exec", "./uaa_lock_user.rb"}, args...)
		command := exec.Command("bundle", args...)

		var err error
		session, err = gexec.Start(command, GinkgoWriter, GinkgoWriter)
		Expect(err).ToNot(HaveOccurred())
	})

	Context("user doesn't exist", func() {
		BeforeEach(func() {
			args = []string{"john.smith@gov.uk"}

			response.Resources = []UserResource{}
			response.TotalResults = 0
		})

		It("returns an error", func() {
			Eventually(session).Should(gexec.Exit(1))
			Expect(session.Err).To(gbytes.Say("Username not found"))
		})
	})

	Context("more than one user", func() {
		BeforeEach(func() {
			response.Resources = append(response.Resources, response.Resources[0])
			response.TotalResults = 2
		})

		It("returns an error", func() {
			Eventually(session).Should(gexec.Exit(1))
			Expect(session.Err).To(gbytes.Say("Username is not unique"))
		})
	})

	Context("lock existing user", func() {
		Context("user is not locked", func() {
			BeforeEach(func() {
				response.Resources[0].Active = true

				server.RouteToHandler("PATCH", "/Users/17742243-4250-4266-ba8b-5d76dc97d52e", ghttp.RespondWith(statusCode, ""))
			})

			It("prints success message", func() {
				Eventually(session).Should(gexec.Exit(0))
				Expect(session.Out).To(gbytes.Say("locked jane.smith@gov.uk"))
			})
		})

		Context("user is locked", func() {
			BeforeEach(func() {
				response.Resources[0].Active = false
			})

			It("prints success message", func() {
				Eventually(session).Should(gexec.Exit(0))
				Expect(session.Out).To(gbytes.Say("locked jane.smith@gov.uk"))
			})
		})
	})

	Context("unlock existing user", func() {
		BeforeEach(func() {
			args = append(args, "-u")
		})

		Context("user is not locked", func() {
			BeforeEach(func() {
				response.Resources[0].Active = true
			})

			It("prints success message", func() {
				Eventually(session).Should(gexec.Exit(0))
				Expect(session.Out).To(gbytes.Say("unlocked jane.smith@gov.uk"))
			})
		})

		Context("user is locked", func() {
			BeforeEach(func() {
				response.Resources[0].Active = false

				server.RouteToHandler("PATCH", "/Users/17742243-4250-4266-ba8b-5d76dc97d52e", ghttp.RespondWith(statusCode, ""))
			})

			It("prints success message", func() {
				Eventually(session).Should(gexec.Exit(0))
				Expect(session.Out).To(gbytes.Say("unlocked jane.smith@gov.uk"))
			})
		})
	})
})
<commit_msg>(dev) increase timeout in uaa_lock_user_test.go<commit_after>package scripts_test

import (
	"net/http"
	"os"
	"os/exec"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"
)

type Response struct {
	Resources    []UserResource `json:"resources"`
	StartIndex   int            `json:"startIndex"`
	ItemsPerPage int            `json:"itemsPerPage"`
	TotalResults int            `json:"totalResults"`
	Schemas      []string       `json:"schemas"`
}

type UserResource struct {
	ID       string `json:"id"`
	Username string `json:"userName"`
	Active   bool   `json:"active"`
}

var _ = Describe("uaa_lock_user.rb", func() {
	var (
		args       []string
		session    *gexec.Session
		server     *ghttp.Server
		response   *Response
		statusCode int
	)

	BeforeEach(func() {
		args = []string{"jane.smith@gov.uk"}

		statusCode = http.StatusOK
		response = &Response{
			Resources: []UserResource{{
				ID:       "17742243-4250-4266-ba8b-5d76dc97d52e",
				Username: "jane.smith@gov.uk",
			}},
			StartIndex:   1,
			ItemsPerPage: 100,
			TotalResults: 1,
			Schemas:      []string{"urn:scim:schemas:core:1.0"},
		}

		server = ghttp.NewServer()
		server.RouteToHandler("GET", "/Users", ghttp.RespondWithJSONEncodedPtr(&statusCode, &response))
	})

	AfterEach(func() {
		server.Close()
	})

	JustBeforeEach(func() {
		os.Setenv("TARGET", server.URL())
		os.Setenv("TOKEN", "faketoken")

		args = append([]string{"exec", "./uaa_lock_user.rb"}, args...)
		command := exec.Command("bundle", args...)

		var err error
		session, err = gexec.Start(command, GinkgoWriter, GinkgoWriter)
		Expect(err).ToNot(HaveOccurred())
	})

	Context("user doesn't exist", func() {
		BeforeEach(func() {
			args = []string{"john.smith@gov.uk"}

			response.Resources = []UserResource{}
			response.TotalResults = 0
		})

		It("returns an error", func() {
			Eventually(session, "5s").Should(gexec.Exit(1))
			Expect(session.Err).To(gbytes.Say("Username not found"))
		})
	})

	Context("more than one user", func() {
		BeforeEach(func() {
			response.Resources = append(response.Resources, response.Resources[0])
			response.TotalResults = 2
		})

		It("returns an error", func() {
			Eventually(session, "5s").Should(gexec.Exit(1))
			Expect(session.Err).To(gbytes.Say("Username is not unique"))
		})
	})

	Context("lock existing user", func() {
		Context("user is not locked", func() {
			BeforeEach(func() {
				response.Resources[0].Active = true

				server.RouteToHandler("PATCH", "/Users/17742243-4250-4266-ba8b-5d76dc97d52e", ghttp.RespondWith(statusCode, ""))
			})

			It("prints success message", func() {
				Eventually(session, "5s").Should(gexec.Exit(0))
				Expect(session.Out).To(gbytes.Say("locked jane.smith@gov.uk"))
			})
		})

		Context("user is locked", func() {
			BeforeEach(func() {
				response.Resources[0].Active = false
			})

			It("prints success message", func() {
				Eventually(session, "5s").Should(gexec.Exit(0))
				Expect(session.Out).To(gbytes.Say("locked jane.smith@gov.uk"))
			})
		})
	})

	Context("unlock existing user", func() {
		BeforeEach(func() {
			args = append(args, "-u")
		})

		Context("user is not locked", func() {
			BeforeEach(func() {
				response.Resources[0].Active = true
			})

			It("prints success message", func() {
				Eventually(session, "5s").Should(gexec.Exit(0))
				Expect(session.Out).To(gbytes.Say("unlocked jane.smith@gov.uk"))
			})
		})

		Context("user is locked", func() {
			BeforeEach(func() {
				response.Resources[0].Active = false

				server.RouteToHandler("PATCH", "/Users/17742243-4250-4266-ba8b-5d76dc97d52e", ghttp.RespondWith(statusCode, ""))
			})

			It("prints success message", func() {
				Eventually(session, "5s").Should(gexec.Exit(0))
				Expect(session.Out).To(gbytes.Say("unlocked jane.smith@gov.uk"))
			})
		})
	})
})
<|endoftext|>