<commit_before>package driverhttp_test

import (
	"net/http"
	"time"

	"bytes"
	"fmt"

	os_http "net/http"

	"io/ioutil"

	"encoding/json"

	"os"
	"os/exec"
	"path"

	"github.com/cloudfoundry-incubator/volman/voldriver"
	"github.com/cloudfoundry-incubator/volman/voldriver/driverhttp"
	"github.com/cloudfoundry/gunk/http_wrap/httpfakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-golang/clock/fakeclock"
	"github.com/pivotal-golang/lager/lagertest"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

var _ = Describe("RemoteClient", func() {

	var (
		testLogger                = lagertest.NewTestLogger("FakeDriver Server Test")
		httpClient                *httpfakes.FakeClient
		driver                    voldriver.Driver
		validHttpMountResponse    *http.Response
		validHttpPathResponse     *http.Response
		validHttpCreateResponse   *http.Response
		validHttpActivateResponse *http.Response
		invalidHttpResponse       *http.Response
		fakeClock                 *fakeclock.FakeClock
	)

	BeforeEach(func() {
		httpClient = new(httpfakes.FakeClient)
		fakeClock = fakeclock.NewFakeClock(time.Now())
		driver = driverhttp.NewRemoteClientWithClient("http://127.0.0.1:8080", httpClient, fakeClock)
		validHttpCreateResponse = &http.Response{
			StatusCode: 200,
		}

		validHttpMountResponse = &http.Response{
			StatusCode: 200,
			Body:       stringCloser{bytes.NewBufferString("{\"Mountpoint\":\"somePath\"}")},
		}

		validHttpPathResponse = &http.Response{
			StatusCode: 200,
			Body:       stringCloser{bytes.NewBufferString("{\"Mountpoint\":\"somePath\"}")},
		}

		validHttpActivateResponse = &http.Response{
			StatusCode: 200,
			Body:       stringCloser{bytes.NewBufferString("{\"Implements\":[\"VolumeDriver\"]}")},
		}
	})

	Context("when the driver returns as error and the transport is TCP", func() {

		BeforeEach(func() {
			fakeClock = fakeclock.NewFakeClock(time.Now())
			httpClient = new(httpfakes.FakeClient)
			driver = driverhttp.NewRemoteClientWithClient("http://127.0.0.1:8080", httpClient, fakeClock)
			invalidHttpResponse = &http.Response{
				StatusCode: 500,
				Body:       stringCloser{bytes.NewBufferString("{\"Err\":\"some error string\"}")},
			}
		})

		It("should not be able to mount", func() {

			httpClient.DoReturns(invalidHttpResponse, nil)

			volumeId := "fake-volume"
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("signaling an error")
			Expect(mountResponse.Err).To(Equal("some error string"))
			Expect(mountResponse.Mountpoint).To(Equal(""))
		})

		It("should not be able to unmount", func() {

			httpClient.DoReturns(invalidHttpResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			By("signaling an error")
			Expect(unmountResponse.Err).To(Equal("some error string"))
		})

	})

	Context("when the driver returns successful and the transport is TCP", func() {
		var volumeId string

		BeforeEach(func() {
			httpClient.DoReturns(validHttpCreateResponse, nil)
			volumeId = "fake-volume"
			createResponse := driver.Create(testLogger, voldriver.CreateRequest{Name: volumeId, Opts: map[string]interface{}{"volume_id": volumeId}})
			Expect(createResponse.Err).To(Equal(""))
		})

		It("should be able to mount", func() {
			httpClient.DoReturns(validHttpMountResponse, nil)

			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("giving back a path with no error")
			Expect(mountResponse.Err).To(Equal(""))
			Expect(mountResponse.Mountpoint).To(Equal("somePath"))
		})

		It("should return mount point", func() {
			httpClient.DoReturns(validHttpPathResponse, nil)

			volumeId := "fake-volume"
			pathResponse := driver.Path(testLogger, voldriver.PathRequest{Name: volumeId})

			Expect(pathResponse.Err).To(Equal(""))
			Expect(pathResponse.Mountpoint).To(Equal("somePath"))
		})

		It("should be able to unmount", func() {

			validHttpUnmountResponse := &http.Response{
				StatusCode: 200,
			}

			httpClient.DoReturns(validHttpUnmountResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			Expect(unmountResponse.Err).To(Equal(""))
		})

		It("should be able to activate", func() {
			httpClient.DoReturns(validHttpActivateResponse, nil)

			activateResponse := driver.Activate(testLogger)

			By("giving back a path with no error")
			Expect(activateResponse.Err).To(Equal(""))
			Expect(activateResponse.Implements).To(Equal([]string{"VolumeDriver"}))
		})
	})

	Context("when the driver is malicious and the transport is TCP", func() {

		BeforeEach(func() {
		})

		It("should not be able to mount", func() {

			validHttpMountResponse = &http.Response{
				StatusCode: 200,
				Body:       stringCloser{bytes.NewBufferString("i am trying to pown your system")},
			}

			httpClient.DoReturns(validHttpMountResponse, nil)

			volumeId := "fake-volume"
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("signaling an error")
			Expect(mountResponse.Err).NotTo(Equal(""))
			Expect(mountResponse.Mountpoint).To(Equal(""))
		})

		It("should still not be able to mount", func() {

			invalidHttpResponse := &http.Response{
				StatusCode: 500,
				Body:       stringCloser{bytes.NewBufferString("i am trying to pown your system")},
			}

			httpClient.DoReturns(invalidHttpResponse, nil)

			volumeId := "fake-volume"
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("signaling an error")
			Expect(mountResponse.Err).NotTo(Equal(""))
			Expect(mountResponse.Mountpoint).To(Equal(""))
		})

		It("should not be able to unmount", func() {

			validHttpUnmountResponse := &http.Response{
				StatusCode: 500,
				Body:       stringCloser{bytes.NewBufferString("i am trying to pown your system")},
			}

			httpClient.DoReturns(validHttpUnmountResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			Expect(unmountResponse.Err).NotTo(Equal(""))
		})

	})

	Context("when the http transport fails and the transport is TCP", func() {

		BeforeEach(func() {
			// all of the tests in this context will perform retry logic over 30 seconds, so we need to
			// simulate time passing.
			go fastForward(fakeClock, 40)
		})

		It("should fail to mount", func() {

			httpClient.DoReturns(nil, fmt.Errorf("connection failed"))

			volumeId := "fake-volume"
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("signaling an error")
			Expect(mountResponse.Err).To(Equal("connection failed"))
		})

		It("should fail to unmount", func() {

			httpClient.DoReturns(nil, fmt.Errorf("connection failed"))

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			By("signaling an error")
			Expect(unmountResponse.Err).NotTo(Equal(""))
		})

		It("should still fail to unmount", func() {

			invalidHttpResponse := &http.Response{
				StatusCode: 500,
				Body:       errCloser{bytes.NewBufferString("")},
			}

			httpClient.DoReturns(invalidHttpResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			Expect(unmountResponse.Err).NotTo(Equal(""))
		})

		It("should fail to activate", func() {
			httpClient.DoReturns(nil, fmt.Errorf("connection failed"))

			activateResponse := driver.Activate(testLogger)

			By("signaling an error")
			Expect(activateResponse.Err).NotTo(Equal(""))
		})

	})

	Context("when the transport is unix", func() {
		var (
			volumeId                    string
			unixRunner                  *ginkgomon.Runner
			fakedriverUnixServerProcess ifrit.Process
			socketPath                  string
		)

		BeforeEach(func() {
			tmpdir, err := ioutil.TempDir(os.TempDir(), "fake-driver-test")
			Expect(err).ShouldNot(HaveOccurred())

			socketPath = path.Join(tmpdir, "fakedriver.sock")

			unixRunner = ginkgomon.New(ginkgomon.Config{
				Name: "fakedriverUnixServer",
				Command: exec.Command(
					fakeDriverPath,
					"-listenAddr", socketPath,
					"-transport", "unix",
				),
				StartCheck: "fakedriverServer.started",
			})

			httpClient = new(httpfakes.FakeClient)
			volumeId = "fake-volume"
			fakedriverUnixServerProcess = ginkgomon.Invoke(unixRunner)

			time.Sleep(time.Millisecond * 1000)

			fakeClock = fakeclock.NewFakeClock(time.Now())
			driver = driverhttp.NewRemoteClientWithClient(socketPath, httpClient, fakeClock)
			validHttpMountResponse = &http.Response{
				StatusCode: 200,
				Body:       stringCloser{bytes.NewBufferString("{\"Mountpoint\":\"somePath\"}")},
			}
		})

		AfterEach(func() {
			ginkgomon.Kill(fakedriverUnixServerProcess)
		})

		It("should be able to mount", func() {
			httpClient.DoReturns(validHttpMountResponse, nil)
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("returning a mountpoint without errors")
			Expect(mountResponse.Err).To(Equal(""))
			Expect(mountResponse.Mountpoint).To(Equal("somePath"))
		})

		It("should be able to unmount", func() {

			validHttpUnmountResponse := &http.Response{
				StatusCode: 200,
			}

			httpClient.DoReturns(validHttpUnmountResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			Expect(unmountResponse.Err).To(Equal(""))
		})

		It("should be able to activate", func() {
			httpClient.DoReturns(validHttpActivateResponse, nil)

			activateResponse := driver.Activate(testLogger)

			By("giving back a activation response with no error")
			Expect(activateResponse.Err).To(Equal(""))
			Expect(activateResponse.Implements).To(Equal([]string{"VolumeDriver"}))
		})

	})

	Context("when the transport fails and back off is required", func() {

		var (
			retryCount    int
			mountResponse voldriver.MountResponse
			volumeId      string
		)

		Context("when it fails first time and then succeeds", func() {

			var (
				requestBody []byte
				err         error
			)

			BeforeEach(func() {
				httpClient.DoStub = func(req *os_http.Request) (resp *os_http.Response, err error) {

					defer func() {
						retryCount = retryCount + 1
					}()

					requestBody, err = ioutil.ReadAll(req.Body)
					Expect(err).NotTo(HaveOccurred())

					if retryCount == 0 {
						return nil, fmt.Errorf("connection failed but read body")
					}

					return validHttpMountResponse, nil
				}

				go fastForward(fakeClock, 10)

				volumeId = "fake-volume"
				mountResponse = driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})
			})

			It("should have the correct number of retries and the correct response", func() {
				Expect(retryCount).To(Equal(2))
			})

			It("should submit the correct request each time", func() {
				var expectedRequestBody []byte
				expectedRequestBody, err = json.Marshal(voldriver.MountRequest{Name: volumeId})
				Expect(err).NotTo(HaveOccurred())

				Expect(requestBody).To(Equal(expectedRequestBody))
			})

			It("should have the correct response", func() {
				Expect(mountResponse.Err).To(Equal(""))
				Expect(mountResponse.Mountpoint).NotTo(Equal(""))
			})

		})

		Context("when it fails and timeout exceeds", func() {

			var timestamp time.Time

			BeforeEach(func() {
				httpClient.DoStub = func(req *os_http.Request) (resp *os_http.Response, err error) {
					return nil, fmt.Errorf("connection failed")
				}

				go fastForward(fakeClock, 40)

				timestamp = fakeClock.Now()
				volumeId = "fake-volume"
				mountResponse = driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})
			})

			It("should return an error after 30 seconds have passed", func() {
				Expect(mountResponse.Err).NotTo(Equal(""))

				elapsed := fakeClock.Now().Sub(timestamp)
				Expect(elapsed.Seconds()).To(BeNumerically(">", 30))
			})
		})
	})
})

func fastForward(fakeClock *fakeclock.FakeClock, seconds int) {
	for i := 0; i < seconds; i++ {
		time.Sleep(time.Millisecond * 1)
		fakeClock.IncrementBySeconds(1)
	}
}
<commit_msg>Fix flaky backoff unit test<commit_after>package driverhttp_test

import (
	"net/http"
	"time"

	"bytes"
	"fmt"

	os_http "net/http"

	"io/ioutil"

	"encoding/json"

	"os"
	"os/exec"
	"path"

	"github.com/cloudfoundry-incubator/volman/voldriver"
	"github.com/cloudfoundry-incubator/volman/voldriver/driverhttp"
	"github.com/cloudfoundry/gunk/http_wrap/httpfakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-golang/clock/fakeclock"
	"github.com/pivotal-golang/lager/lagertest"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

var _ = Describe("RemoteClient", func() {

	var (
		testLogger                = lagertest.NewTestLogger("FakeDriver Server Test")
		httpClient                *httpfakes.FakeClient
		driver                    voldriver.Driver
		validHttpMountResponse    *http.Response
		validHttpPathResponse     *http.Response
		validHttpCreateResponse   *http.Response
		validHttpActivateResponse *http.Response
		invalidHttpResponse       *http.Response
		fakeClock                 *fakeclock.FakeClock
	)

	BeforeEach(func() {
		httpClient = new(httpfakes.FakeClient)
		fakeClock = fakeclock.NewFakeClock(time.Now())
		driver = driverhttp.NewRemoteClientWithClient("http://127.0.0.1:8080", httpClient, fakeClock)
		validHttpCreateResponse = &http.Response{
			StatusCode: 200,
		}

		validHttpMountResponse = &http.Response{
			StatusCode: 200,
			Body:       stringCloser{bytes.NewBufferString("{\"Mountpoint\":\"somePath\"}")},
		}

		validHttpPathResponse = &http.Response{
			StatusCode: 200,
			Body:       stringCloser{bytes.NewBufferString("{\"Mountpoint\":\"somePath\"}")},
		}

		validHttpActivateResponse = &http.Response{
			StatusCode: 200,
			Body:       stringCloser{bytes.NewBufferString("{\"Implements\":[\"VolumeDriver\"]}")},
		}
	})

	Context("when the driver returns as error and the transport is TCP", func() {

		BeforeEach(func() {
			fakeClock = fakeclock.NewFakeClock(time.Now())
			httpClient = new(httpfakes.FakeClient)
			driver = driverhttp.NewRemoteClientWithClient("http://127.0.0.1:8080", httpClient, fakeClock)
			invalidHttpResponse = &http.Response{
				StatusCode: 500,
				Body:       stringCloser{bytes.NewBufferString("{\"Err\":\"some error string\"}")},
			}
		})

		It("should not be able to mount", func() {

			httpClient.DoReturns(invalidHttpResponse, nil)

			volumeId := "fake-volume"
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("signaling an error")
			Expect(mountResponse.Err).To(Equal("some error string"))
			Expect(mountResponse.Mountpoint).To(Equal(""))
		})

		It("should not be able to unmount", func() {

			httpClient.DoReturns(invalidHttpResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			By("signaling an error")
			Expect(unmountResponse.Err).To(Equal("some error string"))
		})

	})

	Context("when the driver returns successful and the transport is TCP", func() {
		var volumeId string

		BeforeEach(func() {
			httpClient.DoReturns(validHttpCreateResponse, nil)
			volumeId = "fake-volume"
			createResponse := driver.Create(testLogger, voldriver.CreateRequest{Name: volumeId, Opts: map[string]interface{}{"volume_id": volumeId}})
			Expect(createResponse.Err).To(Equal(""))
		})

		It("should be able to mount", func() {
			httpClient.DoReturns(validHttpMountResponse, nil)

			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("giving back a path with no error")
			Expect(mountResponse.Err).To(Equal(""))
			Expect(mountResponse.Mountpoint).To(Equal("somePath"))
		})

		It("should return mount point", func() {
			httpClient.DoReturns(validHttpPathResponse, nil)

			volumeId := "fake-volume"
			pathResponse := driver.Path(testLogger, voldriver.PathRequest{Name: volumeId})

			Expect(pathResponse.Err).To(Equal(""))
			Expect(pathResponse.Mountpoint).To(Equal("somePath"))
		})

		It("should be able to unmount", func() {

			validHttpUnmountResponse := &http.Response{
				StatusCode: 200,
			}

			httpClient.DoReturns(validHttpUnmountResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			Expect(unmountResponse.Err).To(Equal(""))
		})

		It("should be able to activate", func() {
			httpClient.DoReturns(validHttpActivateResponse, nil)

			activateResponse := driver.Activate(testLogger)

			By("giving back a path with no error")
			Expect(activateResponse.Err).To(Equal(""))
			Expect(activateResponse.Implements).To(Equal([]string{"VolumeDriver"}))
		})
	})

	Context("when the driver is malicious and the transport is TCP", func() {

		BeforeEach(func() {
		})

		It("should not be able to mount", func() {

			validHttpMountResponse = &http.Response{
				StatusCode: 200,
				Body:       stringCloser{bytes.NewBufferString("i am trying to pown your system")},
			}

			httpClient.DoReturns(validHttpMountResponse, nil)

			volumeId := "fake-volume"
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("signaling an error")
			Expect(mountResponse.Err).NotTo(Equal(""))
			Expect(mountResponse.Mountpoint).To(Equal(""))
		})

		It("should still not be able to mount", func() {

			invalidHttpResponse := &http.Response{
				StatusCode: 500,
				Body:       stringCloser{bytes.NewBufferString("i am trying to pown your system")},
			}

			httpClient.DoReturns(invalidHttpResponse, nil)

			volumeId := "fake-volume"
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("signaling an error")
			Expect(mountResponse.Err).NotTo(Equal(""))
			Expect(mountResponse.Mountpoint).To(Equal(""))
		})

		It("should not be able to unmount", func() {

			validHttpUnmountResponse := &http.Response{
				StatusCode: 500,
				Body:       stringCloser{bytes.NewBufferString("i am trying to pown your system")},
			}

			httpClient.DoReturns(validHttpUnmountResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			Expect(unmountResponse.Err).NotTo(Equal(""))
		})

	})

	Context("when the http transport fails and the transport is TCP", func() {

		BeforeEach(func() {
			// all of the tests in this context will perform retry logic over 30 seconds, so we need to
			// simulate time passing.
			go fastForward(fakeClock, 40)
		})

		It("should fail to mount", func() {

			httpClient.DoReturns(nil, fmt.Errorf("connection failed"))

			volumeId := "fake-volume"
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("signaling an error")
			Expect(mountResponse.Err).To(Equal("connection failed"))
		})

		It("should fail to unmount", func() {

			httpClient.DoReturns(nil, fmt.Errorf("connection failed"))

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			By("signaling an error")
			Expect(unmountResponse.Err).NotTo(Equal(""))
		})

		It("should still fail to unmount", func() {

			invalidHttpResponse := &http.Response{
				StatusCode: 500,
				Body:       errCloser{bytes.NewBufferString("")},
			}

			httpClient.DoReturns(invalidHttpResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			Expect(unmountResponse.Err).NotTo(Equal(""))
		})

		It("should fail to activate", func() {
			httpClient.DoReturns(nil, fmt.Errorf("connection failed"))

			activateResponse := driver.Activate(testLogger)

			By("signaling an error")
			Expect(activateResponse.Err).NotTo(Equal(""))
		})

	})

	Context("when the transport is unix", func() {
		var (
			volumeId                    string
			unixRunner                  *ginkgomon.Runner
			fakedriverUnixServerProcess ifrit.Process
			socketPath                  string
		)

		BeforeEach(func() {
			tmpdir, err := ioutil.TempDir(os.TempDir(), "fake-driver-test")
			Expect(err).ShouldNot(HaveOccurred())

			socketPath = path.Join(tmpdir, "fakedriver.sock")

			unixRunner = ginkgomon.New(ginkgomon.Config{
				Name: "fakedriverUnixServer",
				Command: exec.Command(
					fakeDriverPath,
					"-listenAddr", socketPath,
					"-transport", "unix",
				),
				StartCheck: "fakedriverServer.started",
			})

			httpClient = new(httpfakes.FakeClient)
			volumeId = "fake-volume"
			fakedriverUnixServerProcess = ginkgomon.Invoke(unixRunner)

			time.Sleep(time.Millisecond * 1000)

			fakeClock = fakeclock.NewFakeClock(time.Now())
			driver = driverhttp.NewRemoteClientWithClient(socketPath, httpClient, fakeClock)
			validHttpMountResponse = &http.Response{
				StatusCode: 200,
				Body:       stringCloser{bytes.NewBufferString("{\"Mountpoint\":\"somePath\"}")},
			}
		})

		AfterEach(func() {
			ginkgomon.Kill(fakedriverUnixServerProcess)
		})

		It("should be able to mount", func() {
			httpClient.DoReturns(validHttpMountResponse, nil)
			mountResponse := driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})

			By("returning a mountpoint without errors")
			Expect(mountResponse.Err).To(Equal(""))
			Expect(mountResponse.Mountpoint).To(Equal("somePath"))
		})

		It("should be able to unmount", func() {

			validHttpUnmountResponse := &http.Response{
				StatusCode: 200,
			}

			httpClient.DoReturns(validHttpUnmountResponse, nil)

			volumeId := "fake-volume"
			unmountResponse := driver.Unmount(testLogger, voldriver.UnmountRequest{Name: volumeId})

			Expect(unmountResponse.Err).To(Equal(""))
		})

		It("should be able to activate", func() {
			httpClient.DoReturns(validHttpActivateResponse, nil)

			activateResponse := driver.Activate(testLogger)

			By("giving back a activation response with no error")
			Expect(activateResponse.Err).To(Equal(""))
			Expect(activateResponse.Implements).To(Equal([]string{"VolumeDriver"}))
		})

	})

	Context("when the transport fails and back off is required", func() {

		var (
			retryCount    int
			mountResponse voldriver.MountResponse
			volumeId      string
		)

		Context("when it fails first time and then succeeds", func() {

			var (
				requestBody []byte
				err         error
			)

			BeforeEach(func() {
				httpClient.DoStub = func(req *os_http.Request) (resp *os_http.Response, err error) {

					defer func() {
						retryCount = retryCount + 1
					}()

					requestBody, err = ioutil.ReadAll(req.Body)
					Expect(err).NotTo(HaveOccurred())

					if retryCount == 0 {
						return nil, fmt.Errorf("connection failed but read body")
					}

					return validHttpMountResponse, nil
				}

				go fastForward(fakeClock, 10)

				volumeId = "fake-volume"
				mountResponse = driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})
			})

			It("should have the correct number of retries, submit the same request each time and eventually recieve the correct response", func() {
				// retries
				Expect(retryCount).To(Equal(2))

				// request
				var expectedRequestBody []byte
				expectedRequestBody, err = json.Marshal(voldriver.MountRequest{Name: volumeId})
				Expect(err).NotTo(HaveOccurred())
				Expect(requestBody).To(Equal(expectedRequestBody))

				// response
				Expect(mountResponse.Err).To(Equal(""))
				Expect(mountResponse.Mountpoint).NotTo(Equal(""))
			})
		})

		Context("when it fails and timeout exceeds", func() {

			var timestamp time.Time

			BeforeEach(func() {
				httpClient.DoStub = func(req *os_http.Request) (resp *os_http.Response, err error) {
					return nil, fmt.Errorf("connection failed")
				}

				go fastForward(fakeClock, 40)

				timestamp = fakeClock.Now()
				volumeId = "fake-volume"
				mountResponse = driver.Mount(testLogger, voldriver.MountRequest{Name: volumeId})
			})

			It("should return an error after 30 seconds have passed", func() {
				Expect(mountResponse.Err).NotTo(Equal(""))

				elapsed := fakeClock.Now().Sub(timestamp)
				Expect(elapsed.Seconds()).To(BeNumerically(">", 30))
			})
		})
	})
})

func fastForward(fakeClock *fakeclock.FakeClock, seconds int) {
	for i := 0; i < seconds; i++ {
		time.Sleep(time.Millisecond * 1)
		fakeClock.IncrementBySeconds(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"image"
	"image/png"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/shurcooL/go/u/u4"
	"github.com/shurcooL/trayhost"
	_ "golang.org/x/image/tiff"
)

// TODO: Load from config. Have ability to set config.
var hostFlag = flag.String("host", "", "Target server host.")
var debugFlag = flag.Bool("debug", false, "Adds menu items for debugging purposes.")

var httpClient = &http.Client{Timeout: 3 * time.Second}

var clipboard struct {
	extension string // File extension in lower case: "png", "tiff", "mov", etc. Empty string means no content.
	bytes     []byte
}
var notificationThumbnail trayhost.Image

// fileThumbnail returns a thumbnail image that represents the input file, or an empty image if it cannot.
func fileThumbnail(extension string, bytes []byte) trayhost.Image {
	switch extension {
	case "png":
		return trayhost.Image{Kind: trayhost.ImageKind(extension), Bytes: bytes}
	default:
		return trayhost.Image{}
	}
}

// instantShareEnabled returns true if we have valid content in clipboard that can be instantly shared.
// It also updates values of clipboard and notificationThumbnail if so.
func instantShareEnabled() bool {
	fmt.Println("grab content, content-type of clipboard")

	cc, err := trayhost.GetClipboardContent()
	if err != nil {
		return false
	}

	switch {
	case len(cc.Files) == 1: // Single file.
		b, err := ioutil.ReadFile(cc.Files[0])
		if err != nil {
			return false
		}
		extension := strings.TrimPrefix(filepath.Ext(cc.Files[0]), ".")
		clipboard.extension = extension
		clipboard.bytes = b
		notificationThumbnail = fileThumbnail(extension, b)
		return true
	case cc.Image.Kind != "":
		clipboard.extension = string(cc.Image.Kind)
		clipboard.bytes = cc.Image.Bytes
		notificationThumbnail = cc.Image
		return true
	default:
		return false
	}
}

func instantShareHandler() {
	// Convert some source image types to desired destination format.
	switch clipboard.extension {
	case "tiff":
		// Convert tiff to png.
		m, _, err := image.Decode(bytes.NewReader(clipboard.bytes))
		if err != nil {
			log.Panicln("image.Decode:", err)
		}

		var buf bytes.Buffer
		err = png.Encode(&buf, m)
		if err != nil {
			log.Panicln("png.Encode:", err)
		}

		clipboard.extension = "png"
		clipboard.bytes = buf.Bytes()
	}

	fmt.Println("request URL")

	resp, err := httpClient.Get(*hostFlag + "/api/getfilename?ext=" + clipboard.extension)
	if err != nil {
		trayhost.Notification{Title: "Upload Failed", Body: err.Error()}.Display()
		log.Println(err)
		return
	}
	defer resp.Body.Close()
	filename, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		trayhost.Notification{Title: "Upload Failed", Body: err.Error()}.Display()
		log.Println(err)
		return
	}

	fmt.Println("display/put URL in clipboard")

	url := *hostFlag + "/" + string(filename)
	trayhost.SetClipboardText(url)
	trayhost.Notification{
		Title:   "Success",
		Body:    url,
		Image:   notificationThumbnail,
		Timeout: 3 * time.Second,
		Handler: func() {
			// On click, open the displayed URL.
			u4.Open(url)
		},
	}.Display()

	fmt.Println("upload image in background of size", len(clipboard.bytes))

	go func(b []byte) {
		req, err := http.NewRequest("PUT", url, bytes.NewReader(b))
		if err != nil {
			log.Println(err)
			return
		}
		req.Header.Set("Content-Type", "application/octet-stream")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			log.Println(err)
			return
		}
		_ = resp.Body.Close()
		fmt.Println("done")
	}(clipboard.bytes)
}

func main() {
	flag.Parse()
	runtime.LockOSThread()

	menuItems := []trayhost.MenuItem{
		trayhost.MenuItem{
			Title:   "Instant Share",
			Enabled: instantShareEnabled,
			Handler: instantShareHandler,
		},
		trayhost.SeparatorMenuItem(),
		trayhost.MenuItem{
			Title:   "Quit",
			Handler: trayhost.Exit,
		},
	}
	if *debugFlag {
		menuItems = append(menuItems,
			trayhost.SeparatorMenuItem(),
			trayhost.MenuItem{
				Title: "Debug: Get Clipboard Content",
				Handler: func() {
					cc, err := trayhost.GetClipboardContent()
					fmt.Printf("GetClipboardContent() error: %v\n", err)
					fmt.Printf("Text: %q\n", cc.Text)
					fmt.Printf("Image: %v len(%v)\n", cc.Image.Kind, len(cc.Image.Bytes))
					fmt.Printf("Files: len(%v) %v\n", len(cc.Files), cc.Files)
				},
			},
			trayhost.MenuItem{
				Title: "Debug: Set Clipboard Text",
				Handler: func() {
					trayhost.SetClipboardText("http://www.example.com/image.png")
				},
			},
			trayhost.MenuItem{
				Title: "Debug: Notification",
				Handler: func() {
					handler := func() {
						u4.Open("http://www.example.com/image.png")
					}
					notification := trayhost.Notification{Title: "Upload Complete", Body: "http://www.example.com/image.png", Timeout: 3 * time.Second, Handler: handler}
					//trayhost.Notification{Title: "Upload Failed", Body: "error description goes here"}.Display()
					if cc, err := trayhost.GetClipboardContent(); err == nil && cc.Image.Kind != "" {
						notification.Image = cc.Image
					}
					notification.Display()
				},
			},
		)
	}

	// TODO: Create a real icon and bake it into the binary.
	// TODO: Optionally, if non-Retina pixel perfection is required, generate or load 1x image and supply that as a second representation.
	iconData, err := ioutil.ReadFile("./icon@2x.png")
	if err != nil {
		panic(err)
	}

	fmt.Println("Starting.")

	trayhost.Initialize("Instant Share", iconData, menuItems)

	trayhost.EnterLoop()

	fmt.Println("Exiting.")
}
<commit_msg>Only convert tiff -> png when its coming from clipboard.<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"image"
	"image/png"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/shurcooL/go/u/u4"
	"github.com/shurcooL/trayhost"
	_ "golang.org/x/image/tiff"
)

// TODO: Load from config. Have ability to set config.
var hostFlag = flag.String("host", "", "Target server host.")
var debugFlag = flag.Bool("debug", false, "Adds menu items for debugging purposes.")

var httpClient = &http.Client{Timeout: 3 * time.Second}

var clipboard struct {
	extension string // File extension in lower case: "png", "tiff", "mov", etc. Empty string means no content.
	bytes     []byte
}
var notificationThumbnail trayhost.Image

// fileThumbnail returns a thumbnail image that represents the input file, or an empty image if it cannot.
func fileThumbnail(extension string, bytes []byte) trayhost.Image {
	switch extension {
	case "png":
		return trayhost.Image{Kind: trayhost.ImageKind(extension), Bytes: bytes}
	default:
		return trayhost.Image{}
	}
}

// instantShareEnabled returns true if we have valid content in clipboard that can be instantly shared.
// It also updates values of clipboard and notificationThumbnail if so.
func instantShareEnabled() bool {
	fmt.Println("grab content, content-type of clipboard")

	cc, err := trayhost.GetClipboardContent()
	if err != nil {
		return false
	}

	switch {
	case len(cc.Files) == 1: // Single file.
		b, err := ioutil.ReadFile(cc.Files[0])
		if err != nil {
			return false
		}
		extension := strings.TrimPrefix(filepath.Ext(cc.Files[0]), ".")
		clipboard.extension = extension
		clipboard.bytes = b
		notificationThumbnail = fileThumbnail(extension, b)
		return true
	case cc.Image.Kind != "":
		clipboard.extension = string(cc.Image.Kind)
		clipboard.bytes = cc.Image.Bytes
		notificationThumbnail = cc.Image

		// Convert some source clipboard image types to desired destination format.
		switch clipboard.extension {
		case "tiff":
			// Convert tiff to png.
			m, _, err := image.Decode(bytes.NewReader(clipboard.bytes))
			if err != nil {
				log.Panicln("image.Decode:", err)
			}

			var buf bytes.Buffer
			err = png.Encode(&buf, m)
			if err != nil {
				log.Panicln("png.Encode:", err)
			}

			clipboard.extension = "png"
			clipboard.bytes = buf.Bytes()
		}

		return true
	default:
		return false
	}
}

func instantShareHandler() {
	fmt.Println("request URL")

	resp, err := httpClient.Get(*hostFlag + "/api/getfilename?ext=" + clipboard.extension)
	if err != nil {
		trayhost.Notification{Title: "Upload Failed", Body: err.Error()}.Display()
		log.Println(err)
		return
	}
	defer resp.Body.Close()
	filename, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		trayhost.Notification{Title: "Upload Failed", Body: err.Error()}.Display()
		log.Println(err)
		return
	}

	fmt.Println("display/put URL in clipboard")

	url := *hostFlag + "/" + string(filename)
	trayhost.SetClipboardText(url)
	trayhost.Notification{
		Title:   "Success",
		Body:    url,
		Image:   notificationThumbnail,
		Timeout: 3 * time.Second,
		Handler: func() {
			// On click, open the displayed URL.
			u4.Open(url)
		},
	}.Display()

	fmt.Println("upload image in background of size", len(clipboard.bytes))

	go func(b []byte) {
		req, err := http.NewRequest("PUT", url, bytes.NewReader(b))
		if err != nil {
			log.Println(err)
			return
		}
		req.Header.Set("Content-Type", "application/octet-stream")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			log.Println(err)
			return
		}
		_ = resp.Body.Close()
		fmt.Println("done")
	}(clipboard.bytes)
}

func main() {
	flag.Parse()
	runtime.LockOSThread()

	menuItems := []trayhost.MenuItem{
		trayhost.MenuItem{
			Title:   "Instant Share",
			Enabled: instantShareEnabled,
			Handler: instantShareHandler,
		},
		trayhost.SeparatorMenuItem(),
		trayhost.MenuItem{
			Title:   "Quit",
			Handler: trayhost.Exit,
		},
	}
	if *debugFlag {
		menuItems = append(menuItems,
			trayhost.SeparatorMenuItem(),
			trayhost.MenuItem{
				Title: "Debug: Get Clipboard Content",
				Handler: func() {
					cc, err := trayhost.GetClipboardContent()
					fmt.Printf("GetClipboardContent() error: %v\n", err)
					fmt.Printf("Text: %q\n", cc.Text)
					fmt.Printf("Image: %v len(%v)\n", cc.Image.Kind, len(cc.Image.Bytes))
					fmt.Printf("Files: len(%v) %v\n", len(cc.Files), cc.Files)
				},
			},
			trayhost.MenuItem{
				Title: "Debug: Set Clipboard Text",
				Handler: func() {
					trayhost.SetClipboardText("http://www.example.com/image.png")
				},
			},
			trayhost.MenuItem{
				Title: "Debug: Notification",
				Handler: func() {
					handler := func() {
						u4.Open("http://www.example.com/image.png")
					}
					notification := trayhost.Notification{Title: "Upload Complete", Body: "http://www.example.com/image.png", Timeout: 3 * time.Second, Handler: handler}
					//trayhost.Notification{Title: "Upload Failed", Body: "error description goes here"}.Display()
					if cc, err := trayhost.GetClipboardContent(); err == nil && cc.Image.Kind != "" {
						notification.Image = cc.Image
					}
					notification.Display()
				},
			},
		)
	}

	// TODO: Create a real icon and bake it into the binary.
	// TODO: Optionally, if non-Retina pixel perfection is required, generate or load 1x image and supply that as a second representation.
	iconData, err := ioutil.ReadFile("./icon@2x.png")
	if err != nil {
		panic(err)
	}

	fmt.Println("Starting.")

	trayhost.Initialize("Instant Share", iconData, menuItems)

	trayhost.EnterLoop()

	fmt.Println("Exiting.")
}
<|endoftext|>


<commit_before>package riot

import (
	"encoding/gob"
	"fmt"
	"log"
	"os"
	"testing"

	"github.com/go-ego/riot/types"
	"github.com/vcaesar/tt"
)

func makeDocIds() map[uint64]bool {
	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[3] = true
	docIds[1] = true
	docIds[2] = true

	return docIds
}

func TestEngineIndexWithNewStore(t *testing.T) {
	gob.Register(ScoringFields{})
	var engine = New("./testdata/test_dict.txt", "./riot.new", 8)
	log.Println("new engine start...")
	// engine = engine.New()
	AddDocs(engine)

	engine.RemoveDoc(5, true)
	engine.Flush()

	engine.Close()
	// os.RemoveAll("riot.new")

	// var engine1 = New("./testdata/test_dict.txt", "./riot.new")
	var engine1 = New("./testdata/test_new.toml")
	// engine1 = engine1.New()
	log.Println("test...")
	engine1.Flush()
	log.Println("new engine1 start...")

	outputs := engine1.Search(types.SearchReq{Text: "World人口"})
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "world", outputs.Tokens[0])
	tt.Expect(t, "人口", outputs.Tokens[1])

	outDocs := outputs.Docs.(types.ScoredDocs)
	tt.Expect(t, "2", len(outDocs))

	// tt.Expect(t, "2", outDocs[0].DocId)
	tt.Expect(t, "2500", int(outDocs[0].Scores[0]*1000))
	tt.Expect(t, "[]", outDocs[0].TokenSnippetLocs)

	// tt.Expect(t, "1", outDocs[1].DocId)
	tt.Expect(t, "2215", int(outDocs[1].Scores[0]*1000))
	tt.Expect(t, "[]", outDocs[1].TokenSnippetLocs)

	engine1.Close()
	os.RemoveAll("riot.new")
	// os.RemoveAll("riot-index")
}

var (
	rankTestOpts = rankOptsMax(0, 1)
)

func testRankOpt(idOnly bool) types.EngineOpts {
	return types.EngineOpts{
		Using:       1,
		IDOnly:      idOnly,
		GseDict:     "./testdata/test_dict.txt",
		DefRankOpts: &rankTestOpts,
		IndexerOpts: &types.IndexerOpts{
			IndexType: types.LocsIndex,
		},
	}
}

func lookupReq(engine *Engine) (types.SearchReq, []string, chan rankerReturnReq) {
	request := types.SearchReq{
		Text:   "World人口",
		DocIds: makeDocIds(),
	}

	tokens := engine.Tokens(request)
	// 建立排序器返回的通信通道
	rankerReturnChan := make(
		chan rankerReturnReq, engine.initOptions.NumShards)

	// 生成查找请求
	lookupRequest := indexerLookupReq{
		countDocsOnly:    request.CountDocsOnly,
		tokens:           tokens,
		labels:           request.Labels,
		docIds:           request.DocIds,
		options:          rankTestOpts,
		rankerReturnChan: rankerReturnChan,
		orderless:        request.Orderless,
		logic:            request.Logic,
	}

	// 向索引器发送查找请求
	for shard := 0; shard < engine.initOptions.NumShards; shard++ {
		engine.indexerLookupChans[shard] <- lookupRequest
	}

	return request, tokens, rankerReturnChan
}

func TestDocRankID(t *testing.T) {
	var engine Engine

	engine.Init(testRankOpt(true))
	AddDocs(&engine)

	engine.RemoveDoc(5)
	engine.Flush()

	request, tokens, rankerReturnChan := lookupReq(&engine)
	outputs := engine.RankID(request, rankTestOpts, tokens, rankerReturnChan)

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		tt.Expect(t, "1", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "2", outputs.NumDocs)

	engine.Close()
}

func TestDocRanks(t *testing.T) {
	var engine Engine

	engine.Init(testRankOpt(false))
	AddDocs(&engine)

	engine.RemoveDoc(5)
	engine.Flush()

	request, tokens, rankerReturnChan := lookupReq(&engine)
	outputs := engine.Ranks(request, rankTestOpts, tokens, rankerReturnChan)

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredDocs)
		tt.Expect(t, "1", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "2", outputs.NumDocs)

	// test search
	outputs1 := engine.Search(types.SearchReq{
		Text:    "World人口",
		Timeout: 10,
		DocIds:  makeDocIds()})

	if outputs1.Docs != nil {
		outDocs1 := outputs.Docs.(types.ScoredDocs)
		tt.Expect(t, "1", len(outDocs1))
	}
	tt.Expect(t, "2", len(outputs1.Tokens))
	tt.Expect(t, "2", outputs1.NumDocs)

	engine.Close()
}

func TestDocGetAllDocAndID(t *testing.T) {
	gob.Register(ScoringFields{})

	var engine Engine
	engine.Init(types.EngineOpts{
		Using:     1,
		NumShards: 5,
		UseStore:  true,
		// StoreEngine: "bg",
		StoreFolder: "riot.id",
		IDOnly:      true,
		GseDict:     "./testdata/test_dict.txt",
		DefRankOpts: &rankTestOpts,
		IndexerOpts: &types.IndexerOpts{
			IndexType: types.LocsIndex,
		},
	})

	AddDocs(&engine)

	engine.RemoveDoc(5)
	engine.Flush()

	allIds := engine.GetDBAllIds()
	fmt.Println("all id", allIds)
	tt.Expect(t, "5", len(allIds))
	tt.Expect(t, "[3 4 1 6 2]", allIds)

	allIds = engine.GetAllDocIds()
	fmt.Println("all doc id", allIds)
	tt.Expect(t, "5", len(allIds))
	tt.Expect(t, "[3 4 1 6 2]", allIds)

	ids, docs := engine.GetDBAllDocs()
	fmt.Println("all id and doc", allIds, docs)
	tt.Expect(t, "5", len(ids))
	tt.Expect(t, "5", len(docs))
	tt.Expect(t, "[3 4 1 6 2]", ids)
	allDoc := `[{The world <nil> [] [] <nil>} {有人口 <nil> [] [] {2 3 1}} {The world, 有七十亿人口人口 <nil> [] [] {1 2 3}} {有七十亿人口 <nil> [] [] {2 3 3}} {The world, 人口 <nil> [] [] <nil>}]`
	tt.Expect(t, allDoc, docs)

	has := engine.HasDoc(5)
	tt.Expect(t, "false", has)

	has = engine.HasDoc(2)
	tt.Equal(t, true, has)
	has = engine.HasDoc(3)
	tt.Equal(t, true, has)
	has = engine.HasDoc(4)
	tt.Expect(t, "true", has)

	dbhas := engine.HasDocDB(5)
	tt.Expect(t, "false", dbhas)

	dbhas = engine.HasDocDB(2)
	tt.Equal(t, true, dbhas)
	dbhas = engine.HasDocDB(3)
	tt.Equal(t, true, dbhas)
	dbhas = engine.HasDocDB(4)
	tt.Expect(t, "true", dbhas)

	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[1] = true

	outputs := engine.Search(types.SearchReq{
		Text:   "World人口",
		DocIds: docIds})

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		fmt.Println("output docs: ", outputs)
		tt.Expect(t, "1", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "1", outputs.NumDocs)

	engine.Close()
	os.RemoveAll("riot.id")
}

func testOpts(use int, store string) types.EngineOpts {
	return types.EngineOpts{
		// Using:      1,
		Using:       use,
		UseStore:    true,
		StoreFolder: store,
		IDOnly:      true,
		GseDict:     "./testdata/test_dict.txt",
	}
}

func TestDocPinYin(t *testing.T) {
	var engine Engine
	engine.Init(testOpts(0, "riot.py"))

	// AddDocs(&engine)
	// engine.RemoveDoc(5)

	tokens := engine.PinYin("在路上, in the way")

	fmt.Println("tokens...", tokens)
	tt.Expect(t, "52", len(tokens))

	var tokenDatas []types.TokenData
	// tokens := []string{"z", "zl"}
	for i := 0; i < len(tokens); i++ {
		tokenData := types.TokenData{Text: tokens[i]}
		tokenDatas = append(tokenDatas, tokenData)
	}

	index1 := types.DocData{Tokens: tokenDatas, Fields: "在路上"}
	index2 := types.DocData{Content: "在路上, in the way",
		Tokens: tokenDatas}

	engine.Index(10, index1)
	engine.Index(11, index2)

	engine.Flush()

	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[10] = true
	docIds[11] = true

	outputs := engine.Search(types.SearchReq{
		Text:   "zl",
		DocIds: docIds,
	})

	fmt.Println("outputs", outputs.Docs)
	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		tt.Expect(t, "2", len(outDocs))
		// tt.Expect(t, "11", outDocs[0].DocId)
		// tt.Expect(t, "10", outDocs[1].DocId)
	}
	tt.Expect(t, "1", len(outputs.Tokens))
	tt.Expect(t, "2", outputs.NumDocs)

	engine.Close()
	os.RemoveAll("riot.py")
}

func TestForSplitData(t *testing.T) {
	var engine Engine
	engine.Init(testOpts(4, "riot.data"))

	AddDocs(&engine)

	engine.RemoveDoc(5)
	engine.Flush()

	tokenDatas := engine.PinYin("在路上, in the way")
	tokens, num := engine.ForSplitData(tokenDatas, 52)
	tt.Expect(t, "93", len(tokens))
	tt.Expect(t, "104", num)

	index1 := types.DocData{Content: "在路上"}
	engine.Index(10, index1, true)

	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[1] = true
	outputs := engine.Search(types.SearchReq{
		Text:   "World人口",
		DocIds: docIds})

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		tt.Expect(t, "0", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "0", outputs.NumDocs)

	engine.Close()
	os.RemoveAll("riot.data")
}

func testNum(t *testing.T, numAdd, numInx, numRm uint64) {
	tt.Expect(t, "26", numAdd)
	tt.Expect(t, "6", numInx)
	tt.Expect(t, "8", numRm)
}
func TestDocCounters(t *testing.T) {
	var engine Engine
	engine.Init(testOpts(1, "riot.doc"))

	AddDocs(&engine)
	engine.RemoveDoc(5)
	engine.Flush()

	numAdd := engine.NumTokenAdded()
	numInx := engine.NumIndexed()
	numRm := engine.NumRemoved()
	testNum(t, numAdd, numInx, numRm)

	numAdd = engine.NumTokenIndexAdded()
	numInx = engine.NumDocsIndexed()
	numRm = engine.NumDocsRemoved()
	testNum(t, numAdd, numInx, numRm)

	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[1] = true

	outputs := engine.Search(types.SearchReq{
		Text:   "World人口",
		DocIds: docIds})

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		tt.Expect(t, "1", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "1", outputs.NumDocs)

	engine.Close()
	os.RemoveAll("riot.doc")
}
<commit_msg>add pinyin engine option init test<commit_after>package riot

import (
	"encoding/gob"
	"fmt"
	"log"
	"os"
	"testing"

	"github.com/go-ego/riot/types"
	"github.com/vcaesar/tt"
)

func makeDocIds() map[uint64]bool {
	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[3] = true
	docIds[1] = true
	docIds[2] = true

	return docIds
}

func TestEngineIndexWithNewStore(t *testing.T) {
	gob.Register(ScoringFields{})
	var engine = New("./testdata/test_dict.txt", "./riot.new", 8)
	log.Println("new engine start...")
	// engine = engine.New()
	AddDocs(engine)

	engine.RemoveDoc(5, true)
	engine.Flush()

	engine.Close()
	// os.RemoveAll("riot.new")

	// var engine1 = New("./testdata/test_dict.txt", "./riot.new")
	var engine1 = New("./testdata/test_new.toml")
	// engine1 = engine1.New()
	log.Println("test...")
	engine1.Flush()
	log.Println("new engine1 start...")

	outputs := engine1.Search(types.SearchReq{Text: "World人口"})
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "world", outputs.Tokens[0])
	tt.Expect(t, "人口", outputs.Tokens[1])

	outDocs := outputs.Docs.(types.ScoredDocs)
	tt.Expect(t, "2", len(outDocs))

	// tt.Expect(t, "2", outDocs[0].DocId)
	tt.Expect(t, "2500", int(outDocs[0].Scores[0]*1000))
	tt.Expect(t, "[]", outDocs[0].TokenSnippetLocs)

	// tt.Expect(t, "1", outDocs[1].DocId)
	tt.Expect(t, "2215", int(outDocs[1].Scores[0]*1000))
	tt.Expect(t, "[]", outDocs[1].TokenSnippetLocs)

	engine1.Close()
	os.RemoveAll("riot.new")
	// os.RemoveAll("riot-index")
}

var (
	rankTestOpts = rankOptsMax(0, 1)
)

func testRankOpt(idOnly bool) types.EngineOpts {
	return types.EngineOpts{
		Using:       1,
		IDOnly:      idOnly,
		GseDict:     "./testdata/test_dict.txt",
		DefRankOpts: &rankTestOpts,
		IndexerOpts: &types.IndexerOpts{
			IndexType: types.LocsIndex,
		},
	}
}

func lookupReq(engine *Engine) (types.SearchReq, []string, chan rankerReturnReq) {
	request := types.SearchReq{
		Text:   "World人口",
		DocIds: makeDocIds(),
	}

	tokens := engine.Tokens(request)
	// 建立排序器返回的通信通道
	rankerReturnChan := make(
		chan rankerReturnReq, engine.initOptions.NumShards)

	// 生成查找请求
	lookupRequest := indexerLookupReq{
		countDocsOnly:    request.CountDocsOnly,
		tokens:           tokens,
		labels:           request.Labels,
		docIds:           request.DocIds,
		options:          rankTestOpts,
		rankerReturnChan: rankerReturnChan,
		orderless:        request.Orderless,
		logic:            request.Logic,
	}

	// 向索引器发送查找请求
	for shard := 0; shard < engine.initOptions.NumShards; shard++ {
		engine.indexerLookupChans[shard] <- lookupRequest
	}

	return request, tokens, rankerReturnChan
}

func TestDocRankID(t *testing.T) {
	var engine Engine

	engine.Init(testRankOpt(true))
	AddDocs(&engine)

	engine.RemoveDoc(5)
	engine.Flush()

	request, tokens, rankerReturnChan := lookupReq(&engine)
	outputs := engine.RankID(request, rankTestOpts, tokens, rankerReturnChan)

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		tt.Expect(t, "1", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "2", outputs.NumDocs)

	engine.Close()
}

func TestDocRanks(t *testing.T) {
	var engine Engine

	engine.Init(testRankOpt(false))
	AddDocs(&engine)

	engine.RemoveDoc(5)
	engine.Flush()

	request, tokens, rankerReturnChan := lookupReq(&engine)
	outputs := engine.Ranks(request, rankTestOpts, tokens, rankerReturnChan)

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredDocs)
		tt.Expect(t, "1", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "2", outputs.NumDocs)

	// test search
	outputs1 := engine.Search(types.SearchReq{
		Text:    "World人口",
		Timeout: 10,
		DocIds:  makeDocIds()})

	if outputs1.Docs != nil {
		outDocs1 := outputs.Docs.(types.ScoredDocs)
		tt.Expect(t, "1", len(outDocs1))
	}
	tt.Expect(t, "2", len(outputs1.Tokens))
	tt.Expect(t, "2", outputs1.NumDocs)

	engine.Close()
}

func TestDocGetAllDocAndID(t *testing.T) {
	gob.Register(ScoringFields{})

	var engine Engine
	opts := types.EngineOpts{
		Using:     1,
		NumShards: 5,
		UseStore:  true,
		// StoreEngine: "bg",
		StoreFolder: "riot.id",
		IDOnly:      true,
		GseDict:     "./testdata/test_dict.txt",
		DefRankOpts: &rankTestOpts,
		IndexerOpts: &types.IndexerOpts{
			IndexType: types.LocsIndex,
		},
	}
	engine.Init(opts)

	AddDocs(&engine)

	engine.RemoveDoc(5)
	engine.Flush()

	allIds := engine.GetDBAllIds()
	fmt.Println("all id", allIds)
	tt.Expect(t, "5", len(allIds))
	tt.Expect(t, "[3 4 1 6 2]", allIds)

	allIds = engine.GetAllDocIds()
	fmt.Println("all doc id", allIds)
	tt.Expect(t, "5", len(allIds))
	tt.Expect(t, "[3 4 1 6 2]", allIds)

	ids, docs := engine.GetDBAllDocs()
	fmt.Println("all id and doc", allIds, docs)
	tt.Expect(t, "5", len(ids))
	tt.Expect(t, "5", len(docs))
	tt.Expect(t, "[3 4 1 6 2]", ids)
	allDoc := `[{The world <nil> [] [] <nil>} {有人口 <nil> [] [] {2 3 1}} {The world, 有七十亿人口人口 <nil> [] [] {1 2 3}} {有七十亿人口 <nil> [] [] {2 3 3}} {The world, 人口 <nil> [] [] <nil>}]`
	tt.Expect(t, allDoc, docs)

	has := engine.HasDoc(5)
	tt.Expect(t, "false", has)

	has = engine.HasDoc(2)
	tt.Equal(t, true, has)
	has = engine.HasDoc(3)
	tt.Equal(t, true, has)
	has = engine.HasDoc(4)
	tt.Expect(t, "true", has)

	dbhas := engine.HasDocDB(5)
	tt.Expect(t, "false", dbhas)

	dbhas = engine.HasDocDB(2)
	tt.Equal(t, true, dbhas)
	dbhas = engine.HasDocDB(3)
	tt.Equal(t, true, dbhas)
	dbhas = engine.HasDocDB(4)
	tt.Expect(t, "true", dbhas)

	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[1] = true

	outputs := engine.Search(types.SearchReq{
		Text:   "World人口",
		DocIds: docIds})

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		fmt.Println("output docs: ", outputs)
		tt.Expect(t, "1", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "1", outputs.NumDocs)

	engine.Close()
	os.RemoveAll("riot.id")
}

func testOpts(use int, store string, args ...bool) types.EngineOpts {
	var pinyin bool
	if len(args) > 0 {
		pinyin = args[0]
	}

	return types.EngineOpts{
		// Using:      1,
		Using:       use,
		UseStore:    true,
		StoreFolder: store,
		PinYin:      pinyin,
		IDOnly:      true,
		GseDict:     "./testdata/test_dict.txt",
	}
}

func TestDocPinYin(t *testing.T) {
	var engine, pinyinOpt Engine
	engine.Init(testOpts(0, "riot.py"))
	pinyinOpt.Init(testOpts(0, "riot.py.opt", true))

	// AddDocs(&engine)
	// engine.RemoveDoc(5)

	text := "在路上, in the way"

	tokens := engine.PinYin(text)
	fmt.Println("tokens...", tokens)
	tt.Expect(t, "52", len(tokens))

	var tokenDatas []types.TokenData
	// tokens := []string{"z", "zl"}
	for i := 0; i < len(tokens); i++ {
		tokenData := types.TokenData{Text: tokens[i]}
		tokenDatas = append(tokenDatas, tokenData)
	}

	index1 := types.DocData{Tokens: tokenDatas, Fields: "在路上"}
	index2 := types.DocData{Content: text, Tokens: tokenDatas}

	engine.Index(10, index1)
	engine.Index(11, index2)
	engine.Flush()

	data := types.DocData{Content: text}
	pinyinOpt.Index(10, data)
	pinyinOpt.Index(11, data)
	pinyinOpt.Flush()

	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[10] = true
	docIds[11] = true

	pyOutputs := pinyinOpt.SearchID(types.SearchReq{
		Text:   "zl",
		DocIds: docIds,
	})

	tt.Expect(t, "2", len(pyOutputs.Docs))
	tt.Expect(t, "1", len(pyOutputs.Tokens))
	tt.Expect(t, "2", pyOutputs.NumDocs)

	outputs := engine.Search(types.SearchReq{
		Text:   "zl",
		DocIds: docIds,
	})

	fmt.Println("outputs", outputs.Docs)
	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		tt.Expect(t, "2", len(outDocs))
		// tt.Expect(t, "11", outDocs[0].DocId)
		// tt.Expect(t, "10", outDocs[1].DocId)
	}
	tt.Expect(t, "1", len(outputs.Tokens))
	tt.Expect(t, "2", outputs.NumDocs)

	engine.Close()
	pinyinOpt.Close()
	os.RemoveAll("riot.py")
	os.RemoveAll("riot.py.opt")
}

func TestForSplitData(t *testing.T) {
	var engine Engine
	engine.Init(testOpts(4, "riot.data"))

	AddDocs(&engine)

	engine.RemoveDoc(5)
	engine.Flush()

	tokenDatas := engine.PinYin("在路上, in the way")
	tokens, num := engine.ForSplitData(tokenDatas, 52)
	tt.Expect(t, "93", len(tokens))
	tt.Expect(t, "104", num)

	index1 := types.DocData{Content: "在路上"}
	engine.Index(10, index1, true)

	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[1] = true
	outputs := engine.Search(types.SearchReq{
		Text:   "World人口",
		DocIds: docIds})

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		tt.Expect(t, "0", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "0", outputs.NumDocs)

	engine.Close()
	os.RemoveAll("riot.data")
}

func testNum(t *testing.T, numAdd, numInx, numRm uint64) {
	tt.Expect(t, "26", numAdd)
	tt.Expect(t, "6", numInx)
	tt.Expect(t, "8", numRm)
}
func TestDocCounters(t *testing.T) {
	var engine Engine
	engine.Init(testOpts(1, "riot.doc"))

	AddDocs(&engine)
	engine.RemoveDoc(5)
	engine.Flush()

	numAdd := engine.NumTokenAdded()
	numInx := engine.NumIndexed()
	numRm := engine.NumRemoved()
	testNum(t, numAdd, numInx, numRm)

	numAdd = engine.NumTokenIndexAdded()
	numInx = engine.NumDocsIndexed()
	numRm = engine.NumDocsRemoved()
	testNum(t, numAdd, numInx, numRm)

	docIds := make(map[uint64]bool)
	docIds[5] = true
	docIds[1] = true

	outputs := engine.Search(types.SearchReq{
		Text:   "World人口",
		DocIds: docIds})

	if outputs.Docs != nil {
		outDocs := outputs.Docs.(types.ScoredIDs)
		tt.Expect(t, "1", len(outDocs))
	}
	tt.Expect(t, "2", len(outputs.Tokens))
	tt.Expect(t, "1", outputs.NumDocs)

	engine.Close()
	os.RemoveAll("riot.doc")
}
<|endoftext|>


<commit_before>package bird

import (
	"regexp"
	"strconv"
	"strings"
)

var ParserConf ParserConfig

type Parsed map[string]interface{}

func dirtyContains(l []string, e string) bool {
	for _, c := range l {
		if c == e {
			return true
		}
	}
	return false
}

func emptyLine(line string) bool {
	return len(strings.TrimSpace(line)) == 0
}

func getLinesUnfiltered(input string) []string {
	line_sep := regexp.MustCompile(`((\r?\n)|(\r\n?))`)
	return line_sep.Split(input, -1)
}

func getLinesFromString(input string) []string {
	lines := getLinesUnfiltered(input)

	var filtered []string

	for _, line := range lines {
		if !emptyLine(line) {
			filtered = append(filtered, line)
		}
	}

	return filtered
}

func getLines(input []byte) []string {
	return getLinesFromString(string(input))
}

func specialLine(line string) bool {
	return (strings.HasPrefix(line, "BIRD") ||
		strings.HasPrefix(line, "Access restricted"))
}

func parseStatus(input []byte) Parsed {
	res := Parsed{}
	lines := getLines(input)

	start_line_rx := regexp.MustCompile(`^BIRD\s([0-9\.]+)\s*$`)
	router_id_rx := regexp.MustCompile(`^Router\sID\sis\s([0-9\.]+)\s*$`)
	current_server_rx := regexp.MustCompile(`^Current\sserver\stime\sis\s([0-9\-]+\s[0-9\:]+)\s*$`)
	last_reboot_rx := regexp.MustCompile(`^Last\sreboot\son\s([0-9\-]+\s[0-9\:]+)\s*$`)
	last_reconfig_rx := regexp.MustCompile(`^Last\sreconfiguration\son\s([0-9\-]+\s[0-9\:]+)\s*$`)

	for _, line := range lines {
		if start_line_rx.MatchString(line) {
			res["version"] = start_line_rx.FindStringSubmatch(line)[1]
		} else if router_id_rx.MatchString(line) {
			res["router_id"] = router_id_rx.FindStringSubmatch(line)[1]
		} else if current_server_rx.MatchString(line) {
			res["current_server"] = current_server_rx.FindStringSubmatch(line)[1]
		} else if last_reboot_rx.MatchString(line) {
			res["last_reboot"] = last_reboot_rx.FindStringSubmatch(line)[1]
		} else if last_reconfig_rx.MatchString(line) {
			res["last_reconfig"] = last_reconfig_rx.FindStringSubmatch(line)[1]
		} else {
			res["message"] = line
		}
	}

	for k, _ := range res {
		if dirtyContains(ParserConf.FilterFields, k) {
			res[k] = nil
		}
	}

	return Parsed{"status": res}
}

func parseProtocols(input []byte) Parsed {
	res := Parsed{}
	protocols := []string{}
	lines := getLinesUnfiltered(string(input))

	proto := ""
	for _, line := range lines {
		if emptyLine(line) {
			if !emptyLine(proto) {
				protocols = append(protocols, proto)
			}
			proto = ""
		} else {
			proto += (line + "\n")
		}
	}

	res["protocols"] = protocols
	return res
}

func parseSymbols(input []byte) Parsed {
	res := Parsed{}
	lines := getLines(input)

	key_rx := regexp.MustCompile(`^([^\s]+)\s+(.+)\s*$`)
	for _, line := range lines {
		if specialLine(line) {
			continue
		}

		if key_rx.MatchString(line) {
			groups := key_rx.FindStringSubmatch(line)
			res[groups[2]] = groups[1]
		}
	}

	return Parsed{"symbols": res}
}

func mainRouteDetail(groups []string, route Parsed) Parsed {
	route["network"] = groups[1]
	route["gateway"] = groups[2]
	route["interface"] = groups[3]
	route["from_protocol"] = groups[4]
	route["age"] = groups[5]
	route["learnt_from"] = groups[6]
	route["primary"] = groups[7] == "*"
	route["metric"] = parseInt(groups[8])

	for k, _ := range route {
		if dirtyContains(ParserConf.FilterFields, k) {
			route[k] = nil
		}
	}

	return route
}

func parseRoutes(input []byte) Parsed {
	res := Parsed{}
	lines := getLines(input)

	routes := []Parsed{}

	route := Parsed{}
	start_def_rx := regexp.MustCompile(`^([0-9a-f\.\:\/]+)\s+via\s+([0-9a-f\.\:]+)\s+on\s+([\w\.]+)\s+\[([\w\.:]+)\s+([0-9\-\:\s]+)(?:\s+from\s+([0-9a-f\.\:\/]+)){0,1}\]\s+(?:(\*)\s+){0,1}\((\d+)(?:\/\d+){0,1}\).*`)
	second_rx := regexp.MustCompile(`^\s+via\s+([0-9a-f\.\:]+)\s+on\s+([\w+)\s+\[([\w\.]+)\s+([0-9\-\:\s]+)(?:\s+from\s+([0-9a-f\.\:\/]+)){0,1}\]\s+(?:(\*)\s+){0,1}\((\d+)(?:\/\d+){0,1}\).*$`)
	type_rx := regexp.MustCompile(`^\s+Type:\s+(.*)\s*$`)
	bgp_rx := regexp.MustCompile(`^\s+BGP.(\w+):\s+(.+)\s*$`)
	community_rx := regexp.MustCompile(`^\((\d+),\s*(\d+)\)`)
	large_community_rx := regexp.MustCompile(`^\((\d+),\s*(\d+),\s*(\d+)\)`)
	paren_rx := regexp.MustCompile(`\([^\(]*\)\s*`)
	for _, line := range lines {
		if specialLine(line) || (len(route) == 0 && emptyLine(line)) {
			continue
		}

		if start_def_rx.MatchString(line) {
			if len(route) > 0 {
				routes = append(routes, route)
				route = Parsed{}
			}
			route = mainRouteDetail(start_def_rx.FindStringSubmatch(line), route)
		} else if second_rx.MatchString(line) {
			routes = append(routes, route)
			var network string
			if tmp, ok := route["network"]; ok {
				if val, ok := tmp.(string); ok {
					network = val
				} else {
					continue
				}
			} else {
				continue
			}
			route = Parsed{}

			groups := second_rx.FindStringSubmatch(line)
			first, groups := groups[0], groups[1:]
			groups = append([]string{network}, groups...)
			groups = append([]string{first}, groups...)
			route = mainRouteDetail(groups, route)
		} else if type_rx.MatchString(line) {
			submatch := type_rx.FindStringSubmatch(line)[1]
			route["type"] = strings.Split(submatch, " ")
		} else if bgp_rx.MatchString(line) {
			groups := bgp_rx.FindStringSubmatch(line)
			bgp := Parsed{}

			if tmp, ok := route["bgp"]; ok {
				if val, ok := tmp.(Parsed); ok {
					bgp = val
				}
			}

			if groups[1] == "community" {
				communities := [][]int64{}
				for _, community := range paren_rx.FindAllString(groups[2], -1) {
					if community_rx.MatchString(community) {
						com_groups := community_rx.FindStringSubmatch(community)
						maj := parseInt(com_groups[1])
						min := parseInt(com_groups[2])
						communities = append(communities, []int64{maj, min})
					}
				}
				bgp["communities"] = communities
			} else if groups[1] == "large_community" {
				communities := [][]int64{}
				for _, community := range paren_rx.FindAllString(groups[2], -1) {
					if large_community_rx.MatchString(community) {
						com_groups := large_community_rx.FindStringSubmatch(community)
						maj := parseInt(com_groups[1])
						min := parseInt(com_groups[2])
						pat := parseInt(com_groups[3])
						communities = append(communities, []int64{maj, min, pat})
					}
				}
				bgp["large_communities"] = communities
			} else if groups[1] == "as_path" {
				bgp["as_path"] = strings.Split(groups[2], " ")
			} else {
				bgp[groups[1]] = groups[2]
			}

			route["bgp"] = bgp
		}
	}

	if len(route) > 0 {
		routes = append(routes, route)
	}

	res["routes"] = routes
	return res
}

func parseRoutesCount(input []byte) Parsed {
	res := Parsed{}
	lines := getLines(input)

	count_rx := regexp.MustCompile(`^(\d+)\s+of\s+(\d+)\s+routes.*$`)
	for _, line := range lines {
		if specialLine(line) {
			continue
		}

		if count_rx.MatchString(line) {
			count := count_rx.FindStringSubmatch(line)[1]
			res["routes"] = parseInt(count)
		}
	}

	return res
}

// Will snake_case a value like that:
// I am a Weird stRiNg -> i_am_a_weird_string
func treatKey(key string) string {
	spaces := regexp.MustCompile(`\s+`)
	key = spaces.ReplaceAllString(key, "_")
	return strings.ToLower(key)
}

func parseInt(from string) int64 {
	val, err := strconv.ParseInt(from, 10, 64)
	if err != nil {
		return 0
	}
	return val
}

func parseBgpRoutes(input string) Parsed {
	routes := Parsed{}
	// Input: 1 imported, 0 filtered, 2 exported, 1 preferred
	tokens := strings.Split(input, ",")
	for _, token := range tokens {
		token = strings.TrimSpace(token)
		content := strings.Split(token, " ")
		routes[content[1]] = parseInt(content[0])
	}
	return routes
}

func parseBgp(input string) Parsed {
	res := Parsed{}
	lines := getLinesFromString(input)
	route_changes := Parsed{}

	bgp_rx := regexp.MustCompile(`^([\w\.:]+)\s+BGP\s+(\w+)\s+(\w+)\s+([0-9]{4}-[0-9]{2}-[0-9]{2}\s+[0-9]{2}:[0-9]{2}:[0-9]{2})\s*(\w+)?.*$`)
	num_val_rx := regexp.MustCompile(`^\s+([^:]+):\s+([\d]+)\s*$`)
	str_val_rx := regexp.MustCompile(`^\s+([^:]+):\s+(.+)\s*$`)

	// routes_rx := regexp.MustCompile(`^\s+Routes:\s+(\d+)\s+imported,\s+(\d+)\s+filtered,\s+(\d+)\s+exported,\s+(\d+)\s+preferred\s*$`)
	routes_rx := regexp.MustCompile(`^\s+Routes:\s+(.*)`)

	imp_updates_rx := regexp.MustCompile(`^\s+Import updates:\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s*$`)
	imp_withdraws_rx := regexp.MustCompile(`^\s+Import withdraws:\s+(\d+)\s+(\d+)\s+\-\-\-\s+(\d+)\s+(\d+)\s*$`)
	exp_updates_rx := regexp.MustCompile(`^\s+Export updates:\s+(\d+)\s+(\d+)\s+(\d+)\s+\-\-\-\s+(\d+)\s*$`)
	exp_withdraws_rx := regexp.MustCompile(`^\s+Export withdraws:\s+(\d+)(\s+\-\-\-){2}\s+(\d+)\s*$`)
	for _, line := range lines {
		if bgp_rx.MatchString(line) {
			groups := bgp_rx.FindStringSubmatch(line)

			res["protocol"] = groups[1]
			res["bird_protocol"] = "BGP"
			res["table"] = groups[2]
			res["state"] = groups[3]
			res["state_changed"] = groups[4]
			res["connection"] = groups[5]
		} else if routes_rx.MatchString(line) {
			groups := routes_rx.FindStringSubmatch(line)
			routes := parseBgpRoutes(groups[1])
			res["routes"] = routes
		} else if imp_updates_rx.MatchString(line) {
			updates := Parsed{}
			groups := imp_updates_rx.FindStringSubmatch(line)

			updates["received"] = parseInt(groups[1])
			updates["rejected"] = parseInt(groups[2])
			updates["filtered"] = parseInt(groups[3])
			updates["ignored"] = parseInt(groups[4])
			updates["accepted"] = parseInt(groups[5])

			route_changes["import_updates"] = updates
		} else if imp_withdraws_rx.MatchString(line) {
			updates := Parsed{}
			groups := imp_withdraws_rx.FindStringSubmatch(line)

			updates["received"] = parseInt(groups[1])
			updates["rejected"] = parseInt(groups[2])
			updates["filtered"] = parseInt(groups[3])
			updates["accepted"] = parseInt(groups[4])

			route_changes["import_withdraws"] = updates
		} else if exp_updates_rx.MatchString(line) {
			updates := Parsed{}
			groups := exp_updates_rx.FindStringSubmatch(line)

			updates["received"] = parseInt(groups[1])
			updates["rejected"] = parseInt(groups[2])
			updates["ignored"] = parseInt(groups[3])
			updates["accepted"] = parseInt(groups[4])

			route_changes["export_updates"] = updates
		} else if exp_withdraws_rx.MatchString(line) {
			updates := Parsed{}
			groups := exp_withdraws_rx.FindStringSubmatch(line)

			updates["received"] = parseInt(groups[1])
			updates["accepted"] = parseInt(groups[3])

			route_changes["export_withdraws"] = updates
		} else if num_val_rx.MatchString(line) {
			groups := num_val_rx.FindStringSubmatch(line)

			key := treatKey(groups[1])

			res[key] = parseInt(groups[2])
		} else if str_val_rx.MatchString(line) {
			groups := str_val_rx.FindStringSubmatch(line)

			key := treatKey(groups[1])

			res[key] = groups[2]
		}
	}

	res["route_changes"] = route_changes

	if _, ok := res["routes"]; !ok {
		routes := Parsed{}

		routes["accepted"] = 0
		routes["filtered"] = 0
		routes["exported"] = 0
		routes["preferred"] = 0

		res["routes"] = routes
	}

	return res
}
<commit_msg>fixed subsequent routes parsing<commit_after>package bird

import (
	"regexp"
	"strconv"
	"strings"
)

var ParserConf ParserConfig

type Parsed map[string]interface{}

func dirtyContains(l []string, e string) bool {
	for _, c := range l {
		if c == e {
			return true
		}
	}
	return false
}

func emptyLine(line string) bool {
	return len(strings.TrimSpace(line)) == 0
}

func getLinesUnfiltered(input string) []string {
	line_sep := regexp.MustCompile(`((\r?\n)|(\r\n?))`)
	return line_sep.Split(input, -1)
}

func getLinesFromString(input string) []string {
	lines := getLinesUnfiltered(input)

	var filtered []string

	for _, line := range lines {
		if !emptyLine(line) {
			filtered = append(filtered, line)
		}
	}

	return filtered
}

func getLines(input []byte) []string {
	return getLinesFromString(string(input))
}

func specialLine(line string) bool {
	return (strings.HasPrefix(line, "BIRD") ||
		strings.HasPrefix(line, "Access restricted"))
}

func parseStatus(input []byte) Parsed {
	res := Parsed{}
	lines := getLines(input)

	start_line_rx := regexp.MustCompile(`^BIRD\s([0-9\.]+)\s*$`)
	router_id_rx := regexp.MustCompile(`^Router\sID\sis\s([0-9\.]+)\s*$`)
	current_server_rx := regexp.MustCompile(`^Current\sserver\stime\sis\s([0-9\-]+\s[0-9\:]+)\s*$`)
	last_reboot_rx := regexp.MustCompile(`^Last\sreboot\son\s([0-9\-]+\s[0-9\:]+)\s*$`)
	last_reconfig_rx := regexp.MustCompile(`^Last\sreconfiguration\son\s([0-9\-]+\s[0-9\:]+)\s*$`)

	for _, line := range lines {
		if start_line_rx.MatchString(line) {
			res["version"] = start_line_rx.FindStringSubmatch(line)[1]
		} else if router_id_rx.MatchString(line) {
			res["router_id"] = router_id_rx.FindStringSubmatch(line)[1]
		} else if current_server_rx.MatchString(line) {
			res["current_server"] = current_server_rx.FindStringSubmatch(line)[1]
		} else if last_reboot_rx.MatchString(line) {
			res["last_reboot"] = last_reboot_rx.FindStringSubmatch(line)[1]
		} else if last_reconfig_rx.MatchString(line) {
			res["last_reconfig"] = last_reconfig_rx.FindStringSubmatch(line)[1]
		} else {
			res["message"] = line
		}
	}

	for k, _ := range res {
		if dirtyContains(ParserConf.FilterFields, k) {
			res[k] = nil
		}
	}

	return Parsed{"status": res}
}

func parseProtocols(input []byte) Parsed {
	res := Parsed{}
	protocols := []string{}
	lines := getLinesUnfiltered(string(input))

	proto := ""
	for _, line := range lines {
		if emptyLine(line) {
			if !emptyLine(proto) {
				protocols = append(protocols, proto)
			}
			proto = ""
		} else {
			proto += (line + "\n")
		}
	}

	res["protocols"] = protocols
	return res
}

func parseSymbols(input []byte) Parsed {
	res := Parsed{}
	lines := getLines(input)

	key_rx := regexp.MustCompile(`^([^\s]+)\s+(.+)\s*$`)
	for _, line := range lines {
		if specialLine(line) {
			continue
		}

		if key_rx.MatchString(line) {
			groups := key_rx.FindStringSubmatch(line)
			res[groups[2]] = groups[1]
		}
	}

	return Parsed{"symbols": res}
}

func mainRouteDetail(groups []string, route Parsed) Parsed {
	route["network"] = groups[1]
	route["gateway"] = groups[2]
	route["interface"] = groups[3]
	route["from_protocol"] = groups[4]
	route["age"] = groups[5]
	route["learnt_from"] = groups[6]
	route["primary"] = groups[7] == "*"
	route["metric"] = parseInt(groups[8])

	for k, _ := range route {
		if dirtyContains(ParserConf.FilterFields, k) {
			route[k] = nil
		}
	}

	return route
}

func parseRoutes(input []byte) Parsed {
	res := Parsed{}
	lines := getLines(input)

	routes := []Parsed{}

	route := Parsed{}
	start_def_rx := regexp.MustCompile(`^([0-9a-f\.\:\/]+)\s+via\s+([0-9a-f\.\:]+)\s+on\s+([\w\.]+)\s+\[([\w\.:]+)\s+([0-9\-\:\s]+)(?:\s+from\s+([0-9a-f\.\:\/]+)){0,1}\]\s+(?:(\*)\s+){0,1}\((\d+)(?:\/\d+){0,1}\).*`)
	second_rx := regexp.MustCompile(`^\s+via\s+([0-9a-f\.\:]+)\s+on\s+([\w\.]+)\s+\[([\w\.:]+)\s+([0-9\-\:\s]+)(?:\s+from\s+([0-9a-f\.\:\/]+)){0,1}\]\s+(?:(\*)\s+){0,1}\((\d+)(?:\/\d+){0,1}\).*$`)
	type_rx := regexp.MustCompile(`^\s+Type:\s+(.*)\s*$`)
	bgp_rx := regexp.MustCompile(`^\s+BGP.(\w+):\s+(.+)\s*$`)
	community_rx := regexp.MustCompile(`^\((\d+),\s*(\d+)\)`)
	large_community_rx := regexp.MustCompile(`^\((\d+),\s*(\d+),\s*(\d+)\)`)
	paren_rx := regexp.MustCompile(`\([^\(]*\)\s*`)
	for _, line := range lines {
		if specialLine(line) || (len(route) == 0 && emptyLine(line)) {
			continue
		}

		if start_def_rx.MatchString(line) {
			if len(route) > 0 {
				routes = append(routes, route)
				route = Parsed{}
			}
			route = mainRouteDetail(start_def_rx.FindStringSubmatch(line), route)
		} else if second_rx.MatchString(line) {
			routes = append(routes, route)
			var network string
			if tmp, ok := route["network"]; ok {
				if val, ok := tmp.(string); ok {
					network = val
				} else {
					continue
				}
			} else {
				continue
			}
			route = Parsed{}

			groups := second_rx.FindStringSubmatch(line)
			first, groups := groups[0], groups[1:]
			groups = append([]string{network}, groups...)
			groups = append([]string{first}, groups...)
			route = mainRouteDetail(groups, route)
		} else if type_rx.MatchString(line) {
			submatch := type_rx.FindStringSubmatch(line)[1]
			route["type"] = strings.Split(submatch, " ")
		} else if bgp_rx.MatchString(line) {
			groups := bgp_rx.FindStringSubmatch(line)
			bgp := Parsed{}

			if tmp, ok := route["bgp"]; ok {
				if val, ok := tmp.(Parsed); ok {
					bgp = val
				}
			}

			if groups[1] == "community" {
				communities := [][]int64{}
				for _, community := range paren_rx.FindAllString(groups[2], -1) {
					if community_rx.MatchString(community) {
						com_groups := community_rx.FindStringSubmatch(community)
						maj := parseInt(com_groups[1])
						min := parseInt(com_groups[2])
						communities = append(communities, []int64{maj, min})
					}
				}
				bgp["communities"] = communities
			} else if groups[1] == "large_community" {
				communities := [][]int64{}
				for _, community := range paren_rx.FindAllString(groups[2], -1) {
					if large_community_rx.MatchString(community) {
						com_groups := large_community_rx.FindStringSubmatch(community)
						maj := parseInt(com_groups[1])
						min := parseInt(com_groups[2])
						pat := parseInt(com_groups[3])
						communities = append(communities, []int64{maj, min, pat})
					}
				}
				bgp["large_communities"] = communities
			} else if groups[1] == "as_path" {
				bgp["as_path"] = strings.Split(groups[2], " ")
			} else {
				bgp[groups[1]] = groups[2]
			}

			route["bgp"] = bgp
		}
	}

	if len(route) > 0 {
		routes = append(routes, route)
	}

	res["routes"] = routes
	return res
}

func parseRoutesCount(input []byte) Parsed {
	res := Parsed{}
	lines := getLines(input)

	count_rx := regexp.MustCompile(`^(\d+)\s+of\s+(\d+)\s+routes.*$`)
	for _, line := range lines {
		if specialLine(line) {
			continue
		}

		if count_rx.MatchString(line) {
			count := count_rx.FindStringSubmatch(line)[1]
			res["routes"] = parseInt(count)
		}
	}

	return res
}

// Will snake_case a value like that:
// I am a Weird stRiNg -> i_am_a_weird_string
func treatKey(key string) string {
	spaces := regexp.MustCompile(`\s+`)
	key = spaces.ReplaceAllString(key, "_")
	return strings.ToLower(key)
}

func parseInt(from string) int64 {
	val, err := strconv.ParseInt(from, 10, 64)
	if err != nil {
		return 0
	}
	return val
}

func parseBgpRoutes(input string) Parsed {
	routes := Parsed{}
	// Input: 1 imported, 0 filtered, 2 exported, 1 preferred
	tokens := strings.Split(input, ",")
	for _, token := range tokens {
		token = strings.TrimSpace(token)
		content := strings.Split(token, " ")
		routes[content[1]] = parseInt(content[0])
	}
	return routes
}

func parseBgp(input string) Parsed {
	res := Parsed{}
	lines := getLinesFromString(input)
	route_changes := Parsed{}

	bgp_rx := regexp.MustCompile(`^([\w\.:]+)\s+BGP\s+(\w+)\s+(\w+)\s+([0-9]{4}-[0-9]{2}-[0-9]{2}\s+[0-9]{2}:[0-9]{2}:[0-9]{2})\s*(\w+)?.*$`)
	num_val_rx := regexp.MustCompile(`^\s+([^:]+):\s+([\d]+)\s*$`)
	str_val_rx := regexp.MustCompile(`^\s+([^:]+):\s+(.+)\s*$`)

	// routes_rx := regexp.MustCompile(`^\s+Routes:\s+(\d+)\s+imported,\s+(\d+)\s+filtered,\s+(\d+)\s+exported,\s+(\d+)\s+preferred\s*$`)
	routes_rx := regexp.MustCompile(`^\s+Routes:\s+(.*)`)

	imp_updates_rx := regexp.MustCompile(`^\s+Import updates:\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s*$`)
	imp_withdraws_rx := regexp.MustCompile(`^\s+Import withdraws:\s+(\d+)\s+(\d+)\s+\-\-\-\s+(\d+)\s+(\d+)\s*$`)
	exp_updates_rx := regexp.MustCompile(`^\s+Export updates:\s+(\d+)\s+(\d+)\s+(\d+)\s+\-\-\-\s+(\d+)\s*$`)
	exp_withdraws_rx := regexp.MustCompile(`^\s+Export withdraws:\s+(\d+)(\s+\-\-\-){2}\s+(\d+)\s*$`)
	for _, line := range lines {
		if bgp_rx.MatchString(line) {
			groups := bgp_rx.FindStringSubmatch(line)

			res["protocol"] = groups[1]
			res["bird_protocol"] = "BGP"
			res["table"] = groups[2]
			res["state"] = groups[3]
			res["state_changed"] = groups[4]
			res["connection"] = groups[5]
		} else if routes_rx.MatchString(line) {
			groups := routes_rx.FindStringSubmatch(line)
			routes := parseBgpRoutes(groups[1])
			res["routes"] = routes
		} else if imp_updates_rx.MatchString(line) {
			updates := Parsed{}
			groups := imp_updates_rx.FindStringSubmatch(line)

			updates["received"] = parseInt(groups[1])
			updates["rejected"] = parseInt(groups[2])
			updates["filtered"] = parseInt(groups[3])
			updates["ignored"] = parseInt(groups[4])
			updates["accepted"] = parseInt(groups[5])

			route_changes["import_updates"] = updates
		} else if imp_withdraws_rx.MatchString(line) {
			updates := Parsed{}
			groups := imp_withdraws_rx.FindStringSubmatch(line)

			updates["received"] = parseInt(groups[1])
			updates["rejected"] = parseInt(groups[2])
			updates["filtered"] = parseInt(groups[3])
			updates["accepted"] = parseInt(groups[4])

			route_changes["import_withdraws"] = updates
		} else if exp_updates_rx.MatchString(line) {
			updates := Parsed{}
			groups := exp_updates_rx.FindStringSubmatch(line)

			updates["received"] = parseInt(groups[1])
			updates["rejected"] = parseInt(groups[2])
			updates["ignored"] = parseInt(groups[3])
			updates["accepted"] = parseInt(groups[4])

			route_changes["export_updates"] = updates
		} else if exp_withdraws_rx.MatchString(line) {
			updates := Parsed{}
			groups := exp_withdraws_rx.FindStringSubmatch(line)

			updates["received"] = parseInt(groups[1])
			updates["accepted"] = parseInt(groups[3])

			route_changes["export_withdraws"] = updates
		} else if num_val_rx.MatchString(line) {
			groups := num_val_rx.FindStringSubmatch(line)

			key := treatKey(groups[1])

			res[key] = parseInt(groups[2])
		} else if str_val_rx.MatchString(line) {
			groups := str_val_rx.FindStringSubmatch(line)

			key := treatKey(groups[1])

			res[key] = groups[2]
		}
	}

	res["route_changes"] = route_changes

	if _, ok := res["routes"]; !ok {
		routes := Parsed{}

		routes["accepted"] = 0
		routes["filtered"] = 0
		routes["exported"] = 0
		routes["preferred"] = 0

		res["routes"] = routes
	}

	return res
}
<|endoftext|>


<commit_before>// run

// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

// ----------------------------------------------------------------------------
// Helper functions

func ASSERT(p bool) {
	if !p {
		// panic 0
	}
}

// ----------------------------------------------------------------------------
// Implementation of the HashMap

type KeyType interface {
	Hash() uint32
	Match(other KeyType) bool
}

type ValueType interface {
	// empty interface
}

type Entry struct {
	key KeyType
	value ValueType
}

type Array [1024]Entry

type HashMap struct {
	map_ *Array
	log2_capacity_ uint32
	occupancy_ uint32
}

func (m *HashMap) capacity() uint32 {
	return 1 << m.log2_capacity_
}

func (m *HashMap) Clear() {
	// Mark all entries as empty.
	var i uint32 = m.capacity() - 1
	for i > 0 {
		m.map_[i].key = nil
		i = i - 1
	}
	m.occupancy_ = 0
}

func (m *HashMap) Initialize (initial_log2_capacity uint32) {
	m.log2_capacity_ = initial_log2_capacity
	m.map_ = new(Array)
	m.Clear()
}

func (m *HashMap) Probe (key KeyType) *Entry {
	ASSERT(key != nil)

	var i uint32 = key.Hash() % m.capacity()
	ASSERT(0 <= i && i < m.capacity())

	ASSERT(m.occupancy_ < m.capacity())	// guarantees loop termination
	for m.map_[i].key != nil && !m.map_[i].key.Match(key) {
		i++
		if i >= m.capacity() {
			i = 0
		}
	}

	return &m.map_[i]
}

func (m *HashMap) Lookup (key KeyType, insert bool) *Entry {
	// Find a matching entry.
	var p *Entry = m.Probe(key)
		if p.key != nil {
		return p
	}

	// No entry found; insert one if necessary.
	if insert {
		p.key = key
		p.value = nil
		m.occupancy_++

		// Grow the map if we reached >= 80% occupancy.
		if m.occupancy_ + m.occupancy_/4 >= m.capacity() {
			m.Resize()
			p = m.Probe(key)
		}

		return p
	}

	// No entry found and none inserted.
	return nil
}

func (m *HashMap) Resize() {
	var hmap *Array = m.map_
	var n uint32 = m.occupancy_

	// Allocate a new map of twice the current size.
	m.Initialize(m.log2_capacity_ << 1)

	// Rehash all current entries.
	var i uint32 = 0
	for n > 0 {
		if hmap[i].key != nil {
			m.Lookup(hmap[i].key, true).value = hmap[i].value
			n = n - 1
		}
		i++
	}
}

// ----------------------------------------------------------------------------
// Test code

type Number struct {
	x uint32
}

func (n *Number) Hash() uint32 {
	return n.x * 23
}

func (n *Number) Match(other KeyType) bool {
	// var y *Number = other
	// return n.x == y.x
	return false
}

func MakeNumber (x uint32) *Number {
	var n *Number = new(Number)
	n.x = x
	return n
}

func main() {
	// func (n int) int { return n + 1; }(1)

	//print "HashMap - gri 2/8/2008\n"

	var hmap *HashMap = new(HashMap)
	hmap.Initialize(0)

	var x1 *Number = MakeNumber(1001)
	var x2 *Number = MakeNumber(2002)
	var x3 *Number = MakeNumber(3003)
	_, _, _ = x1, x2, x3

	// this doesn't work I think...
	//hmap.Lookup(x1, true)
	//hmap.Lookup(x2, true)
	//hmap.Lookup(x3, true)

	//print "done\n"
}
<commit_msg>test/hashmap.go: delete It's in an odd style and it's unclear what true purpose it serves as a test other than to be another Go program.<commit_after><|endoftext|>


<commit_before>package mysql

import (
	"container/list"
	"context"
	"database/sql"
	"fmt"
	"strconv"

	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/go-xorm/core"
	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/tsdb"
)

type MysqlQueryEndpoint struct {
	sqlEngine tsdb.SqlEngine
	log       log.Logger
}

func init() {
	tsdb.RegisterTsdbQueryEndpoint("mysql", NewMysqlQueryEndpoint)
}

func NewMysqlQueryEndpoint(datasource *models.DataSource) (tsdb.TsdbQueryEndpoint, error) {
	endpoint := &MysqlQueryEndpoint{
		log: log.New("tsdb.mysql"),
	}

	endpoint.sqlEngine = &tsdb.DefaultSqlEngine{
		MacroEngine: NewMysqlMacroEngine(),
	}

	cnnstr := fmt.Sprintf("%s:%s@%s(%s)/%s?collation=utf8mb4_unicode_ci&parseTime=true&loc=UTC",
		datasource.User,
		datasource.Password,
		"tcp",
		datasource.Url,
		datasource.Database,
	)
	endpoint.log.Debug("getEngine", "connection", cnnstr)

	if err := endpoint.sqlEngine.InitEngine("mysql", datasource, cnnstr); err != nil {
		return nil, err
	}

	return endpoint, nil
}

// Query is the main function for the MysqlExecutor
func (e *MysqlQueryEndpoint) Query(ctx context.Context, dsInfo *models.DataSource, tsdbQuery *tsdb.TsdbQuery) (*tsdb.Response, error) {
	return e.sqlEngine.Query(ctx, dsInfo, tsdbQuery, e.transformToTimeSeries, e.transformToTable)
}

func (e MysqlQueryEndpoint) transformToTable(query *tsdb.Query, rows *core.Rows, result *tsdb.QueryResult) error {
	columnNames, err := rows.Columns()
	columnCount := len(columnNames)

	if err != nil {
		return err
	}

	table := &tsdb.Table{
		Columns: make([]tsdb.TableColumn, columnCount),
		Rows:    make([]tsdb.RowValues, 0),
	}

	for i, name := range columnNames {
		table.Columns[i].Text = name
	}

	columnTypes, err := rows.ColumnTypes()
	if err != nil {
		return err
	}

	rowLimit := 1000000
	rowCount := 0

	for ; rows.Next(); rowCount++ {
		if rowCount > rowLimit {
			return fmt.Errorf("MySQL query row limit exceeded, limit %d", rowLimit)
		}

		values, err := e.getTypedRowData(columnTypes, rows)
		if err != nil {
			return err
		}

		table.Rows = append(table.Rows, values)
	}

	result.Tables = append(result.Tables, table)
	result.Meta.Set("rowCount", rowCount)
	return nil
}

func (e MysqlQueryEndpoint) getTypedRowData(types []*sql.ColumnType, rows *core.Rows) (tsdb.RowValues, error) {
	values := make([]interface{}, len(types))

	for i, stype := range types {
		e.log.Debug("type", "type", stype)
		switch stype.DatabaseTypeName() {
		case mysql.FieldTypeNameTiny:
			values[i] = new(int8)
		case mysql.FieldTypeNameInt24:
			values[i] = new(int32)
		case mysql.FieldTypeNameShort:
			values[i] = new(int16)
		case mysql.FieldTypeNameVarString:
			values[i] = new(string)
		case mysql.FieldTypeNameVarChar:
			values[i] = new(string)
		case mysql.FieldTypeNameLong:
			values[i] = new(int)
		case mysql.FieldTypeNameLongLong:
			values[i] = new(int64)
		case mysql.FieldTypeNameDouble:
			values[i] = new(float64)
		case mysql.FieldTypeNameDecimal:
			values[i] = new(float32)
		case mysql.FieldTypeNameNewDecimal:
			values[i] = new(float64)
		case mysql.FieldTypeNameFloat:
			values[i] = new(float64)
		case mysql.FieldTypeNameTimestamp:
			values[i] = new(time.Time)
		case mysql.FieldTypeNameDateTime:
			values[i] = new(time.Time)
		case mysql.FieldTypeNameTime:
			values[i] = new(string)
		case mysql.FieldTypeNameYear:
			values[i] = new(int16)
		case mysql.FieldTypeNameNULL:
			values[i] = nil
		case mysql.FieldTypeNameBit:
			values[i] = new([]byte)
		case mysql.FieldTypeNameBLOB:
			values[i] = new(string)
		case mysql.FieldTypeNameTinyBLOB:
			values[i] = new(string)
		case mysql.FieldTypeNameMediumBLOB:
			values[i] = new(string)
		case mysql.FieldTypeNameLongBLOB:
			values[i] = new(string)
		case mysql.FieldTypeNameString:
			values[i] = new(string)
		case mysql.FieldTypeNameDate:
			values[i] = new(string)
		default:
			return nil, fmt.Errorf("Database type %s not supported", stype.DatabaseTypeName())
		}
	}

	if err := rows.Scan(values...); err != nil {
		return nil, err
	}

	return values, nil
}

func (e MysqlQueryEndpoint) transformToTimeSeries(query *tsdb.Query, rows *core.Rows, result *tsdb.QueryResult) error {
	pointsBySeries := make(map[string]*tsdb.TimeSeries)
	seriesByQueryOrder := list.New()
	columnNames, err := rows.Columns()

	if err != nil {
		return err
	}

	rowData := NewStringStringScan(columnNames)
	rowLimit := 1000000
	rowCount := 0

	for ; rows.Next(); rowCount++ {
		if rowCount > rowLimit {
			return fmt.Errorf("MySQL query row limit exceeded, limit %d", rowLimit)
		}

		err := rowData.Update(rows.Rows)
		if err != nil {
			e.log.Error("MySQL response parsing", "error", err)
			return fmt.Errorf("MySQL response parsing error %v", err)
		}

		if rowData.metric == "" {
			rowData.metric = "Unknown"
		}

		if !rowData.time.Valid {
			return fmt.Errorf("Found row with no time value")
		}

		if series, exist := pointsBySeries[rowData.metric]; exist {
			series.Points = append(series.Points, tsdb.TimePoint{rowData.value, rowData.time})
		} else {
			series := &tsdb.TimeSeries{Name: rowData.metric}
			series.Points = append(series.Points, tsdb.TimePoint{rowData.value, rowData.time})
			pointsBySeries[rowData.metric] = series
			seriesByQueryOrder.PushBack(rowData.metric)
		}
	}

	for elem := seriesByQueryOrder.Front(); elem != nil; elem = elem.Next() {
		key := elem.Value.(string)
		result.Series = append(result.Series, pointsBySeries[key])
	}

	result.Meta.Set("rowCount", rowCount)
	return nil
}

type stringStringScan struct {
	rowPtrs     []interface{}
	rowValues   []string
	columnNames []string
	columnCount int

	time   null.Float
	value  null.Float
	metric string
}

func NewStringStringScan(columnNames []string) *stringStringScan {
	s := &stringStringScan{
		columnCount: len(columnNames),
		columnNames: columnNames,
		rowPtrs:     make([]interface{}, len(columnNames)),
		rowValues:   make([]string, len(columnNames)),
	}

	for i := 0; i < s.columnCount; i++ {
		s.rowPtrs[i] = new(sql.RawBytes)
	}

	return s
}

func (s *stringStringScan) Update(rows *sql.Rows) error {
	if err := rows.Scan(s.rowPtrs...); err != nil {
		return err
	}

	s.time = null.FloatFromPtr(nil)
	s.value = null.FloatFromPtr(nil)

	for i := 0; i < s.columnCount; i++ {
		if rb, ok := s.rowPtrs[i].(*sql.RawBytes); ok {
			s.rowValues[i] = string(*rb)

			switch s.columnNames[i] {
			case "time_sec":
				if sec, err := strconv.ParseInt(s.rowValues[i], 10, 64); err == nil {
					s.time = null.FloatFrom(float64(sec * 1000))
				}
			case "value":
				if value, err := strconv.ParseFloat(s.rowValues[i], 64); err == nil {
					s.value = null.FloatFrom(value)
				}
			case "metric":
				s.metric = s.rowValues[i]
			}

			*rb = nil // reset pointer to discard current value to avoid a bug
		} else {
			return fmt.Errorf("Cannot convert index %d column %s to type *sql.RawBytes", i, s.columnNames[i])
		}
	}
	return nil
}
<commit_msg>mysql: add data source support for Azure MySql<commit_after>package mysql

import (
	"container/list"
	"context"
	"database/sql"
	"fmt"
	"strconv"

	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/go-xorm/core"
	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/tsdb"
)

type MysqlQueryEndpoint struct {
	sqlEngine tsdb.SqlEngine
	log       log.Logger
}

func init() {
	tsdb.RegisterTsdbQueryEndpoint("mysql", NewMysqlQueryEndpoint)
}

func NewMysqlQueryEndpoint(datasource *models.DataSource) (tsdb.TsdbQueryEndpoint, error) {
	endpoint := &MysqlQueryEndpoint{
		log: log.New("tsdb.mysql"),
	}

	endpoint.sqlEngine = &tsdb.DefaultSqlEngine{
		MacroEngine: NewMysqlMacroEngine(),
	}

	cnnstr := fmt.Sprintf("%s:%s@%s(%s)/%s?collation=utf8mb4_unicode_ci&parseTime=true&loc=UTC&allowNativePasswords=true",
		datasource.User,
		datasource.Password,
		"tcp",
		datasource.Url,
		datasource.Database,
	)
	endpoint.log.Debug("getEngine", "connection", cnnstr)

	if err := endpoint.sqlEngine.InitEngine("mysql", datasource, cnnstr); err != nil {
		return nil, err
	}

	return endpoint, nil
}

// Query is the main function for the MysqlExecutor
func (e *MysqlQueryEndpoint) Query(ctx context.Context, dsInfo *models.DataSource, tsdbQuery *tsdb.TsdbQuery) (*tsdb.Response, error) {
	return e.sqlEngine.Query(ctx, dsInfo, tsdbQuery, e.transformToTimeSeries, e.transformToTable)
}

func (e MysqlQueryEndpoint) transformToTable(query *tsdb.Query, rows *core.Rows, result *tsdb.QueryResult) error {
	columnNames, err := rows.Columns()
	columnCount := len(columnNames)

	if err != nil {
		return err
	}

	table := &tsdb.Table{
		Columns: make([]tsdb.TableColumn, columnCount),
		Rows:    make([]tsdb.RowValues, 0),
	}

	for i, name := range columnNames {
		table.Columns[i].Text = name
	}

	columnTypes, err := rows.ColumnTypes()
	if err != nil {
		return err
	}

	rowLimit := 1000000
	rowCount := 0

	for ; rows.Next(); rowCount++ {
		if rowCount > rowLimit {
			return fmt.Errorf("MySQL query row limit exceeded, limit %d", rowLimit)
		}

		values, err := e.getTypedRowData(columnTypes, rows)
		if err != nil {
			return err
		}

		table.Rows = append(table.Rows, values)
	}

	result.Tables = append(result.Tables, table)
	result.Meta.Set("rowCount", rowCount)
	return nil
}

func (e MysqlQueryEndpoint) getTypedRowData(types []*sql.ColumnType, rows *core.Rows) (tsdb.RowValues, error) {
	values := make([]interface{}, len(types))

	for i, stype := range types {
		e.log.Debug("type", "type", stype)
		switch stype.DatabaseTypeName() {
		case mysql.FieldTypeNameTiny:
			values[i] = new(int8)
		case mysql.FieldTypeNameInt24:
			values[i] = new(int32)
		case mysql.FieldTypeNameShort:
			values[i] = new(int16)
		case mysql.FieldTypeNameVarString:
			values[i] = new(string)
		case mysql.FieldTypeNameVarChar:
			values[i] = new(string)
		case mysql.FieldTypeNameLong:
			values[i] = new(int)
		case mysql.FieldTypeNameLongLong:
			values[i] = new(int64)
		case mysql.FieldTypeNameDouble:
			values[i] = new(float64)
		case mysql.FieldTypeNameDecimal:
			values[i] = new(float32)
		case mysql.FieldTypeNameNewDecimal:
			values[i] = new(float64)
		case mysql.FieldTypeNameFloat:
			values[i] = new(float64)
		case mysql.FieldTypeNameTimestamp:
			values[i] = new(time.Time)
		case mysql.FieldTypeNameDateTime:
			values[i] = new(time.Time)
		case mysql.FieldTypeNameTime:
			values[i] = new(string)
		case mysql.FieldTypeNameYear:
			values[i] = new(int16)
		case mysql.FieldTypeNameNULL:
			values[i] = nil
		case mysql.FieldTypeNameBit:
			values[i] = new([]byte)
		case mysql.FieldTypeNameBLOB:
			values[i] = new(string)
		case mysql.FieldTypeNameTinyBLOB:
			values[i] = new(string)
		case mysql.FieldTypeNameMediumBLOB:
			values[i] = new(string)
		case mysql.FieldTypeNameLongBLOB:
			values[i] = new(string)
		case mysql.FieldTypeNameString:
			values[i] = new(string)
		case mysql.FieldTypeNameDate:
			values[i] = new(string)
		default:
			return nil, fmt.Errorf("Database type %s not supported", stype.DatabaseTypeName())
		}
	}

	if err := rows.Scan(values...); err != nil {
		return nil, err
	}

	return values, nil
}

func (e MysqlQueryEndpoint) transformToTimeSeries(query *tsdb.Query, rows *core.Rows, result *tsdb.QueryResult) error {
	pointsBySeries := make(map[string]*tsdb.TimeSeries)
	seriesByQueryOrder := list.New()
	columnNames, err := rows.Columns()

	if err != nil {
		return err
	}

	rowData := NewStringStringScan(columnNames)
	rowLimit := 1000000
	rowCount := 0

	for ; rows.Next(); rowCount++ {
		if rowCount > rowLimit {
			return fmt.Errorf("MySQL query row limit exceeded, limit %d", rowLimit)
		}

		err := rowData.Update(rows.Rows)
		if err != nil {
			e.log.Error("MySQL response parsing", "error", err)
			return fmt.Errorf("MySQL response parsing error %v", err)
		}

		if rowData.metric == "" {
			rowData.metric = "Unknown"
		}

		if !rowData.time.Valid {
			return fmt.Errorf("Found row with no time value")
		}

		if series, exist := pointsBySeries[rowData.metric]; exist {
			series.Points = append(series.Points, tsdb.TimePoint{rowData.value, rowData.time})
		} else {
			series := &tsdb.TimeSeries{Name: rowData.metric}
			series.Points = append(series.Points, tsdb.TimePoint{rowData.value, rowData.time})
			pointsBySeries[rowData.metric] = series
			seriesByQueryOrder.PushBack(rowData.metric)
		}
	}

	for elem := seriesByQueryOrder.Front(); elem != nil; elem = elem.Next() {
		key := elem.Value.(string)
		result.Series = append(result.Series, pointsBySeries[key])
	}

	result.Meta.Set("rowCount", rowCount)
	return nil
}

type stringStringScan struct {
	rowPtrs     []interface{}
	rowValues   []string
	columnNames []string
	columnCount int

	time   null.Float
	value  null.Float
	metric string
}

func NewStringStringScan(columnNames []string) *stringStringScan {
	s := &stringStringScan{
		columnCount: len(columnNames),
		columnNames: columnNames,
		rowPtrs:     make([]interface{}, len(columnNames)),
		rowValues:   make([]string, len(columnNames)),
	}

	for i := 0; i < s.columnCount; i++ {
		s.rowPtrs[i] = new(sql.RawBytes)
	}

	return s
}

func (s *stringStringScan) Update(rows *sql.Rows) error {
	if err := rows.Scan(s.rowPtrs...); err != nil {
		return err
	}

	s.time = null.FloatFromPtr(nil)
	s.value = null.FloatFromPtr(nil)

	for i := 0; i < s.columnCount; i++ {
		if rb, ok := s.rowPtrs[i].(*sql.RawBytes); ok {
			s.rowValues[i] = string(*rb)

			switch s.columnNames[i] {
			case "time_sec":
				if sec, err := strconv.ParseInt(s.rowValues[i], 10, 64); err == nil {
					s.time = null.FloatFrom(float64(sec * 1000))
				}
			case "value":
				if value, err := strconv.ParseFloat(s.rowValues[i], 64); err == nil {
					s.value = null.FloatFrom(value)
				}
			case "metric":
				s.metric = s.rowValues[i]
			}

			*rb = nil // reset pointer to discard current value to avoid a bug
		} else {
			return fmt.Errorf("Cannot convert index %d column %s to type *sql.RawBytes", i, s.columnNames[i])
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Package md decodes md and smd files.

package md

import (
	"bytes"
	"github.com/sselph/scraper/rom"
	"io"
	"io/ioutil"
	"os"
)

func init() {
	rom.RegisterFormat(".smd", decodeSMD)
	rom.RegisterFormat(".md", decodeMD)
	rom.RegisterFormat(".gen", rom.Noop)
}

func DeInterleave(p []byte) []byte {
	l := len(p)
	m := l / 2
	b := make([]byte, l)
	for i, x := range p {
		if i < m {
			b[i*2+1] = x
		} else {
			b[i*2-l] = x
		}
	}
	return b
}

type SMDReader struct {
	f *os.File
	b []byte
	r *int
}

func (r SMDReader) Read(p []byte) (int, error) {
	ll := len(p)
	rl := ll - *r.r
	l := rl + 16384 - 1 - (rl-1)%16384
	copy(p, r.b[:*r.r])
	if rl <= 0 {
		*r.r = *r.r - ll
		copy(r.b, r.b[ll:])
		return ll, nil
	}
	n := *r.r
	for i := 0; i < l/16384; i++ {
		b := make([]byte, 16384)
		x, err := io.ReadFull(r.f, b)
		if x == 0 || err != nil {
			return n, err
		}
		b = DeInterleave(b)
		if ll < n+x {
			copy(p[n:ll], b)
			copy(r.b, b[ll-n:])
			*r.r = n + x - ll
			return ll, nil
		} else {
			copy(p[n:n+16384], b)
		}
		n += x
	}
	return ll, nil
}

func (r SMDReader) Close() error {
	return r.f.Close()
}

func decodeSMD(f *os.File) (io.ReadCloser, error) {
	f.Seek(512, 0)
	i := 0
	return SMDReader{f, make([]byte, 16384), &i}, nil
}

type MDReader struct {
	r io.Reader
}

func (r MDReader) Read(p []byte) (int, error) {
	n, err := r.r.Read(p)
	return n, err
}

func (r MDReader) Close() error {
	return nil
}

func decodeMD(f *os.File) (io.ReadCloser, error) {
	b, err := ioutil.ReadAll(f)
	f.Close()
	b = DeInterleave(b)
	if err != nil {
		return nil, err
	}
	return MDReader{bytes.NewReader(b)}, nil
}
<commit_msg>Rename MD to MGD and add noop MD decoder.<commit_after>// Package md decodes md and smd files.

package md

import (
	"bytes"
	"github.com/sselph/scraper/rom"
	"io"
	"io/ioutil"
	"os"
)

func init() {
	rom.RegisterFormat(".smd", decodeSMD)
	rom.RegisterFormat(".mgd", decodeMGD)
	rom.RegisterFormat(".gen", rom.Noop)
	rom.RegisterFormat(".md", rom.Noop)
}

func DeInterleave(p []byte) []byte {
	l := len(p)
	m := l / 2
	b := make([]byte, l)
	for i, x := range p {
		if i < m {
			b[i*2+1] = x
		} else {
			b[i*2-l] = x
		}
	}
	return b
}

type SMDReader struct {
	f *os.File
	b []byte
	r *int
}

func (r SMDReader) Read(p []byte) (int, error) {
	ll := len(p)
	rl := ll - *r.r
	l := rl + 16384 - 1 - (rl-1)%16384
	copy(p, r.b[:*r.r])
	if rl <= 0 {
		*r.r = *r.r - ll
		copy(r.b, r.b[ll:])
		return ll, nil
	}
	n := *r.r
	for i := 0; i < l/16384; i++ {
		b := make([]byte, 16384)
		x, err := io.ReadFull(r.f, b)
		if x == 0 || err != nil {
			return n, err
		}
		b = DeInterleave(b)
		if ll < n+x {
			copy(p[n:ll], b)
			copy(r.b, b[ll-n:])
			*r.r = n + x - ll
			return ll, nil
		} else {
			copy(p[n:n+16384], b)
		}
		n += x
	}
	return ll, nil
}

func (r SMDReader) Close() error {
	return r.f.Close()
}

func decodeSMD(f *os.File) (io.ReadCloser, error) {
	f.Seek(512, 0)
	i := 0
	return SMDReader{f, make([]byte, 16384), &i}, nil
}

type MGDReader struct {
	r io.Reader
}

func (r MGDReader) Read(p []byte) (int, error) {
	n, err := r.r.Read(p)
	return n, err
}

func (r MGDReader) Close() error {
	return nil
}

func decodeMGD(f *os.File) (io.ReadCloser, error) {
	b, err := ioutil.ReadAll(f)
	f.Close()
	b = DeInterleave(b)
	if err != nil {
		return nil, err
	}
	return MGDReader{bytes.NewReader(b)}, nil
}
<|endoftext|>