From 07702be4999ea2c4459e5e82c1cb2ce7d6a76325 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_test.go | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/browser/registry_test.go b/browser/registry_test.go
index db21c3f17..dece29d16 100644
--- a/browser/registry_test.go
+++ b/browser/registry_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/xk6-browser/env"
+	"github.com/grafana/xk6-browser/k6ext/k6test"
 )
 
 func TestPidRegistry(t *testing.T) {
@@ -182,3 +183,39 @@ func TestIsRemoteBrowser(t *testing.T) {
 		require.Equal(t, "WS_URL_2", wsURL)
 	})
 }
+
+func TestBrowserRegistry(t *testing.T) {
+	t.Parallel()
+
+	vu := k6test.NewVU(t)
+	remoteRegistry, err := newRemoteRegistry(func(key string) (string, bool) {
+		// No env vars
+		return "", false
+	})
+	require.NoError(t, err)
+
+	// Initialize browser registry and start listening for events
+	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()
+}