websocket/examples: enhance the server and client

In the server, change the listen port to 9101 to prevent conflict
with autobahn test suite.

In the client, add command "chat" and "chatbot".
The "chat" command run client chart as before.
The "chatbot" command create client for all known accounts and send
N messages continuously.

1 files changed, 99 insertions, 8 deletions

diff --git a/lib/websocket/examples/cmd/client/main.go b/lib/websocket/examples/cmd/client/main.go
index 53d43416..013b08b1 100644
--- a/lib/websocket/examples/cmd/client/main.go
+++ b/lib/websocket/examples/cmd/client/main.go
@@ -7,12 +7,12 @@
 //
 // To run the client as user ID 1 (Groot),
 //
-//	$ go run . 1
+//	$ go run . chat 1
 //
 // You can open other terminal and run another clients,
 //
-//	$ go run . 2 # or
-//	$ go run . 3
+//	$ go run . chat 2 # or
+//	$ go run . chat 3
 //
 // and start chatting with each others.
 package main
@@ -20,18 +20,25 @@ package main
 import (
 	"bufio"
 	"encoding/json"
+	"flag"
 	"fmt"
 	"log"
 	"net/http"
 	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/shuLhan/share/lib/websocket"
 	"github.com/shuLhan/share/lib/websocket/examples"
 )
 
+const (
+	cmdChat    = `chat`
+	cmdChatbot = `chatbot`
+)
+
 type ChatClient struct {
 	user *examples.Account
 	conn *websocket.Client
@@ -42,7 +49,7 @@ func NewChatClient(user *examples.Account) (cc *ChatClient) {
 	cc = &ChatClient{
 		user: user,
 		conn: &websocket.Client{
-			Endpoint: "ws://127.0.0.1:9001",
+			Endpoint: `ws://127.0.0.1:9101`,
 			Headers: http.Header{
 				"Key": []string{user.Key},
 			},
@@ -126,14 +133,41 @@ func (cc *ChatClient) handleText(cl *websocket.Client, frame *websocket.Frame) (
 	return nil
 }
 
+func usage() {
+	fmt.Println(`= WebSocket client example
+
+	client <chat | chatbot> <args...>
+
+== USAGE
+
+client ` + cmdChat + ` <id>
+	Connect to chat with others using specific ID: 1, 2, or 3.
+
+client ` + cmdChatbot + ` <N>
+	Connect to the server and sent N messages for each user
+	simultaneously.`)
+}
+
 func main() {
-	log.SetFlags(0)
+	flag.Parse()
 
-	if len(os.Args) <= 1 {
-		log.Printf("client <id>")
+	if len(os.Args) <= 2 {
+		usage()
 		return
 	}
 
+	var cmd = strings.ToLower(flag.Arg(0))
+	switch cmd {
+	case cmdChat:
+		doChat(flag.Arg(1))
+	case cmdChatbot:
+		doChatbot(flag.Arg(1))
+	default:
+		log.Fatalf(`unknown command: %s`, cmd)
+	}
+}
+
+func doChat(userIDStr string) {
 	var (
 		user *examples.Account
 		cc   *ChatClient
@@ -142,7 +176,7 @@ func main() {
 		ok   bool
 	)
 
-	uid, err = strconv.Atoi(os.Args[1])
+	uid, err = strconv.Atoi(userIDStr)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -156,3 +190,60 @@ func main() {
 
 	cc.Start()
 }
+
+func doChatbot(nStr string) {
+	var (
+		wg   sync.WaitGroup
+		user *examples.Account
+		err  error
+		n    int64
+	)
+
+	n, err = strconv.ParseInt(nStr, 10, 64)
+	if err != nil {
+		log.Fatalf(`invalid N: %s`, err)
+	}
+
+	for _, user = range examples.Users {
+		wg.Add(1)
+		go runChatbot(&wg, user, n)
+	}
+	wg.Wait()
+}
+
+func runChatbot(wg *sync.WaitGroup, user *examples.Account, n int64) {
+	var (
+		req = &websocket.Request{
+			Method: http.MethodPost,
+			Target: `/message`,
+		}
+
+		err    error
+		packet []byte
+		x      int64
+	)
+
+	var cc = NewChatClient(user)
+
+	for ; x < n; x++ {
+		req.ID = uint64(time.Now().UnixNano())
+		req.Body = fmt.Sprintf(`#%d Hello from %s at %d`, x, user.Name, req.ID)
+
+		packet, err = json.Marshal(req)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		err = cc.conn.SendText(packet)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		time.Sleep(100 * time.Millisecond)
+	}
+	err = cc.conn.Close()
+	if err != nil {
+		log.Fatal(err)
+	}
+	wg.Done()
+}