package xray

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestRootHandler(t *testing.T) {
	var contentLength int
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		b := []byte(`200 - OK`)
		contentLength = len(b)
		w.Write(b)
	})

	ts := httptest.NewServer(Handler(handler))

	req := httptest.NewRequest("POST", ts.URL, strings.NewReader(""))
	req.Header.Set("User-Agent", "UnitTest")
	req.Header.Set("X-Forwarded-For", "127.0.0.1")

	_, err := http.DefaultTransport.RoundTrip(req)
	assert.NoError(t, err)

	s, e := testDaemon.Recv()
	assert.NoError(t, e)

	assert.Equal(t, http.StatusOK, s.HTTP.Response.Status)
	assert.Equal(t, contentLength, s.HTTP.Response.ContentLength)
	assert.Equal(t, "POST", s.HTTP.Request.Method)
	assert.Equal(t, "/", s.HTTP.Request.URL)
	assert.Equal(t, "127.0.0.1", s.HTTP.Request.ClientIP)
	assert.Equal(t, "UnitTest", s.HTTP.Request.UserAgent)
}

func TestNonRootHandler(t *testing.T) {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	ts := httptest.NewServer(Handler(handler))

	req := httptest.NewRequest("DELETE", ts.URL, strings.NewReader(""))
	req.Header.Set("x-amzn-trace-id", "Root=fakeid; Parent=reqid; Sampled=true")

	_, err := http.DefaultTransport.RoundTrip(req)
	assert.NoError(t, err)

	s, e := testDaemon.Recv()
	assert.NoError(t, e)

	assert.Equal(t, "fakeid", s.TraceID)
	assert.Equal(t, "reqid", s.ParentID)
}
