<commit_before>package heartbeat_test

import (
	. "github.com/ondrej-smola/mesos-go-http/scheduler/stage/heartbeat"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"context"
	"github.com/ondrej-smola/mesos-go-http/flow"
	"github.com/ondrej-smola/mesos-go-http/scheduler"
	"testing"
	"time"
)

func TestAck(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Heartbeat stage suite")
}

var _ = Describe("FrameworkId stage", func() {

	It("Set deadling from subscribed event", func(done Done) {
		fwId := New(WithMaxMissedHeartbeats(1))
		sink := flow.NewTestFlow()
		fwId.Via(sink)

		go func() {
			defer GinkgoRecover()
			id := "5"
			pull := sink.ExpectPull()
			s := scheduler.Subscribed(id)
			s.Subscribed.HeartbeatIntervalSeconds = 0.005 // 5 millis
			pull.Message(s)

			// this pull times out
			pull = sink.ExpectPull()
			start := time.Now()
			<-pull.Ctx.Done()
			Expect(time.Now().Sub(start)).To(BeNumerically("~", 10*time.Millisecond, time.Millisecond))
			pull.Error(pull.Ctx.Err())

			// this pull does not
			pull = sink.ExpectPull()
			time.Sleep(5 * time.Millisecond)
			pull.Message(&scheduler.Event{})
		}()

		_, err := fwId.Pull(context.Background())
		Expect(err).To(Succeed())

		_, err = fwId.Pull(context.Background())
		Expect(err).To(Equal(context.DeadlineExceeded))

		_, err = fwId.Pull(context.Background())
		Expect(err).To(Succeed())

		close(done)
	})

	It("Set deadling from configuration", func(done Done) {
		fwId := New(WithMaxMissedHeartbeats(3), WithHeartbeatDeadline(10*time.Millisecond))
		sink := flow.NewTestFlow()
		fwId.Via(sink)

		go func() {
			defer GinkgoRecover()
			id := "5"
			pull := sink.ExpectPull()
			s := scheduler.Subscribed(id)
			s.Subscribed.HeartbeatIntervalSeconds = 0.005 // 5 millis - should be IGNORED
			pull.Message(s)

			// this pull times out
			pull = sink.ExpectPull()
			start := time.Now()
			<-pull.Ctx.Done()
			Expect(time.Now().Sub(start)).To(BeNumerically("~", 40*time.Millisecond, time.Millisecond))
			pull.Error(pull.Ctx.Err())

			// this pull does not
			pull = sink.ExpectPull()
			time.Sleep(30 * time.Millisecond)
			pull.Message(&scheduler.Event{})<commit_msg>fixed flaky heartbeat tests<commit_after>package heartbeat_test

import (
	. "github.com/ondrej-smola/mesos-go-http/scheduler/stage/heartbeat"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"context"
	"github.com/ondrej-smola/mesos-go-http/flow"
	"github.com/ondrej-smola/mesos-go-http/scheduler"
	"testing"
	"time"
)

func TestAck(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Heartbeat stage suite")
}

var _ = Describe("FrameworkId stage", func() {

	It("Set deadling from subscribed event", func(done Done) {
		fwId := New(WithMaxMissedHeartbeats(1))
		sink := flow.NewTestFlow()
		fwId.Via(sink)

		go func() {
			defer GinkgoRecover()
			id := "5"
			pull := sink.ExpectPull()
			s := scheduler.Subscribed(id)
			s.Subscribed.HeartbeatIntervalSeconds = 0.01 // 10 millis
			pull.Message(s)

			// this pull times out
			pull = sink.ExpectPull()
			start := time.Now()
			<-pull.Ctx.Done()
			Expect(time.Now().Sub(start)).To(BeNumerically(">=", 20*time.Millisecond, 10*time.Millisecond))
			pull.Error(pull.Ctx.Err())

			// this pull does not
			pull = sink.ExpectPull()
			time.Sleep(5 * time.Millisecond)
			pull.Message(&scheduler.Event{})
		}()

		_, err := fwId.Pull(context.Background())
		Expect(err).To(Succeed())

		_, err = fwId.Pull(context.Background())
		Expect(err).To(Equal(context.DeadlineExceeded))

		_, err = fwId.Pull(context.Background())
		Expect(err).To(Succeed())

		close(done)
	})

	It("Set deadling from configuration", func(done Done) {
		fwId := New(WithMaxMissedHeartbeats(3), WithHeartbeatDeadline(10*time.Millisecond))
		sink := flow.NewTestFlow()
		fwId.Via(sink)

		go func() {
			defer GinkgoRecover()
			id := "5"
			pull := sink.ExpectPull()
			s := scheduler.Subscribed(id)
			s.Subscribed.HeartbeatIntervalSeconds = 0.005 // 5 millis - should be IGNORED
			pull.Message(s)

			// this pull times out
			pull = sink.ExpectPull()
			start := time.Now()
			<-pull.Ctx.Done()
			Expect(time.Now().Sub(start)).To(BeNumerically(">=", 40*time.Millisecond, 10*time.Millisecond))
			pull.Error(pull.Ctx.Err())

			// this pull does not
			pull = sink.ExpectPull()
			time.Sleep(30 * time.Millisecond)
			pull.Message(&scheduler.Event{})
{'fraction_non_alphanumeric': 0.16483516483516483, 'fraction_numerical': 0.011446886446886446, 'mean_word_length': 3.8544444444444443, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, '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': '523163', 'n_tokens_mistral': 1870, 'n_tokens_neox': 1612, 'n_words': 270}