From 26c7be5311604bf9aa8fe3e042fecb618f7ef86e Mon Sep 17 00:00:00 2001
From: Christian Kellner <christian@kellner.me>
Date: Sun, 20 Feb 2022 13:23:07 +0100
Subject: [PATCH] disk: add basic check for ForEachEntity

Assure that we visit all the entities.

Co-Authored-By: Achilleas Koutsou <achilleas@koutsou.net>
---
 internal/disk/disk_test.go | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/internal/disk/disk_test.go b/internal/disk/disk_test.go
index 8f397e91e..800ea71c9 100644
--- a/internal/disk/disk_test.go
+++ b/internal/disk/disk_test.go
@@ -154,3 +154,21 @@ func TestDisk_ForEachFilesystem(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, []*Filesystem{bootFs}, expectedFilesystems)
 }
+
+func TestDisk_ForEachEntity(t *testing.T) {
+
+	count := 0
+	err := canonicalPartitionTable.ForEachEntity(func(e Entity, path []Entity) error {
+		assert.NotNil(t, e)
+		assert.NotNil(t, path)
+
+		count += 1
+		return nil
+	})
+
+	assert.NoError(t, err)
+
+	// PartitionTable, 4 partitions, 3 filesystems -> 8 entities
+	assert.Equal(t, 8, count)
+
+}