DAOS-16829 control: fix for daos pool query

src/control/cmd/daos/health.go

@@ -104,7 +104,7 @@ func (cmd *healthCheckCmd) Execute([]string) error {
 		if pool.DisabledTargets > 0 {
 			queryMask.SetOptions(daos.PoolQueryOptionDisabledEngines)
 		}
-		tpi, err := queryPool(poolHdl, queryMask)
+		tpi, err := queryPool(nil, poolHdl, queryMask)
 		if err != nil {
 			cmd.Errorf("failed to query pool %s: %v", pool.Label, err)
 			continue

src/control/cmd/daos/pool.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"unsafe"
 
+	"github.com/golang/mock/gomock"
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
 
@@ -24,6 +25,7 @@ import (
 
 /*
 #include "util.h"
+#include <gurt/common.h>
 */
 import "C"
 
@@ -295,7 +297,7 @@ func convertPoolInfo(pinfo *C.daos_pool_info_t) (*daos.PoolInfo, error) {
 	return poolInfo, nil
 }
 
-func queryPoolRankLists(poolHdl C.daos_handle_t, queryMask daos.PoolQueryMask) (*daos.PoolInfo, error) {
+func queryPoolRankLists(mockDAOS *MockDAOSInterface, poolHdl C.daos_handle_t, queryMask daos.PoolQueryMask) (*daos.PoolInfo, error) {
 	var rlPtr **C.d_rank_list_t = nil
 	var rl *C.d_rank_list_t = nil
 
@@ -308,10 +310,18 @@ func queryPoolRankLists(poolHdl C.daos_handle_t, queryMask daos.PoolQueryMask) (
 		pi_bits: C.uint64_t(queryMask),
 	}
 
-	rc := C.daos_pool_query(poolHdl, rlPtr, &cPoolInfo, nil, nil)
-	defer C.d_rank_list_free(rl)
-	if err := daosError(rc); err != nil {
-		return nil, err
+	if mockDAOS != nil {
+		rc := mockDAOS.daos_pool_query(poolHdl, rlPtr, &cPoolInfo, nil, nil)
+		defer C.d_rank_list_free(rl)
+		if err := daosError(rc); err != nil {
+			return nil, err
+		}
+	} else {
+		rc := C.daos_pool_query(poolHdl, rlPtr, &cPoolInfo, nil, nil)
+		defer C.d_rank_list_free(rl)
+		if err := daosError(rc); err != nil {
+			return nil, err
+		}
 	}
 
 	poolInfo, err := convertPoolInfo(&cPoolInfo)
@@ -337,7 +347,7 @@ func queryPoolRankLists(poolHdl C.daos_handle_t, queryMask daos.PoolQueryMask) (
 
 	return poolInfo, nil
 }
-func queryPool(poolHdl C.daos_handle_t, queryMask daos.PoolQueryMask) (*daos.PoolInfo, error) {
+func queryPool(mockDAOS *MockDAOSInterface, poolHdl C.daos_handle_t, queryMask daos.PoolQueryMask) (*daos.PoolInfo, error) {
 	poolInfo := &daos.PoolInfo{}
 	originalMask := queryMask // Save the original queryMask
 
@@ -347,7 +357,7 @@ func queryPool(poolHdl C.daos_handle_t, queryMask daos.PoolQueryMask) (*daos.Poo
 		queryMask.ClearAll()
 		queryMask.SetOptions(option)
 
-		poolInfo1, err := queryPoolRankLists(poolHdl, queryMask)
+		poolInfo1, err := queryPoolRankLists(mockDAOS, poolHdl, queryMask)
 		if err != nil {
 			return err
 		}
@@ -364,35 +374,71 @@ func queryPool(poolHdl C.daos_handle_t, queryMask daos.PoolQueryMask) (*daos.Poo
 	}
 
 	// Preprocess queryMask, select one option for the first query
-	var firstOption string
-	if originalMask.HasOption(daos.PoolQueryOptionEnabledEngines) {
-		firstOption = daos.PoolQueryOptionEnabledEngines
-	} else if originalMask.HasOption(daos.PoolQueryOptionDisabledEngines) {
-		firstOption = daos.PoolQueryOptionDisabledEngines
-	} else if originalMask.HasOption(daos.PoolQueryOptionDeadEngines) {
-		firstOption = daos.PoolQueryOptionDeadEngines
+	queryOptions := []string{
+		daos.PoolQueryOptionEnabledEngines,
+		daos.PoolQueryOptionDisabledEngines,
+		daos.PoolQueryOptionDeadEngines,
+	}
+	var firstOption string = ""
+	for _, opt := range queryOptions {
+		if queryMask.HasOption(opt) && firstOption == "" {
+			firstOption = opt
+			continue
+		}
+		queryMask.ClearOptions(opt)
 	}
 
 	// Perform the first query to get basic information
-	if err := queryAndUpdate(firstOption); err != nil {
+	poolInfo, err := queryPoolRankLists(mockDAOS, poolHdl, queryMask)
+	if err != nil {
 		return nil, err
 	}
 
 	// Check the original query mask and update fields as needed
-	queryOptions := []string{
-		daos.PoolQueryOptionEnabledEngines,
-		daos.PoolQueryOptionDisabledEngines,
-		daos.PoolQueryOptionDeadEngines,
-	}
-
-	// Process each option sequentially
 	for _, opt := range queryOptions {
 		if originalMask.HasOption(opt) && opt != firstOption {
 			if err := queryAndUpdate(opt); err != nil {
 				return nil, err
 			}
 		}
 	}
+	poolInfo.QueryMask = originalMask
+
+	return poolInfo, nil
+}
+
+func MockQueryPool(mockDAOS *MockDAOSInterface, queryMask daos.PoolQueryMask) (*daos.PoolInfo, error) {
+	var count int = 0
+	if queryMask.HasOption(daos.PoolQueryOptionEnabledEngines) {
+		count++
+	}
+	if queryMask.HasOption(daos.PoolQueryOptionDisabledEngines) {
+		count++
+	}
+	if queryMask.HasOption(daos.PoolQueryOptionDeadEngines) {
+		count++
+	}
+	if count == 0 {
+		count = 1
+	}
+	mockDAOS.EXPECT().daos_pool_query(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(C.int(0)).
+		DoAndReturn(func(poolHdl C.daos_handle_t, rlPtr **C.d_rank_list_t, poolInfo *C.daos_pool_info_t, arg3 unsafe.Pointer, arg4 unsafe.Pointer) C.int {
+			if poolInfo.pi_bits&C.DPI_ENGINES_ENABLED != 0 {
+				*rlPtr = C.d_rank_list_alloc(1)
+			}
+			if poolInfo.pi_bits&C.DPI_ENGINES_DISABLED != 0 {
+				*rlPtr = C.d_rank_list_alloc(2)
+			}
+			if poolInfo.pi_bits&C.DPI_ENGINES_DEAD != 0 {
+				*rlPtr = C.d_rank_list_alloc(3)
+			}
+			return C.int(0)
+		}).Times(count)
+
+	poolInfo, err := queryPool(mockDAOS, C.daos_handle_t{}, queryMask)
+	if err != nil {
+		return nil, err
+	}
 
 	return poolInfo, nil
 }
@@ -413,7 +459,7 @@ func (cmd *poolQueryCmd) Execute(_ []string) error {
 	}
 	defer cleanup()
 
-	poolInfo, err := queryPool(cmd.cPoolHandle, queryMask)
+	poolInfo, err := queryPool(nil, cmd.cPoolHandle, queryMask)
 	if err != nil {
 		return errors.Wrapf(err, "failed to query pool %q", cmd.PoolID())
 	}
@@ -474,7 +520,7 @@ func (cmd *poolQueryTargetsCmd) Execute(_ []string) error {
 	}
 
 	if len(idxList) == 0 {
-		pi, err := queryPool(cmd.cPoolHandle, daos.HealthOnlyPoolQueryMask)
+		pi, err := queryPool(nil, cmd.cPoolHandle, daos.HealthOnlyPoolQueryMask)
 		if err != nil || (pi.TotalTargets == 0 || pi.TotalEngines == 0) {
 			if err != nil {
 				return errors.Wrap(err, "pool query failed")

src/control/cmd/daos/pool_test.go

@@ -0,0 +1,68 @@
+//
+// (C) Copyright 2021-2024 Intel Corporation.
+//
+// SPDX-License-Identifier: BSD-2-Clause-Patent
+//
+
+package main
+
+import (
+	"testing"
+
+	"github.com/daos-stack/daos/src/control/lib/daos"
+	"github.com/golang/mock/gomock"
+
+	"github.com/daos-stack/daos/src/control/common/test"
+)
+
+func TestDaos_PoolQuery(t *testing.T) {
+	for name, tc := range map[string]struct {
+		arg              daos.PoolQueryMask
+		expEnabledCount  int
+		expDisabledCount int
+		expSuspectCount  int
+		expErr           error
+	}{
+		"default Query mask": {
+			arg:              daos.DefaultPoolQueryMask,
+			expEnabledCount:  0,
+			expDisabledCount: 2,
+			expSuspectCount:  0,
+		},
+		"Health Query mask": {
+			arg:              daos.HealthOnlyPoolQueryMask,
+			expEnabledCount:  0,
+			expDisabledCount: 2,
+			expSuspectCount:  3,
+		},
+		"Query All mask": {
+			arg:              daos.PoolQueryMask(^uint64(0)),
+			expEnabledCount:  1,
+			expDisabledCount: 2,
+			expSuspectCount:  3,
+		},
+		"Query nothing": {
+			arg:              daos.PoolQueryMask(uint64(0)),
+			expEnabledCount:  0,
+			expDisabledCount: 0,
+			expSuspectCount:  0,
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+
+			mockDAOS := NewMockDAOSInterface(ctrl)
+			poolInfo, gotErr := MockQueryPool(mockDAOS, tc.arg)
+			test.CmpErr(t, tc.expErr, gotErr)
+			if tc.expErr != nil {
+				return
+			}
+
+			test.AssertEqual(t, tc.expEnabledCount, poolInfo.EnabledRanks.Count(), "unexpected enabled count")
+			test.AssertEqual(t, tc.expDisabledCount, poolInfo.DisabledRanks.Count(), "unexpected disabled count")
+			test.AssertEqual(t, tc.expSuspectCount, poolInfo.DeadRanks.Count(), "unexpected suspect count")
+		})
+	}
+
+}

src/control/go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/Jille/raft-grpc-transport v1.2.0
 	github.com/desertbit/grumble v1.1.3
 	github.com/dustin/go-humanize v1.0.0
+	github.com/golang/mock v1.6.0
 	github.com/google/go-cmp v0.6.0
 	github.com/google/uuid v1.6.0
 	github.com/hashicorp/go-hclog v1.2.2

src/control/go.sum

@@ -123,6 +123,8 @@ github.com/golang/mock v1.4.0/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt
 github.com/golang/mock v1.4.1/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt3cw=
 github.com/golang/mock v1.4.3/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt3cw=
 github.com/golang/mock v1.4.4/go.mod h1:l3mdAwkq5BuhzHwde/uurv3sEJeZMXNpwsxVWU71h+4=
+github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
+github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
@@ -516,6 +518,7 @@ golang.org/x/tools v0.0.0-20200618134242-20370b0cb4b2/go.mod h1:EkVYQZoAsY45+roY
 golang.org/x/tools v0.0.0-20200729194436-6467de6f59a7/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA=
 golang.org/x/tools v0.0.0-20200804011535-6c149bb5ef0d/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA=
 golang.org/x/tools v0.0.0-20200825202427-b303f430e36d/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA=
+golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=