<commit_before>package main

import (
	"bytes"
	crand "crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"os/exec"
	"testing"
	"time"

	"code.google.com/p/go.net/websocket"
)

// sock.js protocol is described here:
// http://sockjs.github.io/sockjs-protocol/sockjs-protocol-0.3.3.html#section-36
const url = "ws://localhost:8008/subscribe/%d/%s/websocket"
const origin = "http://localhost/" // not checked on broker

// returna a new sockjs url for client
func newURL() string {
	return fmt.Sprintf(url, rand.Intn(1000), RandomStringLength(8))
}

func init() {
	rand.Seed(time.Now().UnixNano())
}

func runBroker(t *testing.T) (b *Broker, closer func()) {
	// Run authWorker (Authworker must be running when broker is running.)
	cmd := exec.Command("cake", "authWorker")
	cmd.Dir = "/opt/koding"
	err := cmd.Start()
	if err != nil {
		panic(err)
	}
	t.Log("authWorker is running")

	// Run broker
	broker := NewBroker()
	broker.Start()
	t.Log("broker is running")

	return broker, func() {
		// Close authWorker
		if err := cmd.Process.Kill(); err != nil {
			t.Errorf(err.Error())
		}
		broker.Close()
	}
}

func TestBroker(t *testing.T) {
	_, closer := runBroker(t)
	defer closer()

	client, err := dialSockJS(newURL(), origin)
	if err != nil {
		t.Errorf(err.Error())
		return
	}

	go client.Run()
	defer client.Close()

	type testCase struct{ send, expect string }
	cases := []testCase{
		testCase{`{"action": "ping"}`, `{"routingKey":"broker.pong","payload":null}`},
	}

	for _, tc := range cases {
		err = client.SendAndExpectString(tc.send, tc.expect)
		if err != nil {
			t.Errorf(err.Error())
			return
		}
	}
}

func TestPubSub(t *testing.T) {
	// Run authWorker and broker
	_, closer := runBroker(t)
	defer closer()

	// Run subscriber
	subscriber, err := dialSockJS(newURL(), origin)
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	go subscriber.Run()
	defer subscriber.Close()
	msg, err := subscriber.ReadJSON()
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	if msg["routingKey"].(string) != "broker.connected" {
		t.Errorf(err.Error())
		return
	}
	t.Log("subscriber is running")

	// Run publisher
	publisher, err := dialSockJS(newURL(), origin)
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	go publisher.Run()
	defer publisher.Close()
	msg, err = publisher.ReadJSON()
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	if msg["routingKey"].(string) != "broker.connected" {
		t.Errorf(err.Error())
		return
	}
	t.Log("publisher is running")

	// Subscribe
	err = subscriber.SendString(`{"action": "subscribe", "routingKeyPrefix": "client.foo"}`)
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	str, err := subscriber.ReadString()
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	if str != `{"routingKey":"broker.subscribed","payload":"client.foo"}` {
		t.Errorf("unexpected msg: %s", str)
		return
	}
	t.Log("subscribed")

	// Publish a message
	err = publisher.SendString(`{"action": "publish", "exchange": "broker", "routingKey": "client.foo", "payload": "{\"bar\": \"baz\"}"}`)
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	t.Log("published a message")

	// Receive published message
	str, err = subscriber.ReadString()
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	if str != `{"routingKey":"client.foo","payload":{"bar":"baz"}}` {
		t.Errorf("unexpected msg: %s", str)
		return
	}
}

// cheap imitation of sockjs-client js library
type sockJSClient struct {
	ws       *websocket.Conn
	messages chan []byte
}

func dialSockJS(url, origin string) (*sockJSClient, error) {
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		return nil, err
	}
	return newSockJSClient(ws), nil
}

func newSockJSClient(ws *websocket.Conn) *sockJSClient {
	return &sockJSClient{
		ws:       ws,
		messages: make(chan []byte),
	}
}

// read messages from websocket and put it to the channel
func (c *sockJSClient) Run() error {
	defer close(c.messages)
	for {
		var data []byte
		err := websocket.Message.Receive(c.ws, &data)
		if err != nil {
			return err
		}
		// fmt.Printf("--- read data: %+v\n", string(data))
		c.didMessage(data)
	}
}

func (c *sockJSClient) Close() {
	c.ws.Close()
}

// Send a []byte message to server
func (c *sockJSClient) Send(data []byte) error {
	return websocket.Message.Send(c.ws, data)
}

// Send a string message to server
func (c *sockJSClient) SendString(s string) error {
	return websocket.Message.Send(c.ws, s)
}

// adapted from: https://github.com/sockjs/sockjs-client/blob/master/lib/sockjs.js#L146<commit_msg>broker: add partly functioning benchmarks, will improve later<commit_after>package main

///////////////////////////////////////////////////////////////
// Make sure authWorker is running before running the tests. //
// You can run it with the following command:                //
//   cd /opt/koding && cake -c vagrant authWorker            //
///////////////////////////////////////////////////////////////

import (
	"bytes"
	crand "crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"testing"
	"time"

	"code.google.com/p/go.net/websocket"
)

// sock.js protocol is described here:
// http://sockjs.github.io/sockjs-protocol/sockjs-protocol-0.3.3.html#section-36
const url = "ws://localhost:8008/subscribe/%d/%s/websocket"
const origin = "http://localhost/" // not checked on broker

// returna a new sockjs url for client
func newURL() string {
	return fmt.Sprintf(url, rand.Intn(1000), RandomStringLength(8))
}

// This global instance of broker is run once when the tests are run by init() function.
var broker *Broker

func init() {
	rand.Seed(time.Now().UnixNano())
	broker := NewBroker()
	broker.Start()
}

func TestPingPong(t *testing.T) {
	client, err := dialSockJS(newURL(), origin)
	if err != nil {
		t.Errorf(err.Error())
		return
	}

	go client.Run()
	defer client.Close()

	type testCase struct{ send, expect string }
	cases := []testCase{
		testCase{`{"action": "ping"}`, `{"routingKey":"broker.pong","payload":null}`},
	}

	for _, tc := range cases {
		err = client.SendAndExpectString(tc.send, tc.expect)
		if err != nil {
			t.Errorf(err.Error())
			return
		}
	}
}

func TestPubSub(t *testing.T) {
	// Run subscriber
	subscriber, err := dialSockJS(newURL(), origin)
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	go subscriber.Run()
	defer subscriber.Close()
	msg, err := subscriber.ReadJSON()
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	if msg["routingKey"].(string) != "broker.connected" {
		t.Errorf(err.Error())
		return
	}
	t.Log("subscriber is running")

	// Run publisher
	publisher, err := dialSockJS(newURL(), origin)
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	go publisher.Run()
	defer publisher.Close()
	msg, err = publisher.ReadJSON()
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	if msg["routingKey"].(string) != "broker.connected" {
		t.Errorf(err.Error())
		return
	}
	t.Log("publisher is running")

	// Subscribe
	err = subscriber.SendString(`{"action": "subscribe", "routingKeyPrefix": "client.foo"}`)
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	str, err := subscriber.ReadString()
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	if str != `{"routingKey":"broker.subscribed","payload":"client.foo"}` {
		t.Errorf("unexpected msg: %s", str)
		return
	}
	t.Log("subscribed")

	// Publish a message
	err = publisher.SendString(`{"action": "publish", "exchange": "broker", "routingKey": "client.foo", "payload": "{\"bar\": \"baz\"}"}`)
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	t.Log("published a message")

	// Receive published message
	str, err = subscriber.ReadString()
	if err != nil {
		t.Errorf(err.Error())
		return
	}
	if str != `{"routingKey":"client.foo","payload":{"bar":"baz"}}` {
		t.Errorf("unexpected msg: %s", str)
		return
	}
}

func BenchmarkBroker_1_1(b *testing.B)       { benchmarkBroker(b, 1, 1) }
func BenchmarkBroker_10_10(b *testing.B)     { benchmarkBroker(b, 10, 10) }
func BenchmarkBroker_100_100(b *testing.B)   { benchmarkBroker(b, 100, 100) }
func BenchmarkBroker_1000_1000(b *testing.B) { benchmarkBroker(b, 1000, 1000) }

var nPublished int

func benchmarkBroker(b *testing.B, nClient, nKey int) {
	var err error

	b.Logf("connecting with %d clients", nClient)
	clients := make([]*sockJSClient, nClient)
	for i := 0; i < nClient; i++ {
		clients[i], err = dialSockJS(newURL(), origin)
		if err != nil {
			b.Errorf(err.Error())
			return
		}
		go clients[i].Run()
		defer clients[i].Close()
	}

	b.Logf("generating %d keys", nKey)
	keys := make([]string, nKey)
	for i := 0; i < nKey; i++ {
		keys[i] = "client." + RandomStringLength(8)
	}

	b.Logf("each client subscribes %d keys", nKey)
	for _, client := range clients {
		for _, key := range keys {
			client.SendString(fmt.Sprintf(`{"action": "subscribe", "routingKeyPrefix": "%s"}`, key))
		}
	}

	b.Logf("publishing %d random messages to random keys", b.N)
	// conn := amqputil.CreateConnection("broker")
	// defer conn.Close()
	// ch := amqputil.CreateChannel(conn)
	// defer ch.Close()
	// payload := fmt.Sprintf(`{"random": "%s"}`, RandomStringLength(1024)) // Must be JSON
	body := fmt.Sprintf(`{"action": "publish", "exchange": "broker", "routingKey": "%s", "payload": "{\"random\": \"%s\"}"}`, keys[rand.Intn(nKey)], RandomStringLength(1024))

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		// err := ch.Publish("broker", keys[rand.Intn(nKey)], false, false, amqp.Publishing{Body: []byte(payload)})
		err := clients[rand.Intn(nClient)].SendString(body)
		if err != nil {
			b.Errorf(err.Error())
			return
		}
		nPublished++
	}
	fmt.Println("--- total published:", nPublished)
}

// cheap imitation of sockjs-client js library
type sockJSClient struct {
	ws       *websocket.Conn
	messages chan []byte
}

func dialSockJS(url, origin string) (*sockJSClient, error) {
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		return nil, err
	}
	return newSockJSClient(ws), nil
}

func newSockJSClient(ws *websocket.Conn) *sockJSClient {
	return &sockJSClient{
		ws:       ws,
		messages: make(chan []byte),
	}
}

// read messages from websocket and put it to the channel
func (c *sockJSClient) Run() error {
	defer close(c.messages)
	for {
		var data []byte
		err := websocket.Message.Receive(c.ws, &data)
		if err != nil {
			return err
		}
		c.didMessage(data)
	}
}

func (c *sockJSClient) Close() {
	c.ws.Close()
}

// Send a []byte message to server
func (c *sockJSClient) Send(data []byte) error {
	return websocket.Message.Send(c.ws, data)
}

// Send a string message to server
func (c *sockJSClient) SendString(s string) error {
	return websocket.Message.Send(c.ws, s)
}

// adapted from: https://github.com/sockjs/sockjs-client/blob/master/lib/sockjs.js#L146
{'fraction_non_alphanumeric': 0.1850540439806187, 'fraction_numerical': 0.008386134923592994, 'mean_word_length': 3.569178373776075, 'pattern_counts': {'":': 38, '<': 6, '<?xml version=': 0, '>': 3, 'https://': 2, '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': '4422058', 'n_tokens_mistral': 4120, 'n_tokens_neox': 3863, 'n_words': 941}