diff --git a/pkg/partition_info/partition_info.go b/pkg/partition_info/partition_info.go
index 10b04ae..3bbf986 100644
--- a/pkg/partition_info/partition_info.go
+++ b/pkg/partition_info/partition_info.go
@@ -10,12 +10,56 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
+const DefaultPartitionInfoFile = "/oem/partition_info.yaml"
+
 // PartitionInfo maps a partition label to a partition UUID.
 // It's used in order to be able to ask the kcrypt-challenger for the passphrase
 // using the partition label, even when the label is not accessible (e.g. before
 // decrypting the partition). The UUID can be used to lookup the partition label
 // and make the request.
-type PartitionInfo map[string]string
+type PartitionInfo struct {
+	file    string
+	mapping map[string]string
+}
+
+func NewPartitionInfoFromFile(file string) (*PartitionInfo, error) {
+	mapping, err := ParsePartitionInfoFile(file)
+	if err != nil {
+		return nil, err
+	}
+
+	return &PartitionInfo{
+		file:    file,
+		mapping: mapping,
+	}, nil
+}
+
+func (pi PartitionInfo) LookupUUIDForLabel(l string) string {
+	return pi.mapping[l]
+}
+
+// UpdatePartitionLabelMapping takes partition information as a string argument
+// the the form: `label:name:uuid` (that's what the `kcrypt encrypt` command returns
+// on success. This function stores it in the PartitionInfoFile yaml file for
+// later use.
+func (pi PartitionInfo) UpdateMapping(partitionData string) error {
+	label, uuid := PartitionDataFromString(partitionData)
+	pi.mapping[label] = uuid
+
+	return pi.save()
+}
+
+func (pi PartitionInfo) save() error {
+	data, err := yaml.Marshal(&pi.mapping)
+	if err != nil {
+		return errors.Wrap(err, "marshalling the new partition info to yaml")
+	}
+	err = ioutil.WriteFile(pi.file, data, 0)
+	if err != nil {
+		return errors.Wrap(err, "writing back the partition info file")
+	}
+	return nil
+}
 
 func PartitionToString(p *block.Partition) string {
 	return fmt.Sprintf("%s:%s:%s", p.Label, p.Name, p.UUID)
@@ -29,29 +73,14 @@ func PartitionDataFromString(partitionStr string) (string, string) {
 	return parts[0], parts[2]
 }
 
-// UpdatePartitionLabelMapping takes partition information as a string argument
-// the the form: `label:name:uuid` (that's what the `kcrypt encrypt` command returns
-// on success. This function stores it in the PartitionInfoFile yaml file for
-// later use.
-func UpdatePartitionLabelMapping(partitionData, file string) error {
-	partitionInfo, err := ParsePartitionInfoFile(file)
-	if err != nil {
-		return err
-	}
-
-	label, uuid := PartitionDataFromString(partitionData)
-	partitionInfo[label] = uuid
-
-	return UpdatePartitionInfoFile(partitionInfo, file)
-}
-
-func ParsePartitionInfoFile(file string) (PartitionInfo, error) {
-	var result PartitionInfo
+func ParsePartitionInfoFile(file string) (map[string]string, error) {
+	var result map[string]string
 
 	yamlFile, err := ioutil.ReadFile(file)
 	if err != nil {
 		return result, errors.Wrap(err, "reading the partition info file")
 	}
+
 	err = yaml.Unmarshal(yamlFile, &result)
 	if err != nil {
 		return result, errors.Wrap(err, "unmarshalling partition info file")
@@ -59,15 +88,3 @@ func ParsePartitionInfoFile(file string) (PartitionInfo, error) {
 
 	return result, nil
 }
-
-func UpdatePartitionInfoFile(partitionInfo PartitionInfo, file string) error {
-	data, err := yaml.Marshal(&partitionInfo)
-	if err != nil {
-		return errors.Wrap(err, "marshalling the new partition info to yaml")
-	}
-	err = ioutil.WriteFile(file, data, 0)
-	if err != nil {
-		return errors.Wrap(err, "writing back the partition info file")
-	}
-	return nil
-}
diff --git a/pkg/partition_info/partition_info_test.go b/pkg/partition_info/partition_info_test.go
index eced9d8..1a1e22d 100644
--- a/pkg/partition_info/partition_info_test.go
+++ b/pkg/partition_info/partition_info_test.go
@@ -12,6 +12,20 @@ import (
 )
 
 var _ = Describe("Partition Info file parsing", func() {
+	Describe("NewPartitionInfoFromFile", func() {
+		var file string
+
+		BeforeEach(func() {
+			file = "../../tests/assets/partition_info.yaml"
+		})
+
+		It("returns a PartitionInfo", func() {
+			result, err := pi.NewPartitionInfoFromFile(file)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(result).ToNot(BeNil())
+		})
+	})
+
 	Describe("ParsePartitionInfoFile", func() {
 		var file string
 
@@ -28,79 +42,6 @@ var _ = Describe("Partition Info file parsing", func() {
 		})
 	})
 
-	Describe("UpdatePartitionInfoFile", func() {
-		var file *os.File
-		var err error
-
-		BeforeEach(func() {
-			file, err = ioutil.TempFile("", "parition-info.*.yaml")
-			Expect(err).ToNot(HaveOccurred())
-
-			fileContents := `
-- label: TO_DELETE
-  uuid: old_uuid_1
-- label: TO_BE_UPDATED
-  uuid: old_uuid_2
-`
-			_, err = file.Write([]byte(fileContents))
-			Expect(err).ToNot(HaveOccurred())
-		})
-
-		AfterEach(func() {
-			os.Remove(file.Name())
-		})
-
-		It("Updates the partition file correctly", func() {
-			partitionInfo := pi.PartitionInfo{
-				"TO_BE_UPDATED": "new_uuid_1",
-				"NEW_PARTITION": "new_uuid_2",
-			}
-
-			err = pi.UpdatePartitionInfoFile(partitionInfo, file.Name())
-			Expect(err).ToNot(HaveOccurred())
-
-			dat, err := os.ReadFile(file.Name())
-			Expect(err).ToNot(HaveOccurred())
-
-			expectedContent := `NEW_PARTITION: new_uuid_2
-TO_BE_UPDATED: new_uuid_1
-`
-			Expect(string(dat)).To(Equal(expectedContent))
-		})
-	})
-
-	Describe("UpdatePartitionLabelMapping", func() {
-		var file *os.File
-		var err error
-
-		BeforeEach(func() {
-			file, err = ioutil.TempFile("", "parition-info.*.yaml")
-			Expect(err).ToNot(HaveOccurred())
-
-			_, err = file.Write([]byte("TO_KEEP: old_uuid_1"))
-			Expect(err).ToNot(HaveOccurred())
-		})
-
-		AfterEach(func() {
-			os.Remove(file.Name())
-		})
-
-		It("Updates the file correctly from a `kcrypt encrypt` return value", func() {
-			partitionData := "TO_BE_ADDED:some_name:new_uuid"
-
-			err = pi.UpdatePartitionLabelMapping(partitionData, file.Name())
-			Expect(err).ToNot(HaveOccurred())
-
-			dat, err := os.ReadFile(file.Name())
-			Expect(err).ToNot(HaveOccurred())
-
-			expectedContent := `TO_BE_ADDED: new_uuid
-TO_KEEP: old_uuid_1
-`
-			Expect(string(dat)).To(Equal(expectedContent))
-		})
-	})
-
 	Describe("PartitionToString", func() {
 		var partition *block.Partition
 
@@ -132,4 +73,57 @@ TO_KEEP: old_uuid_1
 			Expect(uuid).To(Equal("the_uuid"))
 		})
 	})
+
+	Describe("UpdateMapping", func() {
+		var file *os.File
+		var err error
+		var partitionInfo *pi.PartitionInfo
+
+		BeforeEach(func() {
+			file, err = ioutil.TempFile("", "partition-info.*.yaml")
+			Expect(err).ToNot(HaveOccurred())
+
+			_, err = file.Write([]byte("TO_KEEP: old_uuid_1"))
+			Expect(err).ToNot(HaveOccurred())
+
+			partitionInfo, err = pi.NewPartitionInfoFromFile(file.Name())
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		AfterEach(func() {
+			os.Remove(file.Name())
+		})
+
+		It("Updates the file correctly from a `kcrypt encrypt` return value", func() {
+			partitionData := "TO_BE_ADDED:some_name:new_uuid"
+
+			err = partitionInfo.UpdateMapping(partitionData)
+			Expect(err).ToNot(HaveOccurred())
+
+			dat, err := os.ReadFile(file.Name())
+			Expect(err).ToNot(HaveOccurred())
+
+			expectedContent := `TO_BE_ADDED: new_uuid
+TO_KEEP: old_uuid_1
+`
+			Expect(string(dat)).To(Equal(expectedContent))
+		})
+	})
+
+	Describe("LookupUUIDForLabel", func() {
+		var file string
+		var partitionInfo *pi.PartitionInfo
+		var err error
+
+		BeforeEach(func() {
+			file = "../../tests/assets/partition_info.yaml"
+			partitionInfo, err = pi.NewPartitionInfoFromFile(file)
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		It("parses the file correctly", func() {
+			uuid := partitionInfo.LookupUUIDForLabel("COS_PERSISTENT")
+			Expect(uuid).To(Equal("some_uuid_1"))
+		})
+	})
 })