From 764f815a5ed71667c31063151e395c81e6dc437d Mon Sep 17 00:00:00 2001
From: ka3de <danijs12@hotmail.com>
Date: Tue, 20 Jun 2023 16:05:51 +0200
Subject: [PATCH] Add browser registry test

---
 browser/registry.go      |  8 ++++
 browser/registry_test.go | 92 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/browser/registry.go b/browser/registry.go
index 789817f66..ebebbf743 100644
--- a/browser/registry.go
+++ b/browser/registry.go
@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/grafana/xk6-browser/api"
 	"github.com/grafana/xk6-browser/chromium"
@@ -164,6 +165,8 @@ type browserRegistry struct {
 	m  map[int64]api.Browser
 
 	buildFn browserBuildFunc
+
+	stopped atomic.Bool // testing purposes
 }
 
 type browserBuildFunc func(ctx context.Context) (api.Browser, error)
@@ -236,6 +239,7 @@ func (r *browserRegistry) handleIterEvents(eventsCh <-chan *k6event.Event, unsub
 		// unsubscribe for the VU events and exit the loop in order to reduce unuseful overhead.
 		if !isBrowserIter(r.vu) {
 			unsubscribeFn()
+			r.stop()
 			e.Done()
 			return
 		}
@@ -318,6 +322,10 @@ func (r *browserRegistry) clear() {
 	}
 }
 
+func (r *browserRegistry) stop() {
+	r.stopped.Store(true)
+}
+
 func isBrowserIter(vu k6modules.VU) bool {
 	opts := k6ext.GetScenarioOpts(vu.Context(), vu)
 	_, ok := opts["type"] // Check if browser type option is set
diff --git a/browser/registry_test.go b/browser/registry_test.go
index db21c3f17..41d008445 100644
--- a/browser/registry_test.go
+++ b/browser/registry_test.go
@@ -1,6 +1,7 @@
 package browser
 
 import (
+	"context"
 	"errors"
 	"strconv"
 	"sync"
@@ -10,6 +11,9 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/xk6-browser/env"
+	"github.com/grafana/xk6-browser/k6ext/k6test"
+
+	k6event "go.k6.io/k6/event"
 )
 
 func TestPidRegistry(t *testing.T) {
@@ -182,3 +186,91 @@ func TestIsRemoteBrowser(t *testing.T) {
 		require.Equal(t, "WS_URL_2", wsURL)
 	})
 }
+
+func TestBrowserRegistry(t *testing.T) {
+	t.Parallel()
+
+	remoteRegistry, err := newRemoteRegistry(func(key string) (string, bool) {
+		// No env vars
+		return "", false
+	})
+	require.NoError(t, err)
+
+	t.Run("init_and_close_browsers_on_iter_events", func(t *testing.T) {
+		t.Parallel()
+
+		vu := k6test.NewVU(t)
+		browserRegistry := newBrowserRegistry(vu, remoteRegistry, &pidRegistry{})
+
+		vu.ActivateVU()
+
+		// Send a few IterStart events
+		vu.StartIteration(t, k6test.WithIteration(0))
+		vu.StartIteration(t, k6test.WithIteration(1))
+		vu.StartIteration(t, k6test.WithIteration(2))
+
+		// Verify browsers are initialized
+		browserRegistry.mu.RLock()
+		assert.Equal(t, 3, len(browserRegistry.m))
+		browserRegistry.mu.RUnlock()
+
+		// Send IterEnd events
+		vu.EndIteration(t, k6test.WithIteration(0))
+		vu.EndIteration(t, k6test.WithIteration(1))
+		vu.EndIteration(t, k6test.WithIteration(2))
+
+		// Verify there are no browsers left
+		browserRegistry.mu.RLock()
+		assert.Equal(t, 0, len(browserRegistry.m))
+		browserRegistry.mu.RUnlock()
+	})
+
+	t.Run("close_browsers_on_exit_event", func(t *testing.T) {
+		t.Parallel()
+
+		vu := k6test.NewVU(t)
+		browserRegistry := newBrowserRegistry(vu, remoteRegistry, &pidRegistry{})
+
+		vu.ActivateVU()
+
+		// Send a few IterStart events
+		vu.StartIteration(t, k6test.WithIteration(0))
+		vu.StartIteration(t, k6test.WithIteration(1))
+		vu.StartIteration(t, k6test.WithIteration(2))
+
+		// Verify browsers are initialized
+		browserRegistry.mu.RLock()
+		assert.Equal(t, 3, len(browserRegistry.m))
+		browserRegistry.mu.RUnlock()
+
+		// Send Exit event
+		events, ok := vu.EventsField.Global.(*k6event.System)
+		require.True(t, ok, "want *k6event.System; got %T", events)
+		waitDone := events.Emit(&k6event.Event{
+			Type: k6event.Exit,
+		})
+		require.NoError(t, waitDone(context.Background()), "error waiting on Exit done")
+
+		// Verify there are no browsers left
+		browserRegistry.mu.RLock()
+		assert.Equal(t, 0, len(browserRegistry.m))
+		browserRegistry.mu.RUnlock()
+	})
+
+	t.Run("unsubscribe_on_non_browser_vu", func(t *testing.T) {
+		t.Parallel()
+
+		vu := k6test.NewVU(t)
+		browserRegistry := newBrowserRegistry(vu, remoteRegistry, &pidRegistry{})
+
+		vu.ActivateVU()
+
+		// Unset browser type option in scenario options in order to represent that VU is not
+		// a browser test VU
+		delete(vu.StateField.Options.Scenarios["default"].GetScenarioOptions().Browser, "type")
+
+		vu.StartIteration(t)
+
+		assert.True(t, browserRegistry.stopped.Load())
+	})
+}