From d35fb1587ac50fa0d6311fb81ce9b48fd6ef2d22 Mon Sep 17 00:00:00 2001
From: Chris Coulson <chris.coulson@canonical.com>
Date: Tue, 28 Nov 2023 10:14:15 +0000
Subject: [PATCH] efi: add tests for grubHasPrefix

---
 efi/efi_test.go              |  3 +++
 efi/export_test.go           |  1 +
 efi/image_rules_defs_test.go |  1 +
 efi/image_rules_test.go      | 46 ++++++++++++++++++++++++++++++++++++
 4 files changed, 51 insertions(+)

diff --git a/efi/efi_test.go b/efi/efi_test.go
index 0663ed71..ad47bda1 100644
--- a/efi/efi_test.go
+++ b/efi/efi_test.go
@@ -216,6 +216,9 @@ type mockGrubImageHandle struct {
 }
 
 func (h *mockGrubImageHandle) Prefix() (string, error) {
+	if h.grubPrefix == "" {
+		return "", errors.New("no prefix")
+	}
 	return h.grubPrefix, nil
 }
 
diff --git a/efi/export_test.go b/efi/export_test.go
index 425a1221..e75402e4 100644
--- a/efi/export_test.go
+++ b/efi/export_test.go
@@ -87,6 +87,7 @@ var (
 // unexported members of some unexported types.
 type FwContext = fwContext
 type GrubFlags = grubFlags
+type GrubHasPrefix = grubHasPrefix
 type GrubImageHandle = grubImageHandle
 type GrubLoadHandler = grubLoadHandler
 type ImageLoadHandler = imageLoadHandler
diff --git a/efi/image_rules_defs_test.go b/efi/image_rules_defs_test.go
index 6ab36195..9f3422aa 100644
--- a/efi/image_rules_defs_test.go
+++ b/efi/image_rules_defs_test.go
@@ -271,6 +271,7 @@ func (s *imageRulesDefsSuite) TestFallbackNewImageLoadHandlerGrub(c *C) {
 	c.Assert(handler, testutil.ConvertibleTo, &GrubLoadHandler{})
 	c.Check(handler.(*GrubLoadHandler), DeepEquals, new(GrubLoadHandler))
 }
+
 func (s *imageRulesDefsSuite) TestFallbackNewImageLoadHandlerNull(c *C) {
 	// verify that an unrecognized leaf image is recognized by the fallback rules
 	image := newMockImage()
diff --git a/efi/image_rules_test.go b/efi/image_rules_test.go
index 45fd962d..2f1f689c 100644
--- a/efi/image_rules_test.go
+++ b/efi/image_rules_test.go
@@ -43,6 +43,17 @@ func (p *mockImagePredicate) Matches(image PeImageHandle) (bool, error) {
 
 type imageRulesSuite struct {
 	mockShimImageHandleMixin
+	mockGrubImageHandleMixin
+}
+
+func (s *imageRulesSuite) SetUpTest(c *C) {
+	s.mockShimImageHandleMixin.SetUpTest(c)
+	s.mockGrubImageHandleMixin.SetUpTest(c)
+}
+
+func (s *imageRulesSuite) TearDownTest(c *C) {
+	s.mockShimImageHandleMixin.TearDownTest(c)
+	s.mockGrubImageHandleMixin.TearDownTest(c)
 }
 
 var _ = Suite(&imageRulesSuite{})
@@ -492,6 +503,41 @@ func (s *imageRulesSuite) TestShimVersionErr(c *C) {
 	c.Check(err, ErrorMatches, `cannot obtain shim version: no version`)
 }
 
+func (s *imageRulesSuite) TestGrubHasPrefixTrue1(c *C) {
+	image := newMockImage().withGrubPrefix("/EFI/ubuntu").newPeImageHandle()
+
+	pred := GrubHasPrefix("/EFI/ubuntu")
+	match, err := pred.Matches(image)
+	c.Check(err, IsNil)
+	c.Check(match, testutil.IsTrue)
+}
+
+func (s *imageRulesSuite) TestGrubHasPrefixTrue2(c *C) {
+	image := newMockImage().withGrubPrefix("/EFI/debian").newPeImageHandle()
+
+	pred := GrubHasPrefix("/EFI/debian")
+	match, err := pred.Matches(image)
+	c.Check(err, IsNil)
+	c.Check(match, testutil.IsTrue)
+}
+
+func (s *imageRulesSuite) TestGrubHasPrefixFalse(c *C) {
+	image := newMockImage().withGrubPrefix("/EFI/debian").newPeImageHandle()
+
+	pred := GrubHasPrefix("/EFI/ubuntu")
+	match, err := pred.Matches(image)
+	c.Check(err, IsNil)
+	c.Check(match, testutil.IsFalse)
+}
+
+func (s *imageRulesSuite) TestGrubHasPrefixErr(c *C) {
+	image := newMockImage().newPeImageHandle()
+
+	pred := GrubHasPrefix("/EFI/ubuntu")
+	_, err := pred.Matches(image)
+	c.Check(err, ErrorMatches, `cannot obtain grub prefix: no prefix`)
+}
+
 func (s *imageRulesSuite) TestImageRulesMatch1(c *C) {
 	image := newMockImage().newPeImageHandle()