<commit_before>package inigo_test

import (
	"fmt"
	"syscall"
	"time"

	"github.com/cloudfoundry-incubator/inigo/fixtures"
	"github.com/cloudfoundry-incubator/inigo/helpers"
	"github.com/cloudfoundry-incubator/inigo/loggredile"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/nu7hatch/gouuid"
	"github.com/tedsuo/ifrit"

	"github.com/cloudfoundry-incubator/inigo/inigo_server"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	archive_helper "github.com/pivotal-golang/archiver/extractor/test_helper"
)

var _ = Describe("Convergence to desired state", func() {
	var desiredAppRequest models.DesireAppRequestFromCC
	var appId string
	var processGuid string

	var tpsProcess ifrit.Process
	var tpsAddr string

	var logOutput *gbytes.Buffer
	var stop chan<- bool

	CONVERGE_REPEAT_INTERVAL := time.Second
	WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL := CONVERGE_REPEAT_INTERVAL * 3

	BeforeEach(func() {
		guid, err := uuid.NewV4()
		if err != nil {
			panic("Failed to generate App ID")
		}
		appId = guid.String()

		guid, err = uuid.NewV4()
		if err != nil {
			panic("Failed to generate Process Guid")
		}
		processGuid = guid.String()

		suiteContext.FileServerRunner.Start()
		suiteContext.AuctioneerRunner.Start(AUCTION_MAX_ROUNDS)
		suiteContext.AppManagerRunner.Start()
		suiteContext.RouteEmitterRunner.Start()
		suiteContext.RouterRunner.Start()
		suiteContext.ConvergerRunner.Start(CONVERGE_REPEAT_INTERVAL, 30*time.Second, 5*time.Minute, 30*time.Second, 300*time.Second)

		tpsProcess = ifrit.Envoke(suiteContext.TPSRunner)
		tpsAddr = fmt.Sprintf("http://127.0.0.1:%d", suiteContext.TPSPort)

		archive_helper.CreateZipArchive("/tmp/simple-echo-droplet.zip", fixtures.HelloWorldIndexApp())
		inigo_server.UploadFile("simple-echo-droplet.zip", "/tmp/simple-echo-droplet.zip")

		suiteContext.FileServerRunner.ServeFile("some-lifecycle-bundle.tgz", suiteContext.SharedContext.CircusZipPath)

		logOutput, stop = loggredile.StreamIntoGBuffer(
			suiteContext.LoggregatorRunner.Config.OutgoingPort,
			fmt.Sprintf("/tail/?app=%s", appId),
			"App",
		)
	})

	AfterEach(func() {
		tpsProcess.Signal(syscall.SIGKILL)
		Eventually(tpsProcess.Wait()).Should(Receive())
		close(stop)
	})

	Describe("Executor fault tolerance", func() {
		Context("When starting a long-running process and then bouncing the executor", func() {
			BeforeEach(func() {
				suiteContext.ExecutorRunner.Start()
				suiteContext.RepRunner.Start()

				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})

			It("Eventually brings the long-running process up", func() {
				suiteContext.ExecutorRunner.Stop()

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(BeEmpty())
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(BeEmpty())

				suiteContext.ExecutorRunner.Start()

				running_lrps_poller = helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller = helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})
		})

		Context("When trying to start a long-running process before the executor is up", func() {
			BeforeEach(func() {
				suiteContext.RepRunner.Start()

				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				time.Sleep(WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL)

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Ω(running_lrps_poller()).Should(BeEmpty())
				Ω(hello_world_instance_poller()).Should(BeEmpty())
			})

			It("Eventually brings the long-running process up", func() {
				suiteContext.ExecutorRunner.Start()

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})
		})

		Context("When there is a runaway long-running process with no corresponding desired process", func() {
			BeforeEach(func() {
				suiteContext.RepRunner.Start()
				suiteContext.ExecutorRunner.Start()

				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})

			It("Eventually brings the long-running process down", func() {
				suiteContext.RepRunner.Stop()

				desiredAppStopRequest := models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 0,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppStopRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				time.Sleep(WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL)
				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				Ω(running_lrps_poller()).Should(HaveLen(1))

				suiteContext.RepRunner.Start()

				running_lrps_poller = helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(BeEmpty())
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(BeEmpty())
			})
		})

		Context("When a stop message for an instance of a long-running process is lost", func() {
			BeforeEach(func() {
				suiteContext.RepRunner.Start()
				suiteContext.ExecutorRunner.Start()

				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 2,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(2))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0", "1"}))
			})

			It("Eventually brings the long-running process down", func() {
				suiteContext.RepRunner.Stop()

				desiredAppStopRequest := models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppStopRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				time.Sleep(WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL)
				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				Ω(running_lrps_poller()).Should(HaveLen(2))

				suiteContext.RepRunner.Start()

				running_lrps_poller = helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})
		})
	})
})
<commit_msg>when an auctioneer turns on after an auction is requested: that auction runs<commit_after>package inigo_test

import (
	"fmt"
	"syscall"
	"time"

	"github.com/cloudfoundry-incubator/inigo/fixtures"
	"github.com/cloudfoundry-incubator/inigo/helpers"
	"github.com/cloudfoundry-incubator/inigo/loggredile"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/nu7hatch/gouuid"
	"github.com/tedsuo/ifrit"

	"github.com/cloudfoundry-incubator/inigo/inigo_server"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	archive_helper "github.com/pivotal-golang/archiver/extractor/test_helper"
)

var _ = Describe("Convergence to desired state", func() {
	var desiredAppRequest models.DesireAppRequestFromCC
	var appId string
	var processGuid string

	var tpsProcess ifrit.Process
	var tpsAddr string

	var logOutput *gbytes.Buffer
	var stop chan<- bool

	CONVERGE_REPEAT_INTERVAL := time.Second
	PENDING_AUCTION_KICK_THRESHOLD := time.Second
	WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL := CONVERGE_REPEAT_INTERVAL * 3

	Describe("Executor fault tolerance", func() {
		BeforeEach(func() {
			guid, err := uuid.NewV4()
			if err != nil {
				panic("Failed to generate App ID")
			}
			appId = guid.String()

			guid, err = uuid.NewV4()
			if err != nil {
				panic("Failed to generate Process Guid")
			}
			processGuid = guid.String()

			suiteContext.FileServerRunner.Start()
			suiteContext.AuctioneerRunner.Start(AUCTION_MAX_ROUNDS)
			suiteContext.AppManagerRunner.Start()
			suiteContext.RouteEmitterRunner.Start()
			suiteContext.RouterRunner.Start()
			suiteContext.ConvergerRunner.Start(CONVERGE_REPEAT_INTERVAL, 30*time.Second, 5*time.Minute, PENDING_AUCTION_KICK_THRESHOLD, 300*time.Second)

			tpsProcess = ifrit.Envoke(suiteContext.TPSRunner)
			tpsAddr = fmt.Sprintf("http://127.0.0.1:%d", suiteContext.TPSPort)

			archive_helper.CreateZipArchive("/tmp/simple-echo-droplet.zip", fixtures.HelloWorldIndexApp())
			inigo_server.UploadFile("simple-echo-droplet.zip", "/tmp/simple-echo-droplet.zip")

			suiteContext.FileServerRunner.ServeFile("some-lifecycle-bundle.tgz", suiteContext.SharedContext.CircusZipPath)

			logOutput, stop = loggredile.StreamIntoGBuffer(
				suiteContext.LoggregatorRunner.Config.OutgoingPort,
				fmt.Sprintf("/tail/?app=%s", appId),
				"App",
			)
		})

		AfterEach(func() {
			tpsProcess.Signal(syscall.SIGKILL)
			Eventually(tpsProcess.Wait()).Should(Receive())
			close(stop)
		})

		Context("When starting a long-running process and then bouncing the executor", func() {
			BeforeEach(func() {
				suiteContext.ExecutorRunner.Start()
				suiteContext.RepRunner.Start()

				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})

			It("Eventually brings the long-running process up", func() {
				suiteContext.ExecutorRunner.Stop()

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(BeEmpty())
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(BeEmpty())

				suiteContext.ExecutorRunner.Start()

				running_lrps_poller = helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller = helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})
		})

		Context("When trying to start a long-running process before the executor is up", func() {
			BeforeEach(func() {
				suiteContext.RepRunner.Start()

				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				time.Sleep(WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL)

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Ω(running_lrps_poller()).Should(BeEmpty())
				Ω(hello_world_instance_poller()).Should(BeEmpty())
			})

			It("Eventually brings the long-running process up", func() {
				suiteContext.ExecutorRunner.Start()

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})
		})

		Context("When there is a runaway long-running process with no corresponding desired process", func() {
			BeforeEach(func() {
				suiteContext.RepRunner.Start()
				suiteContext.ExecutorRunner.Start()

				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})

			It("Eventually brings the long-running process down", func() {
				suiteContext.RepRunner.Stop()

				desiredAppStopRequest := models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 0,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppStopRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				time.Sleep(WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL)
				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				Ω(running_lrps_poller()).Should(HaveLen(1))

				suiteContext.RepRunner.Start()

				running_lrps_poller = helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(BeEmpty())
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(BeEmpty())
			})
		})

		Context("When a stop message for an instance of a long-running process is lost", func() {
			BeforeEach(func() {
				suiteContext.RepRunner.Start()
				suiteContext.ExecutorRunner.Start()

				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 2,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(2))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0", "1"}))
			})

			It("Eventually brings the long-running process down", func() {
				suiteContext.RepRunner.Stop()

				desiredAppStopRequest := models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppStopRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				time.Sleep(WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL)
				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				Ω(running_lrps_poller()).Should(HaveLen(2))

				suiteContext.RepRunner.Start()

				running_lrps_poller = helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})
		})
	})

	Describe("Auctioneer Fault Tolerance", func() {
		BeforeEach(func() {
			guid, err := uuid.NewV4()
			if err != nil {
				panic("Failed to generate App ID")
			}
			appId = guid.String()

			guid, err = uuid.NewV4()
			if err != nil {
				panic("Failed to generate Process Guid")
			}
			processGuid = guid.String()

			suiteContext.FileServerRunner.Start()
			suiteContext.AppManagerRunner.Start()
			suiteContext.RouteEmitterRunner.Start()
			suiteContext.RouterRunner.Start()
			suiteContext.ConvergerRunner.Start(CONVERGE_REPEAT_INTERVAL, 30*time.Second, 5*time.Minute, PENDING_AUCTION_KICK_THRESHOLD, 300*time.Second)
			suiteContext.ExecutorRunner.Start()
			suiteContext.RepRunner.Start()

			tpsProcess = ifrit.Envoke(suiteContext.TPSRunner)
			tpsAddr = fmt.Sprintf("http://127.0.0.1:%d", suiteContext.TPSPort)

			archive_helper.CreateZipArchive("/tmp/simple-echo-droplet.zip", fixtures.HelloWorldIndexApp())
			inigo_server.UploadFile("simple-echo-droplet.zip", "/tmp/simple-echo-droplet.zip")

			suiteContext.FileServerRunner.ServeFile("some-lifecycle-bundle.tgz", suiteContext.SharedContext.CircusZipPath)

			logOutput, stop = loggredile.StreamIntoGBuffer(
				suiteContext.LoggregatorRunner.Config.OutgoingPort,
				fmt.Sprintf("/tail/?app=%s", appId),
				"App",
			)
		})

		AfterEach(func() {
			tpsProcess.Signal(syscall.SIGKILL)
			Eventually(tpsProcess.Wait()).Should(Receive())
			close(stop)
		})

		Context("When trying to start an auction before an Auctioneer is up", func() {
			BeforeEach(func() {
				desiredAppRequest = models.DesireAppRequestFromCC{
					ProcessGuid:  processGuid,
					DropletUri:   inigo_server.DownloadUrl("simple-echo-droplet.zip"),
					Stack:        suiteContext.RepStack,
					Environment:  []models.EnvironmentVariable{{Key: "VCAP_APPLICATION", Value: "{}"}},
					NumInstances: 1,
					Routes:       []string{"route-to-simple"},
					StartCommand: "./run",
					LogGuid:      appId,
				}

				err := suiteContext.NatsRunner.MessageBus.Publish("diego.desire.app", desiredAppRequest.ToJSON())
				Ω(err).ShouldNot(HaveOccurred())

				time.Sleep(PENDING_AUCTION_KICK_THRESHOLD)
				time.Sleep(WAIT_FOR_MULTIPLE_CONVERGE_INTERVAL)

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Ω(running_lrps_poller()).Should(BeEmpty())
				Ω(hello_world_instance_poller()).Should(BeEmpty())
			})

			It("Eventually brings the long-running process up", func() {
				suiteContext.AuctioneerRunner.Start(AUCTION_MAX_ROUNDS)

				running_lrps_poller := helpers.RunningLRPInstancesPoller(tpsAddr, processGuid)
				hello_world_instance_poller := helpers.HelloWorldInstancePoller(suiteContext.RouterRunner.Addr(), "route-to-simple")
				Eventually(running_lrps_poller, LONG_TIMEOUT).Should(HaveLen(1))
				Eventually(hello_world_instance_poller, LONG_TIMEOUT, 1).Should(Equal([]string{"0"}))
			})
		})
	})
})
<|endoftext|>


<commit_before>package core

import (
	"context"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"os"
	"testing"
	"time"

	ethCrypto "github.com/ethereum/go-ethereum/crypto"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/go-livepeer/drivers"
	"github.com/livepeer/go-livepeer/eth"
	"github.com/livepeer/go-livepeer/pm"
	"github.com/livepeer/lpms/ffmpeg"
	"github.com/livepeer/lpms/stream"
)

func Over1Pct(val int, cmp int) bool {
	return float32(val) > float32(cmp)*1.01 || float32(val) < float32(cmp)*0.99
}

func StubSegment() *stream.HLSSegment {
	d, _ := ioutil.ReadFile("./test.ts")
	return &stream.HLSSegment{SeqNo: 100, Name: "test.ts", Data: d[0:402696], Duration: 1}
}

func StubJobId() int64 {
	return int64(1234)
}

var videoProfiles = []ffmpeg.VideoProfile{ffmpeg.P144p30fps16x9, ffmpeg.P240p30fps16x9}

func TestTranscode(t *testing.T) {
	//Set up the node
	drivers.NodeStorage = drivers.NewMemoryDriver("")
	db, _ := common.InitDB("file:TestTranscode?mode=memory&cache=shared")
	defer db.Close()
	seth := &eth.StubClient{}
	tmp, _ := ioutil.TempDir("", "")
	n, _ := NewLivepeerNode(seth, tmp, db)
	defer os.RemoveAll(tmp)
	ffmpeg.InitFFmpeg()

	ss := StubSegment()
	md := &SegTranscodingMetadata{Profiles: videoProfiles}

	// Check nil transcoder.
	tr, err := n.sendToTranscodeLoop(md, ss)
	if err != ErrTranscoderAvail {
		t.Error("Error transcoding ", err)
	}

	// Sanity check full flow.
	n.Transcoder = NewLocalTranscoder(tmp)
	tr, err = n.sendToTranscodeLoop(md, ss)
	if err != nil {
		t.Error("Error transcoding ", err)
	}

	if len(tr.Data) != len(videoProfiles) && len(videoProfiles) != 2 {
		t.Error("Job profile count did not match broadcasters")
	}

	// 	Check transcode result
	if Over1Pct(len(tr.Data[0]), 65424) { // 144p
		t.Error("Unexpected transcode result ", len(tr.Data[0]))
	}
	if Over1Pct(len(tr.Data[1]), 81968) { // 240p
		t.Error("Unexpected transcode result ", len(tr.Data[1]))
	}

	// TODO check transcode loop expiry, storage, sig construction, etc
}

func TestTranscodeLoop_GivenNoSegmentsPastTimeout_CleansSegmentChan(t *testing.T) {
	//Set up the node
	drivers.NodeStorage = drivers.NewMemoryDriver("")
	db, _ := common.InitDB("file:TestTranscode?mode=memory&cache=shared")
	defer db.Close()
	seth := &eth.StubClient{}
	tmp, _ := ioutil.TempDir("", "")
	n, _ := NewLivepeerNode(seth, tmp, db)
	defer os.RemoveAll(tmp)
	ffmpeg.InitFFmpeg()
	ss := StubSegment()
	md := &SegTranscodingMetadata{Profiles: videoProfiles}
	n.Transcoder = NewLocalTranscoder(tmp)

	transcodeLoopTimeout = 100 * time.Millisecond
	assert := assert.New(t)
	require := require.New(t)

	_, err := n.sendToTranscodeLoop(md, ss)
	require.Nil(err)
	segChan := getSegChan(n, md.ManifestID)
	require.NotNil(segChan)

	waitForTranscoderLoopTimeout(n, md.ManifestID)

	segChan = getSegChan(n, md.ManifestID)
	assert.Nil(segChan)
}

func TestTranscodeLoop_GivenOnePMSession_RedeemsOneSession(t *testing.T) {
	recipient := new(pm.MockRecipient)
	//Set up the node
	drivers.NodeStorage = drivers.NewMemoryDriver("")
	db, _ := common.InitDB("file:TestTranscode?mode=memory&cache=shared")
	defer db.Close()
	seth := &eth.StubClient{}
	tmp, _ := ioutil.TempDir("", "")
	n, _ := NewLivepeerNode(seth, tmp, db)
	n.Recipient = recipient
	defer os.RemoveAll(tmp)
	ffmpeg.InitFFmpeg()
	ss := StubSegment()
	md := &SegTranscodingMetadata{Profiles: videoProfiles}
	n.Transcoder = NewLocalTranscoder(tmp)

	transcodeLoopTimeout = 100 * time.Millisecond
	require := require.New(t)

	sessionID := "some session ID"
	n.pmSessionsMutex.Lock()
	n.PMSessions[md.ManifestID] = make(map[string]bool)
	n.PMSessions[md.ManifestID][sessionID] = true
	n.pmSessionsMutex.Unlock()

	recipient.On("RedeemWinningTickets", []string{sessionID}[:]).Return(nil)

	_, err := n.sendToTranscodeLoop(md, ss)
	require.Nil(err)
	waitForTranscoderLoopTimeout(n, md.ManifestID)

	recipient.AssertExpectations(t)
}

func waitForTranscoderLoopTimeout(n *LivepeerNode, m ManifestID) {
	for i := 0; i < 3; i++ {
		time.Sleep(transcodeLoopTimeout * 2)
		segChan := getSegChan(n, m)
		if segChan == nil {
			return
		}
	}
}

func getSegChan(n *LivepeerNode, m ManifestID) SegmentChan {
	n.segmentMutex.Lock()
	defer n.segmentMutex.Unlock()

	return n.SegmentChans[m]
}

// XXX unclear what the tests below check
type Vint interface {
	Call(nums ...int)
}

type Vimp struct{}

func (*Vimp) Call(nums ...int) {
	fmt.Println(nums[0])
}

func TestInterface(t *testing.T) {
	var obj Vint
	obj = &Vimp{}
	obj.Call(4, 5, 6)
}

func TestSync(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	boolChan := make(chan bool)
	intChan := chanRoutine(ctx, boolChan)
	go insertBool(boolChan)
	go monitorChan(intChan)
	time.Sleep(time.Second)
	cancel()

	// time.Sleep(time.Second * 5)
}

func insertBool(boolChan chan bool) {
	for {
		boolChan <- true
		time.Sleep(500 * time.Millisecond)
	}
}

func chanRoutine(ctx context.Context, boolChan chan bool) chan int {
	intChan := make(chan int)
	go func() {
		for i := 0; ; i++ {
			select {
			case <-boolChan:
				intChan <- i
			case <-ctx.Done():
				fmt.Println("Done")
				return
			}
		}
	}()
	return intChan
}

func monitorChan(intChan chan int) {
	for {
		select {
		case i := <-intChan:
			fmt.Printf("i:%v\n", i)
		}
	}
}

func TestCrypto(t *testing.T) {
	blkNumB := make([]byte, 8)
	binary.BigEndian.PutUint64(blkNumB, uint64(9994353847340985734))
	fmt.Printf("%x\n\n", blkNumB)

	newb := make([]byte, 32)
	copy(newb[24:], blkNumB[:])
	fmt.Printf("%x\n\n", newb)

	i, _ := binary.Uvarint(ethCrypto.Keccak256(newb, ethCrypto.Keccak256([]byte("abc"))))
	fmt.Printf("%x\n\n", i%1)
}
<commit_msg>PM: end-of-stream ticket redemption tests<commit_after>package core

import (
	"context"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"os"
	"testing"
	"time"

	ethCrypto "github.com/ethereum/go-ethereum/crypto"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/go-livepeer/drivers"
	"github.com/livepeer/go-livepeer/eth"
	"github.com/livepeer/go-livepeer/pm"
	"github.com/livepeer/lpms/ffmpeg"
	"github.com/livepeer/lpms/stream"
)

func Over1Pct(val int, cmp int) bool {
	return float32(val) > float32(cmp)*1.01 || float32(val) < float32(cmp)*0.99
}

func StubSegment() *stream.HLSSegment {
	d, _ := ioutil.ReadFile("./test.ts")
	return &stream.HLSSegment{SeqNo: 100, Name: "test.ts", Data: d[0:402696], Duration: 1}
}

func StubJobId() int64 {
	return int64(1234)
}

var videoProfiles = []ffmpeg.VideoProfile{ffmpeg.P144p30fps16x9, ffmpeg.P240p30fps16x9}

func TestTranscode(t *testing.T) {
	//Set up the node
	drivers.NodeStorage = drivers.NewMemoryDriver("")
	db, _ := common.InitDB("file:TestTranscode?mode=memory&cache=shared")
	defer db.Close()
	seth := &eth.StubClient{}
	tmp, _ := ioutil.TempDir("", "")
	n, _ := NewLivepeerNode(seth, tmp, db)
	defer os.RemoveAll(tmp)
	ffmpeg.InitFFmpeg()

	ss := StubSegment()
	md := &SegTranscodingMetadata{Profiles: videoProfiles}

	// Check nil transcoder.
	tr, err := n.sendToTranscodeLoop(md, ss)
	if err != ErrTranscoderAvail {
		t.Error("Error transcoding ", err)
	}

	// Sanity check full flow.
	n.Transcoder = NewLocalTranscoder(tmp)
	tr, err = n.sendToTranscodeLoop(md, ss)
	if err != nil {
		t.Error("Error transcoding ", err)
	}

	if len(tr.Data) != len(videoProfiles) && len(videoProfiles) != 2 {
		t.Error("Job profile count did not match broadcasters")
	}

	// 	Check transcode result
	if Over1Pct(len(tr.Data[0]), 65424) { // 144p
		t.Error("Unexpected transcode result ", len(tr.Data[0]))
	}
	if Over1Pct(len(tr.Data[1]), 81968) { // 240p
		t.Error("Unexpected transcode result ", len(tr.Data[1]))
	}

	// TODO check transcode loop expiry, storage, sig construction, etc
}

func TestTranscodeLoop_GivenNoSegmentsPastTimeout_CleansSegmentChan(t *testing.T) {
	//Set up the node
	drivers.NodeStorage = drivers.NewMemoryDriver("")
	db, _ := common.InitDB("file:TestTranscode?mode=memory&cache=shared")
	defer db.Close()
	seth := &eth.StubClient{}
	tmp, _ := ioutil.TempDir("", "")
	n, _ := NewLivepeerNode(seth, tmp, db)
	defer os.RemoveAll(tmp)
	ffmpeg.InitFFmpeg()
	ss := StubSegment()
	md := &SegTranscodingMetadata{Profiles: videoProfiles}
	n.Transcoder = NewLocalTranscoder(tmp)

	transcodeLoopTimeout = 100 * time.Millisecond
	assert := assert.New(t)
	require := require.New(t)

	_, err := n.sendToTranscodeLoop(md, ss)
	require.Nil(err)
	segChan := getSegChan(n, md.ManifestID)
	require.NotNil(segChan)

	waitForTranscoderLoopTimeout(n, md.ManifestID)

	segChan = getSegChan(n, md.ManifestID)
	assert.Nil(segChan)
}

func TestTranscodeLoop_GivenOnePMSession_RedeemsOneSession(t *testing.T) {
	recipient := new(pm.MockRecipient)
	//Set up the node
	drivers.NodeStorage = drivers.NewMemoryDriver("")
	db, _ := common.InitDB("file:TestTranscode?mode=memory&cache=shared")
	defer db.Close()
	seth := &eth.StubClient{}
	tmp, _ := ioutil.TempDir("", "")
	n, _ := NewLivepeerNode(seth, tmp, db)
	n.Recipient = recipient
	defer os.RemoveAll(tmp)
	ffmpeg.InitFFmpeg()
	ss := StubSegment()
	md := &SegTranscodingMetadata{Profiles: videoProfiles}
	n.Transcoder = NewLocalTranscoder(tmp)

	transcodeLoopTimeout = 100 * time.Millisecond
	require := require.New(t)

	sessionID := "some session ID"
	n.pmSessionsMutex.Lock()
	n.PMSessions[md.ManifestID] = make(map[string]bool)
	n.PMSessions[md.ManifestID][sessionID] = true
	n.pmSessionsMutex.Unlock()

	recipient.On("RedeemWinningTickets", []string{sessionID}[:]).Return(nil)

	_, err := n.sendToTranscodeLoop(md, ss)
	require.Nil(err)
	waitForTranscoderLoopTimeout(n, md.ManifestID)

	recipient.AssertExpectations(t)
}

func TestTranscodeLoop_GivenMultiplePMSession_RedeemsAllSessions(t *testing.T) {
	recipient := new(pm.MockRecipient)
	//Set up the node
	drivers.NodeStorage = drivers.NewMemoryDriver("")
	db, _ := common.InitDB("file:TestTranscode?mode=memory&cache=shared")
	defer db.Close()
	seth := &eth.StubClient{}
	tmp, _ := ioutil.TempDir("", "")
	n, _ := NewLivepeerNode(seth, tmp, db)
	n.Recipient = recipient
	defer os.RemoveAll(tmp)
	ffmpeg.InitFFmpeg()
	ss := StubSegment()
	md := &SegTranscodingMetadata{Profiles: videoProfiles}
	n.Transcoder = NewLocalTranscoder(tmp)

	transcodeLoopTimeout = 100 * time.Millisecond
	require := require.New(t)

	sessionIDs := []string{"first session ID", "second session ID"}
	n.pmSessionsMutex.Lock()
	n.PMSessions[md.ManifestID] = make(map[string]bool)
	n.PMSessions[md.ManifestID][sessionIDs[0]] = true
	n.PMSessions[md.ManifestID][sessionIDs[1]] = true
	n.pmSessionsMutex.Unlock()

	recipient.On("RedeemWinningTickets", sessionIDs).Return(nil)

	_, err := n.sendToTranscodeLoop(md, ss)
	require.Nil(err)
	waitForTranscoderLoopTimeout(n, md.ManifestID)

	recipient.AssertExpectations(t)
}

func TestTranscodeLoop_GivenNoPMSession_DoesntTryToRedeem(t *testing.T) {
	recipient := new(pm.MockRecipient)
	//Set up the node
	drivers.NodeStorage = drivers.NewMemoryDriver("")
	db, _ := common.InitDB("file:TestTranscode?mode=memory&cache=shared")
	defer db.Close()
	seth := &eth.StubClient{}
	tmp, _ := ioutil.TempDir("", "")
	n, _ := NewLivepeerNode(seth, tmp, db)
	n.Recipient = recipient
	defer os.RemoveAll(tmp)
	ffmpeg.InitFFmpeg()
	ss := StubSegment()
	md := &SegTranscodingMetadata{Profiles: videoProfiles}
	n.Transcoder = NewLocalTranscoder(tmp)
	transcodeLoopTimeout = 100 * time.Millisecond
	require := require.New(t)
	recipient.On("RedeemWinningTickets", mock.Anything).Return(nil)

	_, err := n.sendToTranscodeLoop(md, ss)
	require.Nil(err)
	waitForTranscoderLoopTimeout(n, md.ManifestID)

	recipient.AssertNotCalled(t, "RedeemWinningTickets", mock.Anything)
}

func waitForTranscoderLoopTimeout(n *LivepeerNode, m ManifestID) {
	for i := 0; i < 3; i++ {
		time.Sleep(transcodeLoopTimeout * 2)
		segChan := getSegChan(n, m)
		if segChan == nil {
			return
		}
	}
}

func getSegChan(n *LivepeerNode, m ManifestID) SegmentChan {
	n.segmentMutex.Lock()
	defer n.segmentMutex.Unlock()

	return n.SegmentChans[m]
}

// XXX unclear what the tests below check
type Vint interface {
	Call(nums ...int)
}

type Vimp struct{}

func (*Vimp) Call(nums ...int) {
	fmt.Println(nums[0])
}

func TestInterface(t *testing.T) {
	var obj Vint
	obj = &Vimp{}
	obj.Call(4, 5, 6)
}

func TestSync(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	boolChan := make(chan bool)
	intChan := chanRoutine(ctx, boolChan)
	go insertBool(boolChan)
	go monitorChan(intChan)
	time.Sleep(time.Second)
	cancel()

	// time.Sleep(time.Second * 5)
}

func insertBool(boolChan chan bool) {
	for {
		boolChan <- true
		time.Sleep(500 * time.Millisecond)
	}
}

func chanRoutine(ctx context.Context, boolChan chan bool) chan int {
	intChan := make(chan int)
	go func() {
		for i := 0; ; i++ {
			select {
			case <-boolChan:
				intChan <- i
			case <-ctx.Done():
				fmt.Println("Done")
				return
			}
		}
	}()
	return intChan
}

func monitorChan(intChan chan int) {
	for {
		select {
		case i := <-intChan:
			fmt.Printf("i:%v\n", i)
		}
	}
}

func TestCrypto(t *testing.T) {
	blkNumB := make([]byte, 8)
	binary.BigEndian.PutUint64(blkNumB, uint64(9994353847340985734))
	fmt.Printf("%x\n\n", blkNumB)

	newb := make([]byte, 32)
	copy(newb[24:], blkNumB[:])
	fmt.Printf("%x\n\n", newb)

	i, _ := binary.Uvarint(ethCrypto.Keccak256(newb, ethCrypto.Keccak256([]byte("abc"))))
	fmt.Printf("%x\n\n", i%1)
}
<|endoftext|>


<commit_before>package common

import (
	"bytes"
	"encoding/binary"
)

type RouterInfo []byte

func (router_info RouterInfo) RouterIdentity() RouterIdentity {
	router_identity, _, _ := readRouterIdentity(router_info)
	return router_identity
}

func (router_info RouterInfo) Published() (d Date) {
	_, remainder, _ := readRouterIdentity(router_info)
	copy(remainder[:8], d[:])
	return
}

func (router_info RouterInfo) RouterAddressCount() int {
	_, remainder, _ := readRouterIdentity(router_info)
	var count int
	buf := bytes.NewReader(
		[]byte{remainder[8]},
	)
	binary.Read(buf, binary.BigEndian, &count)
	return count
}

func (router_info RouterInfo) RouterAddresses() []RouterAddress {
	var router_address RouterAddress
	remaining := router_info[9:]
	var err error
	addresses := make([]RouterAddress, 0)
	for i := 0; i < router_info.RouterAddressCount(); i++ {
		router_address, remaining, err = readRouterAddress(remaining)
		addresses = append(addresses, router_address)
	}
	return addresses
}

func (router_info RouterInfo) PeerSize() int {
	return 0
}

func (router_info RouterInfo) Options() Mapping {
	head := router_info.optionsLocation()
	size := head + router_info.optionsSize()
	return Mapping(router_info[head:size])
}

func (router_info RouterInfo) Signature() []byte {
	offset := router_info.optionsLocation() + router_info.optionsSize()
	sig_size := router_info.
		RouterIdentity().
		Certificate().
		signatureSize()
	return router_info[offset:sig_size]
}

func (router_info RouterInfo) optionsLocation() int {
	offset := 9
	var router_address RouterAddress
	remaining := router_info[9:]
	var err error
	for i := 0; i < router_info.RouterAddressCount(); i++ {
		router_address, remaining, err = readRouterAddress(remaining)
		offset := len(router_address)
	}
	return offset
}

func (router_info RouterInfo) optionsSize() int {
	head := router_info.optionsLocation()
	return int(binary.BigEndian.Uint16(router_info[head : head+1]))
}
<commit_msg>make it compile<commit_after>package common

import (
	"bytes"
	"encoding/binary"
)

type RouterInfo []byte

func (router_info RouterInfo) RouterIdentity() RouterIdentity {
	router_identity, _, _ := readRouterIdentity(router_info)
	return router_identity
}

func (router_info RouterInfo) Published() (d Date) {
	_, remainder, _ := readRouterIdentity(router_info)
	copy(remainder[:8], d[:])
	return
}

func (router_info RouterInfo) RouterAddressCount() int {
	_, remainder, _ := readRouterIdentity(router_info)
	var count int
	buf := bytes.NewReader(
		[]byte{remainder[8]},
	)
	binary.Read(buf, binary.BigEndian, &count)
	return count
}

func (router_info RouterInfo) RouterAddresses() []RouterAddress {
	var router_address RouterAddress
	remaining := router_info[9:]
	var err error
	addresses := make([]RouterAddress, 0)
	for i := 0; i < router_info.RouterAddressCount(); i++ {
		router_address, remaining, err = readRouterAddress(remaining)
		if err == nil {
			addresses = append(addresses, router_address)
		}
	}
	return addresses
}

func (router_info RouterInfo) PeerSize() int {
	return 0
}

func (router_info RouterInfo) Options() Mapping {
	head := router_info.optionsLocation()
	size := head + router_info.optionsSize()
	return Mapping(router_info[head:size])
}

func (router_info RouterInfo) Signature() []byte {
	offset := router_info.optionsLocation() + router_info.optionsSize()
	sig_size := router_info.
		RouterIdentity().
		Certificate().
		signatureSize()
	return router_info[offset:sig_size]
}

func (router_info RouterInfo) optionsLocation() int {
	offset := 9
	var router_address RouterAddress
	remaining := router_info[9:]
	for i := 0; i < router_info.RouterAddressCount(); i++ {
		router_address, remaining, _ = readRouterAddress(remaining)
		offset = len(router_address)
	}
	return offset
}

func (router_info RouterInfo) optionsSize() int {
	head := router_info.optionsLocation()
	return int(binary.BigEndian.Uint16(router_info[head : head+1]))
}
<|endoftext|>


<commit_before>package config

import (
	"github.com/mmmorris1975/aws-config/config"
	"github.com/mmmorris1975/aws-runas/lib/credentials"
	"os"
	"testing"
	"time"
)

func TestNewConfigResolver(t *testing.T) {
	t.Run("nil config", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		if r == nil {
			t.Error("unexpected nil resolver")
			return
		}

		if r.file == nil || len(r.file.Path) < 1 {
			t.Error("config file not set")
		}
	})

	t.Run("good", func(t *testing.T) {
		r, err := NewConfigResolver(&AwsConfig{Region: "us-west-2"})
		if err != nil {
			t.Error(err)
			return
		}

		if r.userConfig.Region != "us-west-2" {
			t.Error("unexpected region")
		}
	})

	t.Run("config file env var", func(t *testing.T) {
		os.Setenv(config.ConfigFileEnvVar, "test/config")
		defer os.Unsetenv(config.ConfigFileEnvVar)

		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		if r.file == nil || r.file.Path != "test/config" {
			t.Error("bad file name")
		}
	})
}

func TestConfigResolver_ResolveConfig(t *testing.T) {
	os.Setenv(config.ConfigFileEnvVar, "test/config")
	defer os.Unsetenv(config.ConfigFileEnvVar)

	t.Run("default only", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("")
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-west-1" {
			t.Error("unexpected region")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("unexpected mfa serial")
		}

		if len(c.RoleArn) > 0 {
			t.Error("unexpected role value")
		}

		if c.SessionDuration != credentials.SessionTokenDefaultDuration {
			t.Error("unexpected session duration")
		}

		if c.RoleDuration != credentials.AssumeRoleDefaultDuration {
			t.Error("unexpected role duration")
		}
	})

	t.Run("source", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("source")
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-east-1" {
			t.Error("unexpected region")
		}

		if c.MfaSerial != "ABCDEFG" {
			t.Error("unexpected mfa serial value")
		}

		if len(c.RoleArn) > 0 {
			t.Error("unexpected role value")
		}

		if c.SessionDuration == 0 {
			t.Error("unexpected session duration")
		}

		if c.RoleDuration != credentials.AssumeRoleDefaultDuration {
			t.Error("unexpected role duration")
		}
	})

	t.Run("role default source", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("role1")
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-west-1" {
			t.Error("unexpected region")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("unexpected mfa serial")
		}

		if c.RoleArn != "role1" {
			t.Error("unexpected role value")
		}

		if c.SessionDuration != credentials.SessionTokenDefaultDuration {
			t.Error("unexpected session duration")
		}

		if c.RoleDuration != credentials.AssumeRoleDefaultDuration {
			t.Error("unexpected role duration")
		}
	})

	t.Run("role non-default source", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("role2")
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-east-1" {
			t.Error("unexpected region")
		}

		if c.MfaSerial != "ABCDEFG" {
			t.Error("unexpected mfa serial value")
		}

		if c.RoleArn != "role2" {
			t.Error("unexpected role value")
		}

		if c.SessionDuration == 0 {
			t.Error("unexpected session duration")
		}

		if c.RoleDuration != credentials.AssumeRoleDefaultDuration {
			t.Error("unexpected role duration")
		}
	})

	t.Run("role bad source", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("role3")
		if err != nil {
			t.Error(err)
			return
		}

		if c.SourceProfile != "other" {
			t.Error("source profile mismatch")
		}
	})

	t.Run("arn", func(t *testing.T) {
		os.Setenv(MfaSerialEnvVar, "654321")
		defer os.Unsetenv(MfaSerialEnvVar)

		arn := "arn:aws:iam::0123456789012:role/Admin"

		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig(arn)
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-west-1" {
			t.Error("unexpected region")
		}

		if c.RoleArn != arn {
			t.Error("bad role value")
		}

		if c.MfaSerial != "654321" {
			t.Error("bad mfa value")
		}
	})

	t.Run("bad arn", func(t *testing.T) {

		arn := "arn:aws:iam::0123456789012:user/Admin"

		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		_, err = r.ResolveConfig(arn)
		if err == nil {
			t.Error("did not receive expected error")
			return
		}
	})
}

func TestConfigResolver_ResolveDefaultConfig(t *testing.T) {
	os.Setenv(config.ConfigFileEnvVar, "test/config")
	defer os.Unsetenv(config.ConfigFileEnvVar)

	r, err := NewConfigResolver(nil)
	if err != nil {
		t.Error(err)
		return
	}

	t.Run("default", func(t *testing.T) {
		c, err := r.ResolveDefaultConfig()
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-west-1" {
			t.Error("unexpected region")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("mfa serial was set for default profile")
		}
	})

	t.Run("good env var", func(t *testing.T) {
		os.Setenv(DefaultProfileEnvVar, "source")
		defer os.Unsetenv(DefaultProfileEnvVar)

		c, err := r.ResolveDefaultConfig()
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-east-1" {
			t.Error("unexpected region")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("mfa serial was set for default profile")
		}

		if c.SessionDuration == 0 {
			t.Error("bad session duration")
		}
	})

	t.Run("bad env var", func(t *testing.T) {
		os.Setenv(DefaultProfileEnvVar, "bad")
		defer os.Unsetenv(DefaultProfileEnvVar)

		_, err := r.ResolveDefaultConfig()
		if err == nil {
			t.Error("did not receive expected error")
			return
		}
	})
}

func TestConfigResolver_ResolveProfileConfig(t *testing.T) {
	os.Setenv(config.ConfigFileEnvVar, "test/config")
	defer os.Unsetenv(config.ConfigFileEnvVar)

	r, err := NewConfigResolver(nil)
	if err != nil {
		t.Error(err)
		return
	}

	t.Run("profile exists", func(t *testing.T) {
		c, err := r.ResolveProfileConfig("role1")
		if err != nil {
			t.Error(err)
			return
		}

		if c.RoleArn != "role1" {
			t.Error("role mismatch")
		}
	})

	t.Run("empty profile", func(t *testing.T) {
		// Will look up default profile
		_, err := r.ResolveProfileConfig("")
		if err != nil {
			t.Error(err)
			return
		}
	})

	t.Run("bad profile", func(t *testing.T) {
		_, err := r.ResolveProfileConfig("bogus")
		if err == nil {
			t.Error("did not receive expected error")
			return
		}
	})
}

func TestConfigResolver_ResolveEnvConfig(t *testing.T) {
	env := make(map[string]string)
	env[RegionEnvVar] = "us-east-2"
	env[ProfileEnvVar] = "profile"
	env[SessionDurationEnvVar] = "1m"
	env[config.CredentialsFileEnvVar] = "test/config"

	for k, v := range env {
		os.Setenv(k, v)
		defer os.Unsetenv(k)
	}

	r, err := NewConfigResolver(nil)
	if err != nil {
		t.Error(err)
		return
	}

	t.Run("good", func(t *testing.T) {
		c, err := r.ResolveEnvConfig()
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != env[RegionEnvVar] {
			t.Error("bad region")
		}

		if c.SessionDuration != 1*time.Minute {
			t.Error("bad session duration")
		}

		if len(c.RoleArn) > 0 {
			t.Error("unexpected role arn")
		}
	})

	t.Run("bad duration", func(t *testing.T) {
		os.Setenv(RoleDurationEnvVar, "ab")
		defer os.Unsetenv(RoleDurationEnvVar)

		_, err := r.ResolveEnvConfig()
		if err == nil {
			t.Error("did not see expected error")
			return
		}
	})
}

func TestMergeConfig(t *testing.T) {
	t.Run("all nil", func(t *testing.T) {
		c := MergeConfig(nil)
		if c == nil {
			t.Error("nil config")
			return
		}

		if len(c.Region) > 0 {
			t.Error("unexpected region value")
		}

		if len(c.RoleArn) > 0 {
			t.Error("unexpected role arn")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("unexpected mfa serial")
		}

		if len(c.ExternalID) > 0 {
			t.Error("unexpected external id")
		}

		if c.SessionDuration > 0 {
			t.Error("unexpeted session duration")
		}

		if c.RoleDuration > 0 {
			t.Error("unexpected role duration")
		}
	})

	t.Run("good", func(t *testing.T) {
		c := MergeConfig(
			nil,
			&AwsConfig{Region: "us-east-1"},
			&AwsConfig{MfaSerial: "123456"},
			nil,
			&AwsConfig{Region: "us-east-2", RoleArn: "my-role"})

		if c == nil {
			t.Error("unexpected nil config")
			return
		}

		if c.Region != "us-east-2" {
			t.Error("bad region")
		}

		if c.MfaSerial != "123456" {
			t.Error("bad mfa serial")
		}

		if c.RoleArn != "my-role" {
			t.Error("bad role")
		}
	})
}
<commit_msg>add test coverage for config module<commit_after>package config

import (
	"github.com/mmmorris1975/aws-config/config"
	"github.com/mmmorris1975/aws-runas/lib/credentials"
	"github.com/mmmorris1975/simple-logger"
	"os"
	"testing"
	"time"
)

func TestNewConfigResolver(t *testing.T) {
	t.Run("nil config", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		if r == nil {
			t.Error("unexpected nil resolver")
			return
		}

		if r.file == nil || len(r.file.Path) < 1 {
			t.Error("config file not set")
		}
	})

	t.Run("good", func(t *testing.T) {
		r, err := NewConfigResolver(&AwsConfig{Region: "us-west-2"})
		if err != nil {
			t.Error(err)
			return
		}

		if r.userConfig.Region != "us-west-2" {
			t.Error("unexpected region")
		}
	})

	t.Run("config file env var", func(t *testing.T) {
		os.Setenv(config.ConfigFileEnvVar, "test/config")
		defer os.Unsetenv(config.ConfigFileEnvVar)

		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		if r.file == nil || r.file.Path != "test/config" {
			t.Error("bad file name")
		}
	})
}

func TestConfigResolver_ResolveConfig(t *testing.T) {
	os.Setenv(config.ConfigFileEnvVar, "test/config")
	defer os.Unsetenv(config.ConfigFileEnvVar)

	t.Run("default only", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("")
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-west-1" {
			t.Error("unexpected region")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("unexpected mfa serial")
		}

		if len(c.RoleArn) > 0 {
			t.Error("unexpected role value")
		}

		if c.SessionDuration != credentials.SessionTokenDefaultDuration {
			t.Error("unexpected session duration")
		}

		if c.RoleDuration != credentials.AssumeRoleDefaultDuration {
			t.Error("unexpected role duration")
		}
	})

	t.Run("source", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("source")
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-east-1" {
			t.Error("unexpected region")
		}

		if c.MfaSerial != "ABCDEFG" {
			t.Error("unexpected mfa serial value")
		}

		if len(c.RoleArn) > 0 {
			t.Error("unexpected role value")
		}

		if c.SessionDuration == 0 {
			t.Error("unexpected session duration")
		}

		if c.RoleDuration != credentials.AssumeRoleDefaultDuration {
			t.Error("unexpected role duration")
		}
	})

	t.Run("role default source", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("role1")
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-west-1" {
			t.Error("unexpected region")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("unexpected mfa serial")
		}

		if c.RoleArn != "role1" {
			t.Error("unexpected role value")
		}

		if c.SessionDuration != credentials.SessionTokenDefaultDuration {
			t.Error("unexpected session duration")
		}

		if c.RoleDuration != credentials.AssumeRoleDefaultDuration {
			t.Error("unexpected role duration")
		}
	})

	t.Run("role non-default source", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("role2")
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-east-1" {
			t.Error("unexpected region")
		}

		if c.MfaSerial != "ABCDEFG" {
			t.Error("unexpected mfa serial value")
		}

		if c.RoleArn != "role2" {
			t.Error("unexpected role value")
		}

		if c.SessionDuration == 0 {
			t.Error("unexpected session duration")
		}

		if c.RoleDuration != credentials.AssumeRoleDefaultDuration {
			t.Error("unexpected role duration")
		}
	})

	t.Run("role bad source", func(t *testing.T) {
		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig("role3")
		if err != nil {
			t.Error(err)
			return
		}

		if c.SourceProfile != "other" {
			t.Error("source profile mismatch")
		}
	})

	t.Run("arn", func(t *testing.T) {
		os.Setenv(MfaSerialEnvVar, "654321")
		defer os.Unsetenv(MfaSerialEnvVar)

		arn := "arn:aws:iam::0123456789012:role/Admin"

		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		c, err := r.ResolveConfig(arn)
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-west-1" {
			t.Error("unexpected region")
		}

		if c.RoleArn != arn {
			t.Error("bad role value")
		}

		if c.MfaSerial != "654321" {
			t.Error("bad mfa value")
		}
	})

	t.Run("bad arn", func(t *testing.T) {

		arn := "arn:aws:iam::0123456789012:user/Admin"

		r, err := NewConfigResolver(nil)
		if err != nil {
			t.Error(err)
			return
		}

		_, err = r.ResolveConfig(arn)
		if err == nil {
			t.Error("did not receive expected error")
			return
		}
	})
}

func TestConfigResolver_ResolveDefaultConfig(t *testing.T) {
	os.Setenv(config.ConfigFileEnvVar, "test/config")
	defer os.Unsetenv(config.ConfigFileEnvVar)

	r, err := NewConfigResolver(nil)
	if err != nil {
		t.Error(err)
		return
	}

	t.Run("default", func(t *testing.T) {
		c, err := r.ResolveDefaultConfig()
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-west-1" {
			t.Error("unexpected region")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("mfa serial was set for default profile")
		}
	})

	t.Run("good env var", func(t *testing.T) {
		os.Setenv(DefaultProfileEnvVar, "source")
		defer os.Unsetenv(DefaultProfileEnvVar)

		c, err := r.ResolveDefaultConfig()
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != "us-east-1" {
			t.Error("unexpected region")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("mfa serial was set for default profile")
		}

		if c.SessionDuration == 0 {
			t.Error("bad session duration")
		}
	})

	t.Run("bad env var", func(t *testing.T) {
		os.Setenv(DefaultProfileEnvVar, "bad")
		defer os.Unsetenv(DefaultProfileEnvVar)

		_, err := r.ResolveDefaultConfig()
		if err == nil {
			t.Error("did not receive expected error")
			return
		}
	})
}

func TestConfigResolver_ResolveProfileConfig(t *testing.T) {
	os.Setenv(config.ConfigFileEnvVar, "test/config")
	defer os.Unsetenv(config.ConfigFileEnvVar)

	r, err := NewConfigResolver(nil)
	if err != nil {
		t.Error(err)
		return
	}

	t.Run("profile exists", func(t *testing.T) {
		c, err := r.ResolveProfileConfig("role1")
		if err != nil {
			t.Error(err)
			return
		}

		if c.RoleArn != "role1" {
			t.Error("role mismatch")
		}
	})

	t.Run("empty profile", func(t *testing.T) {
		// Will look up default profile
		_, err := r.ResolveProfileConfig("")
		if err != nil {
			t.Error(err)
			return
		}
	})

	t.Run("bad profile", func(t *testing.T) {
		_, err := r.ResolveProfileConfig("bogus")
		if err == nil {
			t.Error("did not receive expected error")
			return
		}
	})
}

func TestConfigResolver_ResolveEnvConfig(t *testing.T) {
	env := make(map[string]string)
	env[RegionEnvVar] = "us-east-2"
	env[ProfileEnvVar] = "profile"
	env[SessionDurationEnvVar] = "1m"
	env[config.CredentialsFileEnvVar] = "test/config"

	for k, v := range env {
		os.Setenv(k, v)
		defer os.Unsetenv(k)
	}

	r, err := NewConfigResolver(nil)
	if err != nil {
		t.Error(err)
		return
	}

	t.Run("good", func(t *testing.T) {
		c, err := r.ResolveEnvConfig()
		if err != nil {
			t.Error(err)
			return
		}

		if c.Region != env[RegionEnvVar] {
			t.Error("bad region")
		}

		if c.SessionDuration != 1*time.Minute {
			t.Error("bad session duration")
		}

		if len(c.RoleArn) > 0 {
			t.Error("unexpected role arn")
		}
	})

	t.Run("bad session duration", func(t *testing.T) {
		os.Setenv(SessionDurationEnvVar, "ab")
		defer os.Unsetenv(SessionDurationEnvVar)

		_, err := r.ResolveEnvConfig()
		if err == nil {
			t.Error("did not see expected error")
			return
		}
	})

	t.Run("bad role duration", func(t *testing.T) {
		os.Setenv(RoleDurationEnvVar, "ab")
		defer os.Unsetenv(RoleDurationEnvVar)

		_, err := r.ResolveEnvConfig()
		if err == nil {
			t.Error("did not see expected error")
			return
		}
	})
}

func TestMergeConfig(t *testing.T) {
	t.Run("all nil", func(t *testing.T) {
		c := MergeConfig(nil)
		if c == nil {
			t.Error("nil config")
			return
		}

		if len(c.Region) > 0 {
			t.Error("unexpected region value")
		}

		if len(c.RoleArn) > 0 {
			t.Error("unexpected role arn")
		}

		if len(c.MfaSerial) > 0 {
			t.Error("unexpected mfa serial")
		}

		if len(c.ExternalID) > 0 {
			t.Error("unexpected external id")
		}

		if c.SessionDuration > 0 {
			t.Error("unexpeted session duration")
		}

		if c.RoleDuration > 0 {
			t.Error("unexpected role duration")
		}
	})

	t.Run("good", func(t *testing.T) {
		c := MergeConfig(
			nil,
			&AwsConfig{Region: "us-east-1"},
			&AwsConfig{MfaSerial: "123456"},
			nil,
			&AwsConfig{Region: "us-east-2", RoleArn: "my-role"})

		if c == nil {
			t.Error("unexpected nil config")
			return
		}

		if c.Region != "us-east-2" {
			t.Error("bad region")
		}

		if c.MfaSerial != "123456" {
			t.Error("bad mfa serial")
		}

		if c.RoleArn != "my-role" {
			t.Error("bad role")
		}
	})
}

func TestConfigResolver_WithLogger(t *testing.T) {
	r := new(configResolver).WithLogger(simple_logger.StdLogger)
	if r.log == nil {
		t.Error("unexpected nil logger")
	}
}

func ExampleDebugNilLogger() {
	r := new(configResolver)
	r.debug("test")
	// Output:
	//
}

func ExampleDebug() {
	l := simple_logger.NewLogger(os.Stdout, "", 0)
	l.SetLevel(simple_logger.DEBUG)
	r := new(configResolver).WithLogger(l)
	r.debug("test")
	// Output:
	// DEBUG test
}
<|endoftext|>


<commit_before>// Copyright 2015 Square 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 forecast

import "fmt"

// Forecast uses variants of the Holt-Winters model for data.
// Multiplicative model: Assume y(t) = (a + bt)*f(t)  where f(t) is a periodic function with known period L and mean 1.
// Additive model:       Assume y(t) = a + bt + f(t)  where f(t) is a periodic function with known period L and mean 0.
// "Generalized" model:  Assume y(t) = a(t) + b(t)x where a(t) and b(t) are periodic functions with known period L.

// EstimateGeneralizedHoltWintersModel estimates the corresponding model (as described above)
// given the data and the period of the model parameters. There must be at least 2 complete periods of data,
// but to be even slightly effective, more data MUST be provided.
// The data at the end of the array will be ignored if there is an incomplete period.
// The input slice is not modified.
func TrainGeneralizedHoltWintersModel(ys []float64, period int) (Model, error) {
	if period <= 0 {
		return GeneralizedHoltWintersModel{}, fmt.Errorf("Generalized Holt-Winters model expects a positive period")
	}
	count := len(ys) / period
	alphas := make([]float64, period)
	betas := make([]float64, period)
	for i := range alphas {
		data := make([]float64, count)
		for j := range data {
			data[j] = ys[i+j*period]
		}
		alphas[i], betas[i] = LinearRegression(data)
	}
	return GeneralizedHoltWintersModel{
		Alphas: alphas,
		Betas:  betas,
	}, nil
}

// Trains a multiplicative Holt-Winters model on the given data (using the given period).
// The input slice is not modified.
func TrainMultiplicativeHoltWintersModel(ys []float64, period int) (Model, error) {
	if period <= 0 {
		return MultiplicativeHoltWintersModel{}, fmt.Errorf("Training the multiplicative Holt-Winters model requires a positive period.") // TODO: structured error
	}
	if len(ys) < period*3 {
		return MultiplicativeHoltWintersModel{}, fmt.Errorf("Good results with the Multiplicative Holt-Winters model training require at least 3 periods of data.") // TODO: structured error
	}
	// First we will find the "beta" parameter (the average trend).
	// To do this, we require
	periodMeans := make([]float64, len(ys)/period)
	for i := range periodMeans {
		sum := 0.0
		for t := 0; t < period; t++ {
			sum += ys[i*period+t]
		}
		periodMeans[i] = sum / float64(period)
	}
	// periodMeans now contains the mean of each period of the data.
	_, beta := LinearRegression(periodMeans)
	// This beta is the overall trend of the data, but it needs to be rescaled:
	beta /= float64(period)

	// Next we calculate the "untrended" data, by subtracting beta*t from each point:

	zs := make([]float64, len(ys))
	for i := range ys {
		zs[i] = ys[i] - beta*float64(i)
	}

	// Now we make the following observation. Consider g(t) = f(t) - bt. Then we have
	// g(t) = f(t) - bt = S(t)(a + bt) - bt = a S(t) + b S(t) t - b t = a S(t) + b (S(t) - 1) t
	// So now compute g(np + t) - g(mp + t), where 0 <= t < p, and n, m are two integer. So,
	// g(np+t) - g(mp+t) = aS(np+t) + b(S(np+t)-1)(np+t) - aS(mp+t) + b(S(mp+t)-1)(mp+t).
	// But S is a periodic function, so we can simplify:
	// g(np+t) - g(mp+t) = aS(t) + b(S(t)-1)(np+t) - aS(t) + b(S(t)-1)(mp+t)
	// and a bit more expansion and factoring gives us
	// g(np+t) - g(mp+t) = b(S(t) - 1)(np+t - mp - t) = bp(S(t)-1)(n-m).
	// Thus, by solving for S(t), we can see that
	// (g(np+t) - g(mp+t)) / (bp (n-m)) = S(t) - 1, so
	// S(t) = 1 + (g(np+t) - g(mp+t)) / (bp (n-m))

	// However, this gives us n^2 equations where we have n periods of data. Therefore, we'll use the average across all of these.

	season := make([]float64, period)

	for t := 0; t < period; t++ {
		gs := make([]float64, len(ys)/period)
		for n := range gs {
			gs[n] = zs[n*period+t] // For convenience
		}
		sumS := 0.0
		countS := 0
		for n := range gs {
			for m := range gs {
				if n == m {
					continue
				}
				value := 1 + (gs[n]-gs[m])/beta/float64(period)/float64(n-m)
				sumS += value
				countS++
			}
		}
		season[t] = sumS / float64(countS)
	}

	//log.Printf("Calculated season: %+v", season)

	// Lastly, we'll need to compute 'alpha'. We do this be "deseasonalizing" zs.

	// g(t) = a S(t) + b (S(t) - 1) t
	// So we have to subtract out b(S(t)-1)t
	// and then divide by S(t):
	// a = (g(t) - b(S(t)-1)t) / S(t)

	ds := make([]float64, len(zs))
	for i := range zs {
		s := season[mod(i, period)]
		ds[i] = (zs[i] - beta*(s-1)*float64(i)) / s
	}

	alpha := 0.0
	for i := range ds {
		alpha += ds[i]
	}
	alpha /= float64(len(ds))

	//log.Printf("Calculated alpha: %f", alpha)

	return MultiplicativeHoltWintersModel{
		season: season,
		alpha:  alpha,
		beta:   beta,
	}, nil
}
<commit_msg>check for NaN before accumulating in forecast<commit_after>// Copyright 2015 Square 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 forecast

import (
	"fmt"
	"math"
)

// Forecast uses variants of the Holt-Winters model for data.
// Multiplicative model: Assume y(t) = (a + bt)*f(t)  where f(t) is a periodic function with known period L and mean 1.
// Additive model:       Assume y(t) = a + bt + f(t)  where f(t) is a periodic function with known period L and mean 0.
// "Generalized" model:  Assume y(t) = a(t) + b(t)x where a(t) and b(t) are periodic functions with known period L.

// EstimateGeneralizedHoltWintersModel estimates the corresponding model (as described above)
// given the data and the period of the model parameters. There must be at least 2 complete periods of data,
// but to be even slightly effective, more data MUST be provided.
// The data at the end of the array will be ignored if there is an incomplete period.
// The input slice is not modified.
func TrainGeneralizedHoltWintersModel(ys []float64, period int) (Model, error) {
	if period <= 0 {
		return GeneralizedHoltWintersModel{}, fmt.Errorf("Generalized Holt-Winters model expects a positive period")
	}
	count := len(ys) / period
	alphas := make([]float64, period)
	betas := make([]float64, period)
	for i := range alphas {
		data := make([]float64, count)
		for j := range data {
			data[j] = ys[i+j*period]
		}
		alphas[i], betas[i] = LinearRegression(data)
	}
	return GeneralizedHoltWintersModel{
		Alphas: alphas,
		Betas:  betas,
	}, nil
}

// Trains a multiplicative Holt-Winters model on the given data (using the given period).
// The input slice is not modified.
func TrainMultiplicativeHoltWintersModel(ys []float64, period int) (Model, error) {
	if period <= 0 {
		return MultiplicativeHoltWintersModel{}, fmt.Errorf("Training the multiplicative Holt-Winters model requires a positive period.") // TODO: structured error
	}
	if len(ys) < period*3 {
		return MultiplicativeHoltWintersModel{}, fmt.Errorf("Good results with the Multiplicative Holt-Winters model training require at least 3 periods of data.") // TODO: structured error
	}
	// First we will find the "beta" parameter (the average trend).
	// To do this, we require
	periodMeans := make([]float64, len(ys)/period)
	for i := range periodMeans {
		sum := 0.0
		for t := 0; t < period; t++ {
			sum += ys[i*period+t]
		}
		periodMeans[i] = sum / float64(period)
	}
	// periodMeans now contains the mean of each period of the data.
	_, beta := LinearRegression(periodMeans)
	// This beta is the overall trend of the data, but it needs to be rescaled:
	beta /= float64(period)

	// Next we calculate the "untrended" data, by subtracting beta*t from each point:

	zs := make([]float64, len(ys))
	for i := range ys {
		zs[i] = ys[i] - beta*float64(i)
	}

	// Now we make the following observation. Consider g(t) = f(t) - bt. Then we have
	// g(t) = f(t) - bt = S(t)(a + bt) - bt = a S(t) + b S(t) t - b t = a S(t) + b (S(t) - 1) t
	// So now compute g(np + t) - g(mp + t), where 0 <= t < p, and n, m are two integer. So,
	// g(np+t) - g(mp+t) = aS(np+t) + b(S(np+t)-1)(np+t) - aS(mp+t) + b(S(mp+t)-1)(mp+t).
	// But S is a periodic function, so we can simplify:
	// g(np+t) - g(mp+t) = aS(t) + b(S(t)-1)(np+t) - aS(t) + b(S(t)-1)(mp+t)
	// and a bit more expansion and factoring gives us
	// g(np+t) - g(mp+t) = b(S(t) - 1)(np+t - mp - t) = bp(S(t)-1)(n-m).
	// Thus, by solving for S(t), we can see that
	// (g(np+t) - g(mp+t)) / (bp (n-m)) = S(t) - 1, so
	// S(t) = 1 + (g(np+t) - g(mp+t)) / (bp (n-m))

	// However, this gives us n^2 equations where we have n periods of data. Therefore, we'll use the average across all of these.

	season := make([]float64, period)

	for t := 0; t < period; t++ {
		gs := make([]float64, len(ys)/period)
		for n := range gs {
			gs[n] = zs[n*period+t] // For convenience
		}
		sumS := 0.0
		countS := 0
		for n := range gs {
			for m := range gs {
				if n <= m {
					continue
				}
				value := 1 + (gs[n]-gs[m])/beta/float64(period)/float64(n-m)
				if math.IsNaN(value) {
					continue
				}
				sumS += value
				countS++
			}
		}
		season[t] = sumS / float64(countS)
	}

	//log.Printf("Calculated season: %+v", season)

	// Lastly, we'll need to compute 'alpha'. We do this be "deseasonalizing" zs.

	// g(t) = a S(t) + b (S(t) - 1) t
	// So we have to subtract out b(S(t)-1)t
	// and then divide by S(t):
	// a = (g(t) - b(S(t)-1)t) / S(t)

	ds := make([]float64, len(zs))
	for i := range zs {
		s := season[mod(i, period)]
		ds[i] = (zs[i] - beta*(s-1)*float64(i)) / s
	}

	alpha := 0.0
	for i := range ds {
		alpha += ds[i]
	}
	alpha /= float64(len(ds))

	//log.Printf("Calculated alpha: %f", alpha)

	return MultiplicativeHoltWintersModel{
		season: season,
		alpha:  alpha,
		beta:   beta,
	}, nil
}
<|endoftext|>


<commit_before><commit_msg>console_test: use lowercase for local functions<commit_after><|endoftext|>


<commit_before>// Thread safe engine for MyMySQL
//
// In contrast to native engine:
// - one connection can be used by multiple gorutines,
// - if connection is idle pings are sent to the server (once per minute) to
//   avoid timeout.
//
// See documentation of mymysql/native for details
package thrsafe

import (
	"sync"
	//"log"
	"github.com/ziutek/mymysql/mysql"
	_ "github.com/ziutek/mymysql/native"
	"io"
	"time"
)

type Conn struct {
	mysql.Conn
	mutex *sync.Mutex

	stopPinger chan struct{}
	lastUsed   time.Time
}

func (c *Conn) lock() {
	//log.Println(c, ":: lock @", c.mutex)
	c.mutex.Lock()
	c.lastUsed = time.Now()
}

func (c *Conn) unlock() {
	//log.Println(c, ":: unlock @", c.mutex)
	c.mutex.Unlock()
}

type Result struct {
	mysql.Result
	conn *Conn
}

type Stmt struct {
	mysql.Stmt
	conn *Conn
}

type Transaction struct {
	*Conn
	conn *Conn
}

func New(proto, laddr, raddr, user, passwd string, db ...string) mysql.Conn {
	return &Conn{
		Conn:       orgNew(proto, laddr, raddr, user, passwd, db...),
		mutex:      new(sync.Mutex),
	}
}

func (c *Conn) pinger() {
	c.stopPinger = make(chan struct{})
	defer func() { c.stopPinger = nil }()

	const to = 60 * time.Second
	sleep := to
	for {
		timer := time.After(sleep)
		select {
		case <-c.stopPinger:
			return
		case t := <-timer:
			sleep := to - t.Sub(c.lastUsed)
			if sleep <= 0 {
				if c.Ping() != nil {
					return
				}
				sleep = to
			}
		}
	}
}

func (c *Conn) Connect() error {
	//log.Println("Connect")
	c.lock()
	defer c.unlock()
	go c.pinger()
	return c.Conn.Connect()
}

func (c *Conn) Close() error {
	//log.Println("Close")
	close(c.stopPinger) // Stop pinger before lock connection
	c.lock()
	defer c.unlock()
	return c.Conn.Close()
}

func (c *Conn) Reconnect() error {
	//log.Println("Reconnect")
	c.lock()
	defer c.unlock()
	if c.stopPinger == nil {
		go c.pinger()
	}
	return c.Conn.Reconnect()
}

func (c *Conn) Use(dbname string) error {
	//log.Println("Use")
	c.lock()
	defer c.unlock()
	return c.Conn.Use(dbname)
}

func (c *Conn) Start(sql string, params ...interface{}) (mysql.Result, error) {
	//log.Println("Start")
	c.lock()
	res, err := c.Conn.Start(sql, params...)
	// Unlock if error or OK result (which doesn't provide any fields)
	if err != nil {
		c.unlock()
		return nil, err
	}
	if res.StatusOnly() {
		c.unlock()
	}
	return &Result{Result: res, conn: c}, err
}

func (res *Result) ScanRow(row mysql.Row) error {
	//log.Println("ScanRow")
	err := res.Result.ScanRow(row)
	if err != nil && (err != io.EOF || !res.StatusOnly() && !res.MoreResults()) {
		res.conn.unlock()
	}
	return err
}

func (res *Result) GetRow() (mysql.Row, error) {
	return mysql.GetRow(res)
}

func (res *Result) NextResult() (mysql.Result, error) {
	//log.Println("NextResult")
	next, err := res.Result.NextResult()
	if err != nil {
		return nil, err
	}
	return &Result{next, res.conn}, nil
}

func (c *Conn) Ping() error {
	c.lock()
	defer c.unlock()
	return c.Conn.Ping()
}

func (c *Conn) Prepare(sql string) (mysql.Stmt, error) {
	//log.Println("Prepare")
	c.lock()
	defer c.unlock()
	stmt, err := c.Conn.Prepare(sql)
	if err != nil {
		return nil, err
	}
	return &Stmt{Stmt: stmt, conn: c}, nil
}

func (stmt *Stmt) Run(params ...interface{}) (mysql.Result, error) {
	//log.Println("Run")
	stmt.conn.lock()
	res, err := stmt.Stmt.Run(params...)
	// Unlock if error or OK result (which doesn't provide any fields)
	if err != nil {
		stmt.conn.unlock()
		return nil, err
	}
	if res.StatusOnly() {
		stmt.conn.unlock()
	}
	return &Result{Result: res, conn: stmt.conn}, nil
}

func (stmt *Stmt) Delete() error {
	//log.Println("Delete")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.Delete()
}

func (stmt *Stmt) Reset() error {
	//log.Println("Reset")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.Reset()
}

func (stmt *Stmt) SendLongData(pnum int, data interface{}, pkt_size int) error {
	//log.Println("SendLongData")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.SendLongData(pnum, data, pkt_size)
}

func (c *Conn) Query(sql string, params ...interface{}) ([]mysql.Row, mysql.Result, error) {
	return mysql.Query(c, sql, params...)
}

func (stmt *Stmt) Exec(params ...interface{}) ([]mysql.Row, mysql.Result, error) {
	return mysql.Exec(stmt, params...)
}

func (res *Result) End() error {
	return mysql.End(res)
}

func (res *Result) GetRows() ([]mysql.Row, error) {
	return mysql.GetRows(res)
}

// Begins a new transaction. No any other thread can send command on this
// connection until Commit or Rollback will be called.
// Periodical pinging the server is disabled during transaction.

func (c *Conn) Begin() (mysql.Transaction, error) {
	//log.Println("Begin")
	c.lock()
	tr := Transaction{
		&Conn{Conn: c.Conn, mutex: new(sync.Mutex)},
		c,
	}
	_, err := c.Conn.Start("START TRANSACTION")
	if err != nil {
		c.unlock()
		return nil, err
	}
	return &tr, nil
}

func (tr *Transaction) end(cr string) error {
	tr.lock()
	_, err := tr.conn.Conn.Start(cr)
	tr.conn.unlock()
	// Invalidate this transaction
	m := tr.Conn.mutex
	tr.Conn = nil
	tr.conn = nil
	m.Unlock() // One goorutine which still uses this transaction will panic
	return err
}

func (tr *Transaction) Commit() error {
	//log.Println("Commit")
	return tr.end("COMMIT")
}

func (tr *Transaction) Rollback() error {
	//log.Println("Rollback")
	return tr.end("ROLLBACK")
}

func (tr *Transaction) Do(st mysql.Stmt) mysql.Stmt {
	if s, ok := st.(*Stmt); ok && s.conn == tr.conn {
		// Returns new statement which uses statement mutexes
		return &Stmt{s.Stmt, tr.Conn}
	}
	panic("Transaction and statement doesn't belong to the same connection")
}

var orgNew func(proto, laddr, raddr, user, passwd string, db ...string) mysql.Conn

func init() {
	orgNew = mysql.New
	mysql.New = New
}
<commit_msg>Move a lastUsed update from lock to unlock<commit_after>// Thread safe engine for MyMySQL
//
// In contrast to native engine:
// - one connection can be used by multiple gorutines,
// - if connection is idle pings are sent to the server (once per minute) to
//   avoid timeout.
//
// See documentation of mymysql/native for details
package thrsafe

import (
	"sync"
	//"log"
	"github.com/ziutek/mymysql/mysql"
	_ "github.com/ziutek/mymysql/native"
	"io"
	"time"
)

type Conn struct {
	mysql.Conn
	mutex *sync.Mutex

	stopPinger chan struct{}
	lastUsed   time.Time
}

func (c *Conn) lock() {
	//log.Println(c, ":: lock @", c.mutex)
	c.mutex.Lock()
}

func (c *Conn) unlock() {
	//log.Println(c, ":: unlock @", c.mutex)
	c.lastUsed = time.Now()
	c.mutex.Unlock()
}

type Result struct {
	mysql.Result
	conn *Conn
}

type Stmt struct {
	mysql.Stmt
	conn *Conn
}

type Transaction struct {
	*Conn
	conn *Conn
}

func New(proto, laddr, raddr, user, passwd string, db ...string) mysql.Conn {
	return &Conn{
		Conn:       orgNew(proto, laddr, raddr, user, passwd, db...),
		mutex:      new(sync.Mutex),
	}
}

func (c *Conn) pinger() {
	c.stopPinger = make(chan struct{})
	defer func() { c.stopPinger = nil }()

	const to = 60 * time.Second
	sleep := to
	for {
		timer := time.After(sleep)
		select {
		case <-c.stopPinger:
			return
		case t := <-timer:
			sleep := to - t.Sub(c.lastUsed)
			if sleep <= 0 {
				if c.Ping() != nil {
					return
				}
				sleep = to
			}
		}
	}
}

func (c *Conn) Connect() error {
	//log.Println("Connect")
	c.lock()
	defer c.unlock()
	go c.pinger()
	return c.Conn.Connect()
}

func (c *Conn) Close() error {
	//log.Println("Close")
	close(c.stopPinger) // Stop pinger before lock connection
	c.lock()
	defer c.unlock()
	return c.Conn.Close()
}

func (c *Conn) Reconnect() error {
	//log.Println("Reconnect")
	c.lock()
	defer c.unlock()
	if c.stopPinger == nil {
		go c.pinger()
	}
	return c.Conn.Reconnect()
}

func (c *Conn) Use(dbname string) error {
	//log.Println("Use")
	c.lock()
	defer c.unlock()
	return c.Conn.Use(dbname)
}

func (c *Conn) Start(sql string, params ...interface{}) (mysql.Result, error) {
	//log.Println("Start")
	c.lock()
	res, err := c.Conn.Start(sql, params...)
	// Unlock if error or OK result (which doesn't provide any fields)
	if err != nil {
		c.unlock()
		return nil, err
	}
	if res.StatusOnly() {
		c.unlock()
	}
	return &Result{Result: res, conn: c}, err
}

func (res *Result) ScanRow(row mysql.Row) error {
	//log.Println("ScanRow")
	err := res.Result.ScanRow(row)
	if err != nil && (err != io.EOF || !res.StatusOnly() && !res.MoreResults()) {
		res.conn.unlock()
	}
	return err
}

func (res *Result) GetRow() (mysql.Row, error) {
	return mysql.GetRow(res)
}

func (res *Result) NextResult() (mysql.Result, error) {
	//log.Println("NextResult")
	next, err := res.Result.NextResult()
	if err != nil {
		return nil, err
	}
	return &Result{next, res.conn}, nil
}

func (c *Conn) Ping() error {
	c.lock()
	defer c.unlock()
	return c.Conn.Ping()
}

func (c *Conn) Prepare(sql string) (mysql.Stmt, error) {
	//log.Println("Prepare")
	c.lock()
	defer c.unlock()
	stmt, err := c.Conn.Prepare(sql)
	if err != nil {
		return nil, err
	}
	return &Stmt{Stmt: stmt, conn: c}, nil
}

func (stmt *Stmt) Run(params ...interface{}) (mysql.Result, error) {
	//log.Println("Run")
	stmt.conn.lock()
	res, err := stmt.Stmt.Run(params...)
	// Unlock if error or OK result (which doesn't provide any fields)
	if err != nil {
		stmt.conn.unlock()
		return nil, err
	}
	if res.StatusOnly() {
		stmt.conn.unlock()
	}
	return &Result{Result: res, conn: stmt.conn}, nil
}

func (stmt *Stmt) Delete() error {
	//log.Println("Delete")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.Delete()
}

func (stmt *Stmt) Reset() error {
	//log.Println("Reset")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.Reset()
}

func (stmt *Stmt) SendLongData(pnum int, data interface{}, pkt_size int) error {
	//log.Println("SendLongData")
	stmt.conn.lock()
	defer stmt.conn.unlock()
	return stmt.Stmt.SendLongData(pnum, data, pkt_size)
}

func (c *Conn) Query(sql string, params ...interface{}) ([]mysql.Row, mysql.Result, error) {
	return mysql.Query(c, sql, params...)
}

func (stmt *Stmt) Exec(params ...interface{}) ([]mysql.Row, mysql.Result, error) {
	return mysql.Exec(stmt, params...)
}

func (res *Result) End() error {
	return mysql.End(res)
}

func (res *Result) GetRows() ([]mysql.Row, error) {
	return mysql.GetRows(res)
}

// Begins a new transaction. No any other thread can send command on this
// connection until Commit or Rollback will be called.
// Periodical pinging the server is disabled during transaction.

func (c *Conn) Begin() (mysql.Transaction, error) {
	//log.Println("Begin")
	c.lock()
	tr := Transaction{
		&Conn{Conn: c.Conn, mutex: new(sync.Mutex)},
		c,
	}
	_, err := c.Conn.Start("START TRANSACTION")
	if err != nil {
		c.unlock()
		return nil, err
	}
	return &tr, nil
}

func (tr *Transaction) end(cr string) error {
	tr.lock()
	_, err := tr.conn.Conn.Start(cr)
	tr.conn.unlock()
	// Invalidate this transaction
	m := tr.Conn.mutex
	tr.Conn = nil
	tr.conn = nil
	m.Unlock() // One goorutine which still uses this transaction will panic
	return err
}

func (tr *Transaction) Commit() error {
	//log.Println("Commit")
	return tr.end("COMMIT")
}

func (tr *Transaction) Rollback() error {
	//log.Println("Rollback")
	return tr.end("ROLLBACK")
}

func (tr *Transaction) Do(st mysql.Stmt) mysql.Stmt {
	if s, ok := st.(*Stmt); ok && s.conn == tr.conn {
		// Returns new statement which uses statement mutexes
		return &Stmt{s.Stmt, tr.Conn}
	}
	panic("Transaction and statement doesn't belong to the same connection")
}

var orgNew func(proto, laddr, raddr, user, passwd string, db ...string) mysql.Conn

func init() {
	orgNew = mysql.New
	mysql.New = New
}
<|endoftext|>


<commit_before>package timelearn

import (
	"context"
	"math"
	"math/rand"
	"time"
)

// A problem is a single problem retrieved.
type Problem struct {
	id       int64
	Question string        // The question to be asked.
	Answer   string        // The correct answer.
	Next     time.Time     // When to next ask this question.
	Interval time.Duration // Current interval to next question
}

// GetNexts queries for `count` upcoming problems that are ready to be
// asked.  Will return an array of problems, with element 0 being the
// next problem that should be asked.
func (t *T) GetNexts(count int) ([]*Problem, error) {
	rows, err := t.conn.QueryContext(context.Background(), `
		SELECT id, question, answer, next, interval
		FROM probs JOIN learning
		WHERE probs.id = learning.probid
			AND next <= ?
		ORDER BY next
		LIMIT ?`, timeToDb(t.now()), count)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var result []*Problem

	for rows.Next() {
		var next float64
		var interval float64
		var p Problem
		err = rows.Scan(&p.id, &p.Question, &p.Answer, &next, &interval)
		if err != nil {
			return nil, err
		}

		p.Next = dbToTime(next)
		p.Interval = dbToDur(interval)

		result = append(result, &p)
	}

	// If we got no rows back, fetch an unlearned problem.  It
	// doesn't make any sense to return more than one, because
	// they will usually be incorrect (time will pass causing
	// other problems to become ready).
	if len(result) == 0 {
		prob, err := t.GetNew()
		if err != nil {
			return nil, err
		}

		if prob != nil {
			result = append(result, prob)
		}
	}

	return result, nil
}

// Get a problem that hasn't started being learned.  The interval and
// next will be set appropriately for a new problem.
// TODO: Set the interval based on a configurable value, as the
// default depends on the problem types.
func (t *T) GetNew() (*Problem, error) {
	var p Problem
	err := t.conn.QueryRow(`
		SELECT id, question, answer
		FROM probs
		WHERE ID NOT IN (SELECT probid FROM learning)
		ORDER BY id
		LIMIT 1`).Scan(&p.id, &p.Question, &p.Answer)
	if err != nil {
		return nil, err
	}

	p.Next = t.now()
	p.Interval = 5 * time.Second

	return &p, nil
}

// Update updates a problem, based on a learning factor.  The scale is
// 1-4, with 1 being totally incorrect, and 4 being totally correct.
func (t *T) Update(prob *Problem, factor int) error {
	var adj float64
	switch factor {
	case 1:
		adj = 0.25
	case 2:
		adj = 0.9
	case 3:
		adj = 1.2
	case 4:
		adj = 2.2
	default:
		panic("Invalid factor, should be 1-4")
	}

	// TODO: Cap this based on the minimum interval, not this
	// arbitrary 5 second value.
	tweak := rand.Float64()*(1.25-0.75) + 0.75
	new_interval := time.Duration(float64(prob.Interval) * adj * tweak)
	if new_interval < 5*time.Second {
		new_interval = 5 * time.Second
	}

	prob.Interval = new_interval
	prob.Next = t.now().Add(prob.Interval)

	tx, err := t.conn.Begin()
	if err != nil {
		return err
	}

	_, err = tx.Exec(`INSERT OR REPLACE INTO learning VALUES (?, ?, ?)`,
		prob.id, timeToDb(prob.Next), durToDb(prob.Interval))
	if err != nil {
		return err
	}

	_, err = tx.Exec(`INSERT INTO log VALUES (?, ?, ?)`,
		timeToDb(t.now()), prob.id, factor)
	if err != nil {
		return err
	}
	err = tx.Commit()
	if err != nil {
		return err
	}

	return nil
}

// Convert time to a float64 so that it can be round-tripped to the
// database.
func timeToDb(t time.Time) float64 {
	return float64(t.Unix()) + float64(t.UnixNano())/1.0e9
}

// Convert time from float64 for round-tripping
func dbToTime(t float64) time.Time {
	if t <= 0.0 {
		panic("Don't expect negative time")
	}

	sec := math.Trunc(t)
	nsec := (t - sec) * 1.0e9

	return time.Unix(int64(sec), int64(nsec))
}

func durToDb(d time.Duration) float64 {
	return float64(d) / 1.0e9
}

func dbToDur(d float64) time.Duration {
	return time.Duration(d * 1.0e9)
}
<commit_msg>go: Flag new problems with " NEW"<commit_after>package timelearn

import (
	"context"
	"math"
	"math/rand"
	"time"
)

// A problem is a single problem retrieved.
type Problem struct {
	id       int64
	Question string        // The question to be asked.
	Answer   string        // The correct answer.
	Next     time.Time     // When to next ask this question.
	Interval time.Duration // Current interval to next question
}

// GetNexts queries for `count` upcoming problems that are ready to be
// asked.  Will return an array of problems, with element 0 being the
// next problem that should be asked.
func (t *T) GetNexts(count int) ([]*Problem, error) {
	rows, err := t.conn.QueryContext(context.Background(), `
		SELECT id, question, answer, next, interval
		FROM probs JOIN learning
		WHERE probs.id = learning.probid
			AND next <= ?
		ORDER BY next
		LIMIT ?`, timeToDb(t.now()), count)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var result []*Problem

	for rows.Next() {
		var next float64
		var interval float64
		var p Problem
		err = rows.Scan(&p.id, &p.Question, &p.Answer, &next, &interval)
		if err != nil {
			return nil, err
		}

		p.Next = dbToTime(next)
		p.Interval = dbToDur(interval)

		result = append(result, &p)
	}

	// If we got no rows back, fetch an unlearned problem.  It
	// doesn't make any sense to return more than one, because
	// they will usually be incorrect (time will pass causing
	// other problems to become ready).
	if len(result) == 0 {
		prob, err := t.GetNew()
		if err != nil {
			return nil, err
		}

		if prob != nil {
			// Add a marker to the problem to make it
			// evident it is new.
			prob.Question = prob.Question + " NEW"
			result = append(result, prob)
		}
	}

	return result, nil
}

// Get a problem that hasn't started being learned.  The interval and
// next will be set appropriately for a new problem.
// TODO: Set the interval based on a configurable value, as the
// default depends on the problem types.
func (t *T) GetNew() (*Problem, error) {
	var p Problem
	err := t.conn.QueryRow(`
		SELECT id, question, answer
		FROM probs
		WHERE ID NOT IN (SELECT probid FROM learning)
		ORDER BY id
		LIMIT 1`).Scan(&p.id, &p.Question, &p.Answer)
	if err != nil {
		return nil, err
	}

	p.Next = t.now()
	p.Interval = 5 * time.Second

	return &p, nil
}

// Update updates a problem, based on a learning factor.  The scale is
// 1-4, with 1 being totally incorrect, and 4 being totally correct.
func (t *T) Update(prob *Problem, factor int) error {
	var adj float64
	switch factor {
	case 1:
		adj = 0.25
	case 2:
		adj = 0.9
	case 3:
		adj = 1.2
	case 4:
		adj = 2.2
	default:
		panic("Invalid factor, should be 1-4")
	}

	// TODO: Cap this based on the minimum interval, not this
	// arbitrary 5 second value.
	tweak := rand.Float64()*(1.25-0.75) + 0.75
	new_interval := time.Duration(float64(prob.Interval) * adj * tweak)
	if new_interval < 5*time.Second {
		new_interval = 5 * time.Second
	}

	prob.Interval = new_interval
	prob.Next = t.now().Add(prob.Interval)

	tx, err := t.conn.Begin()
	if err != nil {
		return err
	}

	_, err = tx.Exec(`INSERT OR REPLACE INTO learning VALUES (?, ?, ?)`,
		prob.id, timeToDb(prob.Next), durToDb(prob.Interval))
	if err != nil {
		return err
	}

	_, err = tx.Exec(`INSERT INTO log VALUES (?, ?, ?)`,
		timeToDb(t.now()), prob.id, factor)
	if err != nil {
		return err
	}
	err = tx.Commit()
	if err != nil {
		return err
	}

	return nil
}

// Convert time to a float64 so that it can be round-tripped to the
// database.
func timeToDb(t time.Time) float64 {
	return float64(t.Unix()) + float64(t.UnixNano())/1.0e9
}

// Convert time from float64 for round-tripping
func dbToTime(t float64) time.Time {
	if t <= 0.0 {
		panic("Don't expect negative time")
	}

	sec := math.Trunc(t)
	nsec := (t - sec) * 1.0e9

	return time.Unix(int64(sec), int64(nsec))
}

func durToDb(d time.Duration) float64 {
	return float64(d) / 1.0e9
}

func dbToDur(d float64) time.Duration {
	return time.Duration(d * 1.0e9)
}
<|endoftext|>