Fix ZFS storage reporting on TrueNAS SCALE (#718)

- Refactor collector to support mocking
- Fix ZFS detection to support 'fuse.zfs' and case-insensitivity
- Add regression tests for ZFS dataset deduplication

go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/rs/zerolog v1.34.0
 	github.com/shirou/gopsutil/v4 v4.25.9
 	github.com/spf13/cobra v1.9.1
+	github.com/stretchr/testify v1.11.1
 	golang.org/x/crypto v0.45.0
 	golang.org/x/oauth2 v0.31.0
 	golang.org/x/sys v0.38.0
@@ -33,6 +34,7 @@ require (
 	github.com/containerd/errdefs v1.0.0 // indirect
 	github.com/containerd/errdefs/pkg v0.3.0 // indirect
 	github.com/containerd/log v0.1.0 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/docker/go-connections v0.6.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
@@ -55,6 +57,7 @@ require (
 	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/opencontainers/image-spec v1.1.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect
 	github.com/prometheus/client_model v0.6.2 // indirect
 	github.com/prometheus/common v0.66.1 // indirect

internal/hostmetrics/collector.go

@@ -16,6 +16,18 @@ import (
 	gonet "github.com/shirou/gopsutil/v4/net"
 )
 
+// System call wrappers for testing
+var (
+	cpuCounts      = gocpu.CountsWithContext
+	cpuPercent     = gocpu.PercentWithContext
+	loadAvg        = goload.AvgWithContext
+	virtualMemory  = gomem.VirtualMemoryWithContext
+	diskPartitions = godisk.PartitionsWithContext
+	diskUsage      = godisk.UsageWithContext
+	netInterfaces  = gonet.InterfacesWithContext
+	netIOCounters  = gonet.IOCountersWithContext
+)
+
 // Snapshot represents a host resource utilisation sample.
 type Snapshot struct {
 	CPUUsagePercent float64
@@ -33,7 +45,7 @@ func Collect(ctx context.Context) (Snapshot, error) {
 
 	var snapshot Snapshot
 
-	if cpuCount, err := gocpu.CountsWithContext(collectCtx, true); err == nil {
+	if cpuCount, err := cpuCounts(collectCtx, true); err == nil {
 		snapshot.CPUCount = cpuCount
 	}
 
@@ -41,11 +53,11 @@ func Collect(ctx context.Context) (Snapshot, error) {
 		snapshot.CPUUsagePercent = cpuUsage
 	}
 
-	if loadAvg, err := goload.AvgWithContext(collectCtx); err == nil && loadAvg != nil {
+	if loadAvg, err := loadAvg(collectCtx); err == nil && loadAvg != nil {
 		snapshot.LoadAverage = []float64{loadAvg.Load1, loadAvg.Load5, loadAvg.Load15}
 	}
 
-	memStats, err := gomem.VirtualMemoryWithContext(collectCtx)
+	memStats, err := virtualMemory(collectCtx)
 	if err != nil {
 		return Snapshot{}, fmt.Errorf("memory stats: %w", err)
 	}
@@ -71,7 +83,7 @@ func Collect(ctx context.Context) (Snapshot, error) {
 }
 
 func collectCPUUsage(ctx context.Context) (float64, error) {
-	percentages, err := gocpu.PercentWithContext(ctx, time.Second, false)
+	percentages, err := cpuPercent(ctx, time.Second, false)
 	if err != nil {
 		return 0, err
 	}
@@ -90,7 +102,7 @@ func collectCPUUsage(ctx context.Context) (float64, error) {
 }
 
 func collectDisks(ctx context.Context) []agentshost.Disk {
-	partitions, err := godisk.PartitionsWithContext(ctx, true)
+	partitions, err := diskPartitions(ctx, true)
 	if err != nil {
 		return nil
 	}
@@ -108,7 +120,7 @@ func collectDisks(ctx context.Context) []agentshost.Disk {
 		}
 		seen[part.Mountpoint] = struct{}{}
 
-		usage, err := godisk.UsageWithContext(ctx, part.Mountpoint)
+		usage, err := diskUsage(ctx, part.Mountpoint)
 		if err != nil {
 			continue
 		}
@@ -116,7 +128,7 @@ func collectDisks(ctx context.Context) []agentshost.Disk {
 			continue
 		}
 
-		if part.Fstype == "zfs" {
+		if strings.EqualFold(part.Fstype, "zfs") || strings.EqualFold(part.Fstype, "fuse.zfs") {
 			pool := zfsPoolFromDevice(part.Device)
 			if pool == "" {
 				continue
@@ -162,12 +174,12 @@ func collectDisks(ctx context.Context) []agentshost.Disk {
 }
 
 func collectNetwork(ctx context.Context) []agentshost.NetworkInterface {
-	ifaces, err := gonet.InterfacesWithContext(ctx)
+	ifaces, err := netInterfaces(ctx)
 	if err != nil {
 		return nil
 	}
 
-	ioCounters, err := gonet.IOCountersWithContext(ctx, true)
+	ioCounters, err := netIOCounters(ctx, true)
 	if err != nil {
 		ioCounters = nil
 	}

internal/hostmetrics/collector_test.go

@@ -0,0 +1,133 @@
+package hostmetrics
+
+import (
+	"context"
+	"testing"
+
+	godisk "github.com/shirou/gopsutil/v4/disk"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCollectDisks_ZFS_Deduplication(t *testing.T) {
+	// Save original functions and restore after test
+	origDiskPartitions := diskPartitions
+	origDiskUsage := diskUsage
+	defer func() {
+		diskPartitions = origDiskPartitions
+		diskUsage = origDiskUsage
+	}()
+
+	// Mock partitions
+	diskPartitions = func(ctx context.Context, all bool) ([]godisk.PartitionStat, error) {
+		return []godisk.PartitionStat{
+			{Device: "tank/dataset1", Mountpoint: "/mnt/dataset1", Fstype: "zfs"},
+			{Device: "tank/dataset2", Mountpoint: "/mnt/dataset2", Fstype: "zfs"},
+			{Device: "tank/dataset3", Mountpoint: "/mnt/dataset3", Fstype: "zfs"},
+		}, nil
+	}
+
+	// Mock usage
+	diskUsage = func(ctx context.Context, path string) (*godisk.UsageStat, error) {
+		// All datasets on the same pool share the same free space
+		return &godisk.UsageStat{
+			Total:       1000,
+			Used:        500,
+			Free:        500,
+			UsedPercent: 50.0,
+		}, nil
+	}
+
+	// Mock zpool stats query to fail so we test the fallback logic (which uses the datasets)
+	// We can't easily mock queryZpoolStats because it's in zfs.go and not exported/variable-ized in the same way,
+	// but we can rely on the fact that `exec.LookPath("zpool")` will likely fail or `zpool list` will fail in the test environment.
+	// However, to be sure, we should probably mock it.
+	// But wait, `summarizeZFSPools` calls `queryZpoolStats`.
+	// `queryZpoolStats` is a variable in zfs.go! `var queryZpoolStats = fetchZpoolStats`
+
+	// Let's mock queryZpoolStats too.
+	origQueryZpoolStats := queryZpoolStats
+	defer func() { queryZpoolStats = origQueryZpoolStats }()
+
+	queryZpoolStats = func(ctx context.Context, pools []string) (map[string]zpoolStats, error) {
+		return nil, assert.AnError // Simulate failure
+	}
+
+	ctx := context.Background()
+	disks := collectDisks(ctx)
+
+	// Should return 1 disk (the pool), not 3
+	assert.Len(t, disks, 1)
+	assert.Equal(t, "tank", disks[0].Device)
+	assert.Equal(t, int64(1000), disks[0].TotalBytes)
+}
+
+func TestCollectDisks_ZFS_CaseInsensitive(t *testing.T) {
+	// Save original functions
+	origDiskPartitions := diskPartitions
+	origDiskUsage := diskUsage
+	defer func() {
+		diskPartitions = origDiskPartitions
+		diskUsage = origDiskUsage
+	}()
+
+	// Mock partitions with "ZFS" (uppercase)
+	diskPartitions = func(ctx context.Context, all bool) ([]godisk.PartitionStat, error) {
+		return []godisk.PartitionStat{
+			{Device: "tank/dataset1", Mountpoint: "/mnt/dataset1", Fstype: "ZFS"},
+			{Device: "tank/dataset2", Mountpoint: "/mnt/dataset2", Fstype: "ZFS"},
+		}, nil
+	}
+
+	diskUsage = func(ctx context.Context, path string) (*godisk.UsageStat, error) {
+		return &godisk.UsageStat{Total: 1000, Used: 500, Free: 500}, nil
+	}
+
+	// Mock queryZpoolStats to fail
+	origQueryZpoolStats := queryZpoolStats
+	defer func() { queryZpoolStats = origQueryZpoolStats }()
+	queryZpoolStats = func(ctx context.Context, pools []string) (map[string]zpoolStats, error) {
+		return nil, assert.AnError
+	}
+
+	ctx := context.Background()
+	disks := collectDisks(ctx)
+
+	// If case-sensitive check fails, it will return 2 disks
+	// If we fix it, it should return 1
+	assert.Len(t, disks, 1)
+}
+
+func TestCollectDisks_ZFS_Fuse(t *testing.T) {
+	// Save original functions
+	origDiskPartitions := diskPartitions
+	origDiskUsage := diskUsage
+	defer func() {
+		diskPartitions = origDiskPartitions
+		diskUsage = origDiskUsage
+	}()
+
+	// Mock partitions with "fuse.zfs"
+	diskPartitions = func(ctx context.Context, all bool) ([]godisk.PartitionStat, error) {
+		return []godisk.PartitionStat{
+			{Device: "tank/dataset1", Mountpoint: "/mnt/dataset1", Fstype: "fuse.zfs"},
+			{Device: "tank/dataset2", Mountpoint: "/mnt/dataset2", Fstype: "fuse.zfs"},
+		}, nil
+	}
+
+	diskUsage = func(ctx context.Context, path string) (*godisk.UsageStat, error) {
+		return &godisk.UsageStat{Total: 1000, Used: 500, Free: 500}, nil
+	}
+
+	// Mock queryZpoolStats to fail
+	origQueryZpoolStats := queryZpoolStats
+	defer func() { queryZpoolStats = origQueryZpoolStats }()
+	queryZpoolStats = func(ctx context.Context, pools []string) (map[string]zpoolStats, error) {
+		return nil, assert.AnError
+	}
+
+	ctx := context.Background()
+	disks := collectDisks(ctx)
+
+	// If check fails, it will return 2 disks
+	assert.Len(t, disks, 1)
+}