broad-announce/testfakes/tailcount/main.go

// Command tailcount subscribes to the M5 live tail
// (GET /v1/tail/ws?token=…&company_id=…) and counts how many
// events arrive within a timeout. Optionally filters to a
// specific dedupe_key. Used by the M6.5 smoke (Step 3) to
// prove the tail sees the storm even though the recipient
// only gets 1 message.
//
// Usage:
//
//	tailcount -url ws://localhost:8800/v1/tail/ws \
//	    -token tail-dev-token-please-change-in-prod \
//	    -company acme-001 \
//	    -filter "m65-step3-1234" \
//	    -timeout 15s
//
// On exit (timeout or signal), prints the count to stdout
// in the form `count=<N>` so the smoke script can grep.
package main

import (
	"flag"
	"fmt"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

func main() {
	urlFlag := flag.String("url", "ws://localhost:8800/v1/tail/ws", "tail endpoint")
	token := flag.String("token", "", "tail token (BA_INGESTD_TAIL_TOKEN)")
	company := flag.String("company", "", "company_id filter (required)")
	filter := flag.String("filter", "", "sub-string filter: only count events containing this string")
	timeout := flag.Duration("timeout", 15*time.Second, "max wall time to listen")
	flag.Parse()

	if *token == "" || *company == "" {
		log.Fatal("token and company are required")
	}

	u, err := url.Parse(*urlFlag)
	if err != nil {
		log.Fatalf("parse url: %v", err)
	}
	q := u.Query()
	q.Set("token", *token)
	q.Set("company_id", *company)
	u.RawQuery = q.Encode()

	dialer := *websocket.DefaultDialer
	dialer.HandshakeTimeout = 5 * time.Second
	conn, resp, err := dialer.Dial(u.String(), nil)
	if err != nil {
		st := "<nil>"
		if resp != nil {
			st = resp.Status
		}
		log.Fatalf("dial: %v (status=%s)", err, st)
	}
	defer conn.Close()

	count := 0
	deadline := time.Now().Add(*timeout)
	conn.SetReadDeadline(time.Now().Add(*timeout + 1*time.Second))
	for time.Now().Before(deadline) {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			break
		}
		s := string(msg)
		if *filter == "" || strings.Contains(s, *filter) {
			count++
		}
	}
	fmt.Printf("count=%d\n", count)
}