<commit_before>package ingester

import (
	"reflect"
	"runtime"
	"testing"
	"time"

	"github.com/prometheus/common/log"
	"github.com/stretchr/testify/require"
	"github.com/weaveworks/cortex/ring"
)

func defaultIngesterTestConfig() Config {
	consul := ring.NewMockConsulClient()
	return Config{
		ringConfig: ring.Config{
			ConsulConfig: ring.ConsulConfig{
				Mock: consul,
			},
		},

		NumTokens:       1,
		HeartbeatPeriod: 5 * time.Second,
		ListenPort:      func(i int) *int { return &i }(0),
		addr:            "localhost",
		id:              "localhost",

		FlushCheckPeriod: 99999 * time.Hour,
		MaxChunkIdle:     99999 * time.Hour,
	}
}

// TestIngesterRestart tests a restarting ingester doesn't keep adding more tokens.
func TestIngesterRestart(t *testing.T) {
	config := defaultIngesterTestConfig()
	config.skipUnregister = true

	{
		ingester, err := New(config, nil)
		require.NoError(t, err)
		time.Sleep(100 * time.Millisecond)
		ingester.Shutdown() // doesn't actually unregister due to skipUnregister: true
	}

	poll(t, 100*time.Millisecond, 1, func() interface{} {
		return numTokens(config.ringConfig.ConsulConfig.Mock, "localhost")
	})

	{
		ingester, err := New(config, nil)
		require.NoError(t, err)
		time.Sleep(100 * time.Millisecond)
		ingester.Shutdown() // doesn't actually unregister due to skipUnregister: true
	}

	time.Sleep(200 * time.Millisecond)

	poll(t, 100*time.Millisecond, 1, func() interface{} {
		return numTokens(config.ringConfig.ConsulConfig.Mock, "localhost")
	})
}

func TestIngesterTransfer(t *testing.T) {
	cfg := defaultIngesterTestConfig()
	store := &testStore{
		chunks: map[string][]chunk.Chunk{},
	}
	ing1, err := New(cfg, store)
	require.NoError(t, err)

	const userID = "1"

	var (
		ts  = model.TimeFromUnix(123)
		val = model.SampleValue(456)
	)

	ctx := user.Inject(context.Background(), userID)
	_, err = ing1.Push(ctx, util.ToWriteRequest([]model.Sample{
		{
			Metric: model.Metric{
				model.MetricNameLabel: "foo",
			},
			Timestamp: ts,
			Value:     val,
		},
	}))
	require.NoError(t, err)

}

func numTokens(c ring.ConsulClient, name string) int {
	ringDesc, err := c.Get(ring.ConsulKey)
	if err != nil {
		log.Errorf("Error reading consul: %v", err)
		return 0
	}

	count := 0
	for _, token := range ringDesc.(*ring.Desc).Tokens {
		if token.Ingester == name {
			count++
		}
	}
	return count
}

// poll repeatedly evaluates condition until we either timeout, or it succeeds.
func poll(t *testing.T, d time.Duration, want interface{}, have func() interface{}) {
	deadline := time.Now().Add(d)
	for {
		if time.Now().After(deadline) {
			break
		}
		if reflect.DeepEqual(want, have()) {
			return
		}
		time.Sleep(d / 10)
	}<commit_msg>Never meant to commit that.<commit_after>package ingester

import (
	"reflect"
	"runtime"
	"testing"
	"time"

	"github.com/prometheus/common/log"
	"github.com/stretchr/testify/require"
	"github.com/weaveworks/cortex/ring"
)

func defaultIngesterTestConfig() Config {
	consul := ring.NewMockConsulClient()
	return Config{
		ringConfig: ring.Config{
			ConsulConfig: ring.ConsulConfig{
				Mock: consul,
			},
		},

		NumTokens:       1,
		HeartbeatPeriod: 5 * time.Second,
		ListenPort:      func(i int) *int { return &i }(0),
		addr:            "localhost",
		id:              "localhost",

		FlushCheckPeriod: 99999 * time.Hour,
		MaxChunkIdle:     99999 * time.Hour,
	}
}

// TestIngesterRestart tests a restarting ingester doesn't keep adding more tokens.
func TestIngesterRestart(t *testing.T) {
	config := defaultIngesterTestConfig()
	config.skipUnregister = true

	{
		ingester, err := New(config, nil)
		require.NoError(t, err)
		time.Sleep(100 * time.Millisecond)
		ingester.Shutdown() // doesn't actually unregister due to skipUnregister: true
	}

	poll(t, 100*time.Millisecond, 1, func() interface{} {
		return numTokens(config.ringConfig.ConsulConfig.Mock, "localhost")
	})

	{
		ingester, err := New(config, nil)
		require.NoError(t, err)
		time.Sleep(100 * time.Millisecond)
		ingester.Shutdown() // doesn't actually unregister due to skipUnregister: true
	}

	time.Sleep(200 * time.Millisecond)

	poll(t, 100*time.Millisecond, 1, func() interface{} {
		return numTokens(config.ringConfig.ConsulConfig.Mock, "localhost")
	})
}

func numTokens(c ring.ConsulClient, name string) int {
	ringDesc, err := c.Get(ring.ConsulKey)
	if err != nil {
		log.Errorf("Error reading consul: %v", err)
		return 0
	}

	count := 0
	for _, token := range ringDesc.(*ring.Desc).Tokens {
		if token.Ingester == name {
			count++
		}
	}
	return count
}

// poll repeatedly evaluates condition until we either timeout, or it succeeds.
func poll(t *testing.T, d time.Duration, want interface{}, have func() interface{}) {
	deadline := time.Now().Add(d)
	for {
		if time.Now().After(deadline) {
			break
		}
		if reflect.DeepEqual(want, have()) {
			return
		}
		time.Sleep(d / 10)
	}
{'fraction_non_alphanumeric': 0.13817138171381713, 'fraction_numerical': 0.015785157851578517, 'mean_word_length': 3.718568665377176, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3934341', 'n_tokens_mistral': 1912, 'n_tokens_neox': 1710, 'n_words': 395}