<commit_before>package spinner

import (
    "net/http"
    "time"
    "sync"
    "fmt"
)

type TestWrapper struct {
    Spec *TestSpec
    Response *http.Response
    TimeElapsed float64
    Attempt int
    Err error
}

func getRequest(wrapper *TestWrapper) {
    start := time.Now()
    resp, err := http.Get(wrapper.Spec.Request.FullUrl())
    elapsed := time.Since(start)

    wrapper.Response = resp
    wrapper.Err = err
    wrapper.TimeElapsed = elapsed.Seconds()
}

func postRequest(wrapper *TestWrapper) {
}

func requestHandler(reqChan <-chan *TestWrapper, outChan chan<- *TestWrapper,
        waitGroup *sync.WaitGroup) {
    defer waitGroup.Done()

    for reqWrapper := range reqChan {
        for reqWrapper.Attempt < reqWrapper.Spec.Options.MaxAttempts {
            reqWrapper.Attempt += 1

            switch reqWrapper.Spec.Request.Method {
            case "GET":
                getRequest(reqWrapper)
            case "POST":
                postRequest(reqWrapper)
            default:
                panic("Invalid method encountered")
            }

            if reqWrapper.Err == nil {
                // Succeeded, so we don't need any more attempts
                break
            }

            // Wait some small but significant amount of time before hitting
            // the server again
            time.Sleep(100 * time.Millisecond)
        }

        outChan <- reqWrapper
    }
}

func outputHandler(outChan <-chan *TestWrapper, waitGroup *sync.WaitGroup) {
    defer waitGroup.Done()

    for respWrapper := range outChan {
        fmt.Printf("%v (%v)\n", respWrapper.Spec.Request.FullUrl(),
                respWrapper.Spec.Request.Method)

        if respWrapper.Err != nil {
            printStatus(FAILURE, "Connection")
        } else {
            response := respWrapper.Response
            responseSpec := respWrapper.Spec.Response

            printStatus(SUCCESS, "Connection")
            printStatus(responseSpec.checkStatusCode(response), "Status code")
            printStatus(responseSpec.checkHeaders(response), "Headers")
            printStatus(responseSpec.checkTimeElapsed(respWrapper.TimeElapsed), "Time elapsed")
            printStatus(responseSpec.checkAttempts(respWrapper.Attempt), "Attempts")
        }
    }
}

func ExecuteTestConfig(config *TestConfig) {
    var concurrentRequests int
    if config.Settings.ConcurrentRequests < 1 {
        concurrentRequests = 1
    } else {
        concurrentRequests = config.Settings.ConcurrentRequests
    }

    reqChan := make(chan *TestWrapper, 10)
    outChan := make(chan *TestWrapper)

    reqWaitGroup := new(sync.WaitGroup)
    reqWaitGroup.Add(concurrentRequests)

    for i := 0; i < concurrentRequests; i++ {
        go requestHandler(reqChan, outChan, reqWaitGroup)
    }

    outWaitGroup := new(sync.WaitGroup)
    outWaitGroup.Add(1)

    go outputHandler(outChan, outWaitGroup)

    for _, spec := range config.Specs {
        wrapper := new(TestWrapper)
        wrapper.Spec = spec
        reqChan <- wrapper
    }

    close(reqChan)
    reqWaitGroup.Wait()

    close(outChan)
    outWaitGroup.Wait()
}

// TODO Make this a method on the TestStatus so it can change the way it prints
func printStatus(status TestStatus, attribute string) {
    if status == UNKNOWN {
        return
    }
    if status == SUCCESS {
        println("  \u2713", attribute)
    }
    if status == WARNING {
        println("  \u2713", attribute)
    }
    if status == FAILURE {
        println("  \u2718", attribute)
    }
}
<commit_msg>Make requests more generic and enable body.<commit_after>package spinner

import (
    "bytes"
    "fmt"
    "io"
    "net/http"
    "net/url"
    "sync"
    "time"

    "github.com/wsxiaoys/terminal/color"
)

type TestWrapper struct {
    Spec *TestSpec
    Request *http.Request
    Response *http.Response
    TimeElapsed float64
    Attempt int
    Err error
}

// Create a custom ReadCloser because that's the only way we can get a string
// into the request body (because Go doesn't have type unions).
type BodyWrapper struct {
    io.Reader
}
func (BodyWrapper) Close() error { return nil }

func requestHandler(reqChan <-chan *TestWrapper, outChan chan<- *TestWrapper,
        waitGroup *sync.WaitGroup) {
    defer waitGroup.Done()

    for reqWrapper := range reqChan {
        client := http.Client{}

        for reqWrapper.Attempt < reqWrapper.Spec.Options.MaxAttempts {
            reqWrapper.Attempt += 1

            req := new(http.Request)
            req.Method = reqWrapper.Spec.Request.Method

            // TODO This conversion should be moved to the config read stage
            reqUrl, err := url.Parse(reqWrapper.Spec.Request.FullUrl())
            if err != nil {
                panic("Invalid URL")
            }
            req.URL = reqUrl

            req.Header = reqWrapper.Spec.Request.Headers
            req.Body = BodyWrapper{bytes.NewBufferString(reqWrapper.Spec.Request.Data)}

            start := time.Now()
            resp, err := client.Do(req)
            elapsed := time.Since(start)

            reqWrapper.Response = resp
            reqWrapper.Err = err
            reqWrapper.TimeElapsed = elapsed.Seconds()

            if reqWrapper.Err == nil {
                // Succeeded, so we don't need any more attempts
                break
            }

            // Wait some small but significant amount of time before hitting
            // the server again
            time.Sleep(100 * time.Millisecond)
        }

        outChan <- reqWrapper
    }
}

func outputHandler(outChan <-chan *TestWrapper, waitGroup *sync.WaitGroup) {
    defer waitGroup.Done()

    for respWrapper := range outChan {
        fmt.Printf("%v (%v)\n", respWrapper.Spec.Request.FullUrl(),
                respWrapper.Spec.Request.Method)

        if respWrapper.Err != nil {
            printStatus(FAILURE, "Connection")
        } else {
            response := respWrapper.Response
            responseSpec := respWrapper.Spec.Response

            printStatus(SUCCESS, "Connection")
            printStatus(responseSpec.checkStatusCode(response), "Status code")
            printStatus(responseSpec.checkHeaders(response), "Headers")
            printStatus(responseSpec.checkTimeElapsed(respWrapper.TimeElapsed), "Time elapsed")
            printStatus(responseSpec.checkAttempts(respWrapper.Attempt), "Attempts")
        }
    }
}

func ExecuteTestConfig(config *TestConfig) {
    var concurrentRequests int
    if config.Settings.ConcurrentRequests < 1 {
        concurrentRequests = 1
    } else {
        concurrentRequests = config.Settings.ConcurrentRequests
    }

    reqChan := make(chan *TestWrapper, 10)
    outChan := make(chan *TestWrapper)

    reqWaitGroup := new(sync.WaitGroup)
    reqWaitGroup.Add(concurrentRequests)

    for i := 0; i < concurrentRequests; i++ {
        go requestHandler(reqChan, outChan, reqWaitGroup)
    }

    outWaitGroup := new(sync.WaitGroup)
    outWaitGroup.Add(1)

    go outputHandler(outChan, outWaitGroup)

    for _, spec := range config.Specs {
        wrapper := new(TestWrapper)
        wrapper.Spec = spec
        reqChan <- wrapper
    }

    close(reqChan)
    reqWaitGroup.Wait()

    close(outChan)
    outWaitGroup.Wait()
}

// TODO Make this a method on the TestStatus so it can change the way it prints
func printStatus(status TestStatus, attribute string) {
    if status == UNKNOWN {
        return
    }
    if status == SUCCESS {
        color.Println("@g  \u2713", attribute)
    }
    if status == WARNING {
        color.Println("@y  \u2713", attribute)
    }
    if status == FAILURE {
        color.Println("@r  \u2718", attribute)
    }
}

{'fraction_non_alphanumeric': 0.09823284823284824, 'fraction_numerical': 0.006237006237006237, 'mean_word_length': 2.257300042319086, 'pattern_counts': {'":': 2, '<': 19, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3261493', 'n_tokens_mistral': 2252, 'n_tokens_neox': 2104, 'n_words': 529}