go/src/infra/fleetconsole/internal/database/browser_devicedb/count_devices_test.go

// Copyright 2026 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package browserdevicedb

import (
	"context"
	"database/sql/driver"
	"regexp"
	"testing"

	"github.com/DATA-DOG/go-sqlmock"

	"go.chromium.org/luci/common/testing/truth/assert"
	"go.chromium.org/luci/common/testing/truth/should"
	"go.chromium.org/luci/server/sqldb"

	"go.chromium.org/infra/fleetconsole/api/fleetconsolerpc"
)

func TestCountDevicesSQL(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name               string
		filter             string
		realms             []string
		expectedSQL        string
		expectedParameters []driver.Value
	}{
		{
			name:   "no filter, no realms",
			filter: "",
			realms: nil,
			expectedSQL: `SELECT COUNT(*) AS "total", ` +
				`COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $1) AS "alive", ` +
				`COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $2) AS "dead", ` +
				`COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $3) AS "quarantined", ` +
				`COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $4) AS "maintenance" ` +
				`FROM "browser_devices"`,
			expectedParameters: []driver.Value{`"alive"`, `"dead"`, `"quarantined"`, `"maintenance"`},
		},
		{
			name:   "with filter, with realms",
			filter: "id = 'device1'",
			realms: []string{"realm1"},
			expectedSQL: `SELECT COUNT(*) AS "total", ` +
				`COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $1) AS "alive", ` +
				`COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $2) AS "dead", ` +
				`COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $3) AS "quarantined", ` +
				`COUNT(*) FILTER (WHERE ("swarming_labels"->'state')::jsonb @> $4) AS "maintenance" ` +
				`FROM "browser_devices" ` +
				`WHERE "id" = $5 AND ("browser_devices"."realm" IS NULL OR "browser_devices"."realm" IN ($6))`,
			expectedParameters: []driver.Value{`"alive"`, `"dead"`, `"quarantined"`, `"maintenance"`, "'device1'", "realm1"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			db, mock, err := sqlmock.New()
			if err != nil {
				t.Fatalf("an error '%q' was not expected when opening a stub database connection", err)
			}
			defer db.Close()
			ctx = sqldb.UseDB(ctx, db)

			mock.ExpectQuery(regexp.QuoteMeta(tt.expectedSQL)).
				WithArgs(tt.expectedParameters...).
				WillReturnRows(sqlmock.NewRows([]string{"total", "alive", "dead", "quarantined", "maintenance"}).
					AddRow(10, 5, 2, 1, 2))

			res, err := CountDevices(ctx, tt.filter, tt.realms)
			assert.NoErr(t, err)
			assert.Loosely(t, res, should.Match(&fleetconsolerpc.CountBrowserDevicesResponse{
				Total: 10,
				SwarmingState: &fleetconsolerpc.SwarmingStateCounts{
					Alive:       5,
					Dead:        2,
					Quarantined: 1,
					Maintenance: 2,
				},
			}))

			if err := mock.ExpectationsWereMet(); err != nil {
				t.Errorf("there were unfulfilled expectations: %s", err)
			}
		})
	}
}