From fef1ce579ed5a30c0cd6dcb4ddc5cf9809bc4131 Mon Sep 17 00:00:00 2001
From: Zach Wasserman <zach@fleetdm.com>
Date: Tue, 25 May 2021 15:30:01 -0700
Subject: [PATCH] Fix host count for list teams (#861)

- Add test
---
 server/datastore/datastore.go       |  1 +
 server/datastore/datastore_teams.go | 58 ++++++++++++++++++++++++++++-
 server/datastore/mysql/teams.go     |  6 +--
 3 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/server/datastore/datastore.go b/server/datastore/datastore.go
index 1ea8eeebdc..1ace266c96 100644
--- a/server/datastore/datastore.go
+++ b/server/datastore/datastore.go
@@ -98,6 +98,7 @@ var TestFunctions = [...]func(*testing.T, kolide.Datastore){
 	testCarveUpdateCarve,
 	testTeamGetSetDelete,
 	testTeamUsers,
+	testTeamListTeams,
 	testUserTeams,
 	testUserCreateWithTeams,
 	testSaveHostSoftware,
diff --git a/server/datastore/datastore_teams.go b/server/datastore/datastore_teams.go
index 8efeb4df58..0208d9d914 100644
--- a/server/datastore/datastore_teams.go
+++ b/server/datastore/datastore_teams.go
@@ -1,9 +1,12 @@
 package datastore
 
 import (
+	"sort"
 	"testing"
+	"time"
 
 	"github.com/fleetdm/fleet/server/kolide"
+	"github.com/fleetdm/fleet/server/test"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -98,5 +101,58 @@ func testTeamUsers(t *testing.T, ds kolide.Datastore) {
 	team2, err = ds.Team(team2.ID)
 	require.NoError(t, err)
 	assert.ElementsMatch(t, team2Users, team2.Users)
-
+}
+
+func testTeamListTeams(t *testing.T, ds kolide.Datastore) {
+	users := createTestUsers(t, ds)
+	user1 := kolide.User{Name: users[0].Name, Email: users[0].Email, ID: users[0].ID}
+	user2 := kolide.User{Name: users[1].Name, Email: users[1].Email, ID: users[1].ID}
+
+	team1, err := ds.NewTeam(&kolide.Team{Name: "team1"})
+	require.NoError(t, err)
+	team2, err := ds.NewTeam(&kolide.Team{Name: "team2"})
+	require.NoError(t, err)
+
+	teams, err := ds.ListTeams(kolide.ListOptions{})
+	require.NoError(t, err)
+	sort.Slice(teams, func(i, j int) bool { return teams[i].Name < teams[j].Name })
+
+	assert.Equal(t, "team1", teams[0].Name)
+	assert.Equal(t, 0, teams[0].HostCount)
+	assert.Equal(t, 0, teams[0].UserCount)
+
+	assert.Equal(t, "team2", teams[1].Name)
+	assert.Equal(t, 0, teams[1].HostCount)
+	assert.Equal(t, 0, teams[1].UserCount)
+
+	host1 := test.NewHost(t, ds, "1", "1", "1", "1", time.Now())
+	host2 := test.NewHost(t, ds, "2", "2", "2", "2", time.Now())
+	host3 := test.NewHost(t, ds, "3", "3", "3", "3", time.Now())
+	require.NoError(t, ds.AddHostsToTeam(&team1.ID, []uint{host1.ID}))
+	require.NoError(t, ds.AddHostsToTeam(&team2.ID, []uint{host2.ID, host3.ID}))
+
+	team1.Users = []kolide.TeamUser{
+		{User: user1, Role: "maintainer"},
+		{User: user2, Role: "observer"},
+	}
+	team1, err = ds.SaveTeam(team1)
+	require.NoError(t, err)
+
+	team2.Users = []kolide.TeamUser{
+		{User: user1, Role: "maintainer"},
+	}
+	team1, err = ds.SaveTeam(team2)
+	require.NoError(t, err)
+
+	teams, err = ds.ListTeams(kolide.ListOptions{})
+	require.NoError(t, err)
+	sort.Slice(teams, func(i, j int) bool { return teams[i].Name < teams[j].Name })
+
+	assert.Equal(t, "team1", teams[0].Name)
+	assert.Equal(t, 1, teams[0].HostCount)
+	assert.Equal(t, 2, teams[0].UserCount)
+
+	assert.Equal(t, "team2", teams[1].Name)
+	assert.Equal(t, 2, teams[1].HostCount)
+	assert.Equal(t, 1, teams[1].UserCount)
 }
diff --git a/server/datastore/mysql/teams.go b/server/datastore/mysql/teams.go
index dd5134f8b7..cefb5210fe 100644
--- a/server/datastore/mysql/teams.go
+++ b/server/datastore/mysql/teams.go
@@ -151,9 +151,9 @@ func (d *Datastore) SaveTeam(team *kolide.Team) (*kolide.Team, error) {
 func (d *Datastore) ListTeams(opt kolide.ListOptions) ([]*kolide.Team, error) {
 	query := `
 		SELECT *,
-			(SELECT count(*) FROM user_teams WHERE team_id = id) AS user_count,
-			(SELECT count(*) FROM hosts WHERE team_id = id) AS host_count
-		FROM teams
+			(SELECT count(*) FROM user_teams WHERE team_id = t.id) AS user_count,
+			(SELECT count(*) FROM hosts WHERE team_id = t.id) AS host_count
+		FROM teams t
 		WHERE TRUE
 	`
 	query, params := searchLike(query, nil, opt.MatchQuery, teamSearchColumns...)