Corrected comments, consistency, and integration test organization

SpecterOps · Aug 26, 2024 · 3f884c2 · 3f884c2
2 parents 8274d24 + 048d9ad
commit 3f884c2
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 26 deletions.
diff --git a/cmd/api/src/database/saved_queries_test.go b/cmd/api/src/database/saved_queries_test.go
@@ -24,8 +24,6 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-
 	"github.com/gofrs/uuid"
 	"github.com/specterops/bloodhound/src/model"
 	"github.com/specterops/bloodhound/src/test/integration"
@@ -68,25 +66,3 @@ func TestSavedQueries_ListSavedQueries(t *testing.T) {
 		t.Fatalf("Expected 3 saved queries to be returned")
 	}
 }
-
-func TestSavedQueries_IsSavedQuerySharedToUser(t *testing.T) {
-	var (
-		testCtx = context.Background()
-		dbInst  = integration.SetupDB(t)
-	)
-
-	user1, err := dbInst.CreateUser(testCtx, model.User{
-		PrincipalName: userPrincipal,
-	})
-	require.NoError(t, err)
-
-	query, err := dbInst.CreateSavedQuery(testCtx, user1.ID, "Test Query", "TESTING", "Example")
-	require.NoError(t, err)
-
-	_, err = dbInst.CreateSavedQueryPermissionsToUsers(testCtx, query.ID, user1.ID)
-	require.NoError(t, err)
-
-	isShared, err := dbInst.IsSavedQuerySharedToUser(testCtx, query.ID, user1.ID)
-	require.NoError(t, err)
-	assert.True(t, isShared)
-}
diff --git a/cmd/api/src/database/savedqueriespermissions.go b/cmd/api/src/database/savedqueriespermissions.go
@@ -76,7 +76,7 @@ func (s *BloodhoundDB) CreateSavedQueryPermissionsToUsers(ctx context.Context, q
 	return newPermissions, CheckError(result)
 }
 
-// DeleteSavedQueryPermissionsForUsers batch deletes permissions associated a query id and a list of users
+// DeleteSavedQueryPermissionsForUsers batch deletes permissions associated with a query id and a list of users
 // If no user ids are supplied, all records for query id are deleted
 func (s *BloodhoundDB) DeleteSavedQueryPermissionsForUsers(ctx context.Context, queryID int64, userIDs ...uuid.UUID) error {
 	result := s.db.WithContext(ctx).Table("saved_queries_permissions").Where("query_id = ?", queryID)
@@ -125,7 +125,7 @@ func (s *BloodhoundDB) GetScopeForSavedQuery(ctx context.Context, queryID int64,
 // IsSavedQueryPublic returns true or false whether a provided saved query is public
 func (s *BloodhoundDB) IsSavedQueryPublic(ctx context.Context, queryID int64) (bool, error) {
 	rows := int64(0)
-	result := s.db.WithContext(ctx).Select("saved_queries_permissions.*").Where("public = true AND query_id=", queryID).Count(&rows)
+	result := s.db.WithContext(ctx).Table("saved_queries_permissions").Where("public = true AND query_id = ?", queryID).Count(&rows)
 
 	return rows > 0, CheckError(result)
 }

diff --git a/cmd/api/src/database/savedqueriespermissions_test.go b/cmd/api/src/database/savedqueriespermissions_test.go
@@ -27,6 +27,7 @@ import (
 	"github.com/specterops/bloodhound/src/database"
 	"github.com/specterops/bloodhound/src/model"
 	"github.com/specterops/bloodhound/src/test/integration"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -291,3 +292,47 @@ func TestSavedQueriesPermissions_DeleteSavedQueryPermissionsForUsers(t *testing.
 		model.SavedQueryScopeShared: false,
 	}, scope2)
 }
+
+func TestSavedQueriesPermissions_IsSavedQueryPublic(t *testing.T) {
+	var (
+		testCtx = context.Background()
+		dbInst  = integration.SetupDB(t)
+	)
+
+	user1, err := dbInst.CreateUser(testCtx, model.User{
+		PrincipalName: userPrincipal,
+	})
+	require.NoError(t, err)
+
+	query, err := dbInst.CreateSavedQuery(testCtx, user1.ID, "Test Query", "TESTING", "Example")
+	require.NoError(t, err)
+
+	_, err = dbInst.CreateSavedQueryPermissionToPublic(testCtx, query.ID)
+	require.NoError(t, err)
+
+	isPublic, err := dbInst.IsSavedQueryPublic(testCtx, query.ID)
+	require.NoError(t, err)
+	assert.True(t, isPublic)
+}
+
+func TestSavedQueriesPermissions_IsSavedQuerySharedToUser(t *testing.T) {
+	var (
+		testCtx = context.Background()
+		dbInst  = integration.SetupDB(t)
+	)
+
+	user1, err := dbInst.CreateUser(testCtx, model.User{
+		PrincipalName: userPrincipal,
+	})
+	require.NoError(t, err)
+
+	query, err := dbInst.CreateSavedQuery(testCtx, user1.ID, "Test Query", "TESTING", "Example")
+	require.NoError(t, err)
+
+	_, err = dbInst.CreateSavedQueryPermissionsToUsers(testCtx, query.ID, user1.ID)
+	require.NoError(t, err)
+
+	isShared, err := dbInst.IsSavedQuerySharedToUser(testCtx, query.ID, user1.ID)
+	require.NoError(t, err)
+	assert.True(t, isShared)
+}