github.com/kubeshop/testkube@v1.17.23/pkg/event/kind/websocket/loader.go (about)

     1  package websocket
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/gofiber/websocket/v2"
     7  	"github.com/google/uuid"
     8  
     9  	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
    10  	"github.com/kubeshop/testkube/pkg/event/kind/common"
    11  )
    12  
    13  func NewWebsocketLoader() *WebsocketLoader {
    14  	return &WebsocketLoader{
    15  		Listener: NewWebsocketListener(),
    16  	}
    17  }
    18  
    19  type WebsocketLoader struct {
    20  	mutex    sync.Mutex
    21  	Listener *WebsocketListener
    22  }
    23  
    24  func (l *WebsocketLoader) Kind() string {
    25  	return "websocket"
    26  }
    27  
    28  func (l *WebsocketLoader) Load() (listeners common.Listeners, err error) {
    29  	return common.Listeners{l.Listener}, nil
    30  }
    31  
    32  func (l *WebsocketLoader) Add(conn *websocket.Conn) chan bool {
    33  	var end chan bool
    34  	id := uuid.NewString()
    35  
    36  	conn.SetCloseHandler(func(code int, text string) error {
    37  		for i, websocket := range l.Listener.Websockets {
    38  			if websocket.Id == id {
    39  				l.mutex.Lock()
    40  				l.Listener.Websockets = append(l.Listener.Websockets[:i], l.Listener.Websockets[i+1:]...)
    41  				l.mutex.Unlock()
    42  				break
    43  			}
    44  		}
    45  
    46  		end <- true
    47  		return nil
    48  	})
    49  
    50  	l.mutex.Lock()
    51  	defer l.mutex.Unlock()
    52  	l.Listener.Websockets = append(l.Listener.Websockets, Websocket{Id: id, Conn: conn, Events: testkube.AllEventTypes})
    53  
    54  	conn.WriteJSON(map[string]string{"message": "connected to Testkube Events", "id": id})
    55  	return end
    56  }