old_console/correlator/util/test.go

package util

import (
	"context"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"io/ioutil"
	"iwarma.ru/console/correlator/config"
	"iwarma.ru/console/correlator/es"
	"net/http"
	"testing"
	"time"
)

func PrepareElastic() {
	viper.Set(config.ElasticUrl, "http://elasticsearch:9200")
	viper.Set(config.ElasticRetryCount, 1)
	viper.Set(config.ElasticUsername, "elastic")
	viper.Set(config.ElasticPassword, "changeme")
	viper.Set(config.Verbose, true)
	viper.Set(config.ElasticAggregatedIndexName, "test-aggregated")
	viper.Set(config.ElasticNormalizedIndexName, "test-normalized")
	viper.Set(config.AggregatorIterationDuration, time.Second*2)
	viper.Set(config.Threads, 10)
	eventsFields := []string{
		"event_severity",
		"event_protocol",
		"message",
		"device_vendor",
		"device_product",
		"device_action",
		"device_version",
		"device_timezone",
		"sign_id",
		"sign_category",
		"sign_subcategory",
		"application",
		"source_ip",
		"source_host",
		"source_port",
		"source_mac",
		"source_timezone",
		"source_software",
		"source_action",
		"destination_ip",
		"destination_mac",
		"destination_timezone",
		"destination_software",
		"destination_action",
		"destination_host",
		"destination_port",
		"destination_user",
		"cs1",
		"cs1Label",
		"cs2",
		"cs2Label",
		"object_type",
	}
	viper.Set(config.AggregatedFields, eventsFields)
}

// TestUrl One server endpoint description
type TestUrl struct {
	Url             string            // Url to handle
	Response        []byte            // Response to send
	ResponseHeaders map[string]string // Response headers to send
	ResponseStatus  int               // Status code for response. Default is 200
	ReceivedMessage []byte            // Received message from client
	Error           error             // Error from handler function
	CallCount       int               // Counter of requests to this endpoint
}

// TestServer Test HTTP server
type TestServer struct {
	ServerAddress string         // Address for server
	Urls          []*TestUrl     // Url endpoints
	Server        *http.Server   // Http server object
	ServerMux     *http.ServeMux // Http multiplexer for server
}

// GetRequestAsString Get client request as string
func (url TestUrl) GetRequestAsString() string {
	return string(url.ReceivedMessage)
}

// Start test server
func (server *TestServer) Start() error {
	return server.Server.ListenAndServe()
}

// Stop test server
func (server *TestServer) Stop() error {
	return server.Server.Shutdown(context.Background())
}

// AddUrl Add url to server
func (server *TestServer) AddUrl(url string, response []byte, responseHeaders map[string]string) *TestUrl {
	result := new(TestUrl)
	result.Url = url
	result.Response = response
	result.ResponseHeaders = responseHeaders
	result.ResponseStatus = http.StatusOK

	handler := func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		result.ReceivedMessage, result.Error = ioutil.ReadAll(r.Body)
		result.CallCount = result.CallCount + 1
		for key, value := range result.ResponseHeaders {
			w.Header().Set(key, value)
		}
		w.Write(result.Response)
	}

	server.ServerMux.HandleFunc(result.Url, handler)
	server.Urls = append(server.Urls, result)

	return result
}

func (server *TestServer) AddUrlWithStatus(url string, status int) *TestUrl {
	result := new(TestUrl)
	result.Url = url
	result.ResponseStatus = status

	handler := func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		result.ReceivedMessage, result.Error = ioutil.ReadAll(r.Body)
		result.CallCount = result.CallCount + 1
		w.WriteHeader(result.ResponseStatus)
	}

	server.ServerMux.HandleFunc(result.Url, handler)
	server.Urls = append(server.Urls, result)

	return result
}

// NewTestServer Create new test server
func NewTestServer(address string) *TestServer {
	result := new(TestServer)
	result.ServerAddress = address
	result.Urls = make([]*TestUrl, 0)
	result.ServerMux = http.NewServeMux()
	result.Server = &http.Server{Addr: address, Handler: result.ServerMux}

	return result
}

func SetupTest(t *testing.T) {
	PrepareElastic()

}

func ClearIndex(el *es.Elastic, index ...string) error {
	for _, cur := range index {
		exists, err := el.CheckIndex(cur)
		if err != nil {
			log.Errorf("%+v", err)
			return err
		}

		if exists {
			err = el.DeleteIndex(cur)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func TearDownTest(t *testing.T) {
	client, err := es.NewClient()
	if err != nil {
		t.Error(err)
		return
	}
	del, err := client.DeleteIndex("*").Do(context.Background())
	if err != nil {
		t.Errorf("%v", err)
	}
	if !del.Acknowledged {
		t.Errorf("Got bad response. Indexs, wasn't deleted")
	}
}