Refactor and simplify consent manager functions

internal/consent/consent.go

@@ -19,18 +19,6 @@ type Manager struct {
 	path string
 }
 
-// States is a struct that represents the consent states for a list of sources and the global consent state.
-type States struct {
-	SourceStates map[string]consentStateResult
-	GlobalState  consentStateResult
-}
-
-type consentStateResult struct {
-	Source  string // the source for the consent state
-	State   bool   // the consent state for the source
-	ReadErr bool   // true if there was an error reading the consent file
-}
-
 // consentFile is a struct that represents a consent file.
 type consentFile struct {
 	ConsentState bool `toml:"consent_state"`
@@ -42,130 +30,72 @@ func New(path string) *Manager {
 	return &Manager{path: path}
 }
 
-// GetConsentStates gets the consent state for the given sources and the global consent state.
-// If any of the sources do not have a consent file, it will be considered as a false state.
-// If a specified source does not have a consent file, it will not be included in the returned ConsentStates struct.
-// TODO: Simplify to single source, update spec
-func (cm Manager) GetConsentStates(sources []string) (consentStates States, err error) {
-	consentStates = States{SourceStates: make(map[string]consentStateResult)}
-
-	sourceFiles, globalFile, err := getMatchingConsentFiles(sources, cm.path)
+// GetConsentStates gets the consent state for the given source.
+// If the source do not have a consent file, it will be considered as a false state.
+// If the source is an empty string, then the global consent state will be returned.
+// If the target consent file does not exist, it will not be created.
+func (cm Manager) GetConsentState(source string) (bool, error) {
+	sourceConsent, err := readConsentFile(cm.getConsentFile(source))
 	if err != nil {
-		slog.Error("Error getting consent files", "error", err)
-		return States{}, err
-	}
-
-	results := make(chan consentStateResult, len(sourceFiles))
-	defer close(results)
-
-	// Global consent file
-	if globalFile != "" {
-		globalResult := make(chan consentStateResult, 1)
-		defer close(globalResult)
-
-		go func() {
-			globalConsent, err := readConsentFile(globalFile)
-			if err != nil {
-				slog.Error("Error reading global consent file", "file", globalFile, "error", err)
-				globalResult <- consentStateResult{Source: "global", State: false, ReadErr: true}
-				return
-			}
-			globalResult <- consentStateResult{Source: "global", State: globalConsent.ConsentState, ReadErr: false}
-		}()
-
-		consentStates.GlobalState = <-globalResult
-	}
-
-	// Goroutine to read the consent files for each source, excluding the global consent file.
-	for source, filePath := range sourceFiles {
-		go func(source, filePath string) {
-			consent, err := readConsentFile(filePath)
-			if err != nil {
-				slog.Error("Error reading consent file", "source", source, "error", err)
-				results <- consentStateResult{Source: source, State: false, ReadErr: true}
-				return
-			}
-			results <- consentStateResult{Source: source, State: consent.ConsentState, ReadErr: false}
-		}(source, filePath)
-	}
-
-	for range sourceFiles {
-		res := <-results
-		consentStates.SourceStates[res.Source] = res
+		slog.Error("Error reading source consent file", "source", source, "error", err)
+		return false, err
 	}
 
-	return consentStates, nil
+	return sourceConsent.ConsentState, nil
 }
 
 var consentSourceFilePattern = `%s` + constants.ConsentSourceBaseSeparator + constants.GlobalFileName
 
 // SetConsentState updates the consent state for the given source.
 // If the source is an empty string, then the global consent state will be set.
 // If the target consent file does not exist, it will be created.
-func (cm *Manager) SetConsentState(source string, state bool) (err error) {
+func (cm Manager) SetConsentState(source string, state bool) (err error) {
 	defer decorate.OnError(&err, "could not set consent state:")
 
+	consent := consentFile{ConsentState: state}
+	return consent.write(cm.getConsentFile(source))
+}
+
+// getConsentFile returns the expected path to the consent file for the given source.
+// If source is blank, it returns the path to the global consent file.
+// It does not check if the file exists, or if it is valid.
+func (cm Manager) getConsentFile(source string) string {
 	p := filepath.Join(cm.path, constants.GlobalFileName)
 	if source != "" {
 		p = filepath.Join(cm.path, fmt.Sprintf(consentSourceFilePattern, source))
 	}
 
-	consent := consentFile{ConsentState: state}
-	return consent.write(p)
+	return p
 }
 
-// getMatchingConsentFiles returns a map of all paths to consent files matching the given sources and a path
-// to the global consent file.
-// If sources is empty, all consent files in the folder will be returned.
-// If a source does not have a consent file, it will be represented as an empty string
-// Does not traverse subdirectories.
-func getMatchingConsentFiles(sources []string, folderPath string) (sourceFiles map[string]string, globalFile string, err error) {
-	sourceFiles = make(map[string]string)
+// getSourceConsentFiles returns a map of all paths to validly named consent files in the folder, other than the global file.
+func (cm Manager) getConsentFiles() (map[string]string, error) {
+	sourceFiles := make(map[string]string)
 
-	entries, err := os.ReadDir(folderPath)
+	entries, err := os.ReadDir(cm.path)
 	if err != nil {
-		return sourceFiles, globalFile, err
+		return sourceFiles, err
 	}
 
 	for _, entry := range entries {
 		if entry.IsDir() {
 			continue
 		}
 
-		// Global file
-		if entry.Name() == constants.GlobalFileName {
-			globalFile = filepath.Join(folderPath, entry.Name())
-			slog.Debug("Found global consent file", "file", globalFile)
+		// Source file
+		if !strings.HasSuffix(entry.Name(), constants.ConsentSourceBaseSeparator+constants.GlobalFileName) {
 			continue
 		}
-
-		if len(sources) == 0 {
-			// Source file
-			if !strings.HasSuffix(entry.Name(), constants.ConsentSourceBaseSeparator+constants.GlobalFileName) {
-				continue
-			}
-			source := strings.TrimSuffix(entry.Name(), constants.ConsentSourceBaseSeparator+constants.GlobalFileName)
-			sourceFiles[source] = filepath.Join(folderPath, entry.Name())
-			slog.Debug("Found source consent file", "file", sourceFiles[source])
-			continue
-		}
-
-		for _, source := range sources {
-			if entry.Name() == fmt.Sprintf(consentSourceFilePattern, source) {
-				sourceFiles[source] = filepath.Join(folderPath, entry.Name())
-				slog.Debug("Found matching source consent file", "file", sourceFiles[source])
-				break
-			}
-		}
+		source := strings.TrimSuffix(entry.Name(), constants.ConsentSourceBaseSeparator+constants.GlobalFileName)
+		sourceFiles[source] = filepath.Join(cm.path, entry.Name())
+		slog.Debug("Found source consent file", "file", sourceFiles[source])
 	}
-	return sourceFiles, globalFile, err
+
+	return sourceFiles, nil
 }
 
 func readConsentFile(path string) (*consentFile, error) {
 	var consent consentFile
-	if _, err := os.Stat(path); os.IsNotExist(err) {
-		return &consent, nil
-	}
 	_, err := toml.DecodeFile(path, &consent)
 	slog.Debug("Read consent file", "file", path, "consent", consent.ConsentState)
 
@@ -181,7 +111,7 @@ func (cf consentFile) write(path string) (err error) {
 	}
 	defer func() {
 		_ = tmp.Close()
-		if err := os.Remove(tmp.Name()); err != nil {
+		if err := os.Remove(tmp.Name()); err != nil && !os.IsNotExist(err) {
 			slog.Warn("Failed to remove temporary file when writing consent file", "file", tmp.Name(), "error", err)
 		}
 	}()

internal/consent/consent_test.go

@@ -13,37 +13,35 @@ import (
 	"github.com/ubuntu/ubuntu-insights/internal/testutils"
 )
 
-// consentDir is a struct that holds a test's temporary directory and its locks.
+// consentDir is a struct that holds a test's temporary directory.
 // It should be cleaned up after the test is done.
 type consentDir struct {
 	dir string
 }
 
-func TestGetConsentStates(t *testing.T) {
+func TestGetConsentState(t *testing.T) {
 	t.Parallel()
 
 	tests := map[string]struct {
-		sources    []string
+		source     string
 		globalFile string
 
 		wantErr bool
 	}{
-		"No Global File": {},
+		"No Global File": {wantErr: true},
 
 		// Global File Tests
 		"Valid True Global File":    {globalFile: "valid_true-consent.toml"},
 		"Valid False Global File":   {globalFile: "valid_false-consent.toml"},
-		"Invalid Value Global File": {globalFile: "invalid_value-consent.toml"},
-		"Invalid File Global File":  {globalFile: "invalid_file-consent.toml"},
+		"Invalid Value Global File": {globalFile: "invalid_value-consent.toml", wantErr: true},
+		"Invalid File Global File":  {globalFile: "invalid_file-consent.toml", wantErr: true},
 
 		// Source Specific Tests
-		"Valid True Global File, Valid True Source":                       {globalFile: "valid_true-consent.toml", sources: []string{"valid_true"}},
-		"Valid True Global File, Valid False Source":                      {globalFile: "valid_true-consent.toml", sources: []string{"valid_false"}},
-		"Valid True Global File, Invalid Value Source":                    {globalFile: "valid_true-consent.toml", sources: []string{"invalid_value"}},
-		"Valid True Global File, Invalid File Source":                     {globalFile: "valid_true-consent.toml", sources: []string{"invalid_file"}},
-		"Valid True Global File, No File Source":                          {globalFile: "valid_true-consent.toml", sources: []string{"not_a_file"}},
-		"Valid True Global File, 2 Multiple Sources (VTrue, VFalse)":      {globalFile: "valid_true-consent.toml", sources: []string{"valid_true", "valid_false"}},
-		"Valid True Global File, 3 Multiple Sources (VTrue, VFalse, NAF)": {globalFile: "valid_true-consent.toml", sources: []string{"valid_true", "valid_false", "not_a_file"}},
+		"Valid True Global File, Valid True Source":    {globalFile: "valid_true-consent.toml", source: "valid_true"},
+		"Valid True Global File, Valid False Source":   {globalFile: "valid_true-consent.toml", source: "valid_false"},
+		"Valid True Global File, Invalid Value Source": {globalFile: "valid_true-consent.toml", source: "invalid_value", wantErr: true},
+		"Valid True Global File, Invalid File Source":  {globalFile: "valid_true-consent.toml", source: "invalid_file", wantErr: true},
+		"Valid True Global File, No File Source":       {globalFile: "valid_true-consent.toml", source: "not_a_file", wantErr: true},
 	}
 
 	for name, tc := range tests {
@@ -54,15 +52,15 @@ func TestGetConsentStates(t *testing.T) {
 			defer cDir.cleanup(t)
 			cm := consent.New(cDir.dir)
 
-			got, err := cm.GetConsentStates(tc.sources)
+			got, err := cm.GetConsentState(tc.source)
 			if tc.wantErr {
 				require.Error(t, err, "expected an error but got none")
 				return
 			}
 			require.NoError(t, err, "got an unexpected error")
 
 			want := testutils.LoadWithUpdateFromGoldenYAML(t, got)
-			require.Equal(t, want, got, "GetConsentStates should return expected consent states")
+			require.Equal(t, want, got, "GetConsentState should return expected consent state")
 		})
 	}
 }
@@ -71,7 +69,6 @@ func TestSetConsentStates(t *testing.T) {
 	t.Parallel()
 
 	tests := map[string]struct {
-		sources       []string
 		consentStates map[string]bool
 		globalFile    string
 
@@ -100,8 +97,8 @@ func TestSetConsentStates(t *testing.T) {
 	}
 
 	type goldenFile struct {
-		States    consent.States
-		FileCount uint
+		States    map[string]bool
+		FileCount int
 	}
 
 	for name, tc := range tests {
@@ -119,12 +116,12 @@ func TestSetConsentStates(t *testing.T) {
 			}
 			require.NoError(t, err, "got an unexpected error")
 
-			states, err := cm.GetConsentStates(tc.sources)
+			states, err := cm.GetAllSourceConsentStates(true)
 			require.NoError(t, err, "got an unexpected error while getting consent states")
 
 			d, err := os.ReadDir(cDir.dir)
 			require.NoError(t, err, "failed to read temporary directory")
-			got := goldenFile{States: states, FileCount: uint(len(d))}
+			got := goldenFile{States: states, FileCount: len(d)}
 
 			want := testutils.LoadWithUpdateFromGoldenYAML(t, got)
 			require.Equal(t, want, got, "GetConsentStates should return expected consent states")

internal/consent/export_test.go

@@ -0,0 +1,26 @@
+package consent
+
+// GetAllSourceConsentStates gets the consent states for all sources.
+// It does not get the global consent state.
+// If continueOnErr is true, it will continue to the next source if an error occurs.
+func (cm Manager) GetAllSourceConsentStates(continueOnErr bool) (map[string]bool, error) {
+	p, err := cm.getConsentFiles()
+	if err != nil {
+		return nil, err
+	}
+
+	consentStates := make(map[string]bool)
+	for source, path := range p {
+		consent, err := readConsentFile(path)
+		if err != nil && !continueOnErr {
+			return nil, err
+		}
+		if err != nil {
+			continue
+		}
+
+		consentStates[source] = consent.ConsentState
+	}
+
+	return consentStates, nil
+}

internal/consent/testdata/TestGetConsentState/golden/valid_false_global_file

@@ -0,0 +1 @@
+false

internal/consent/testdata/TestGetConsentState/golden/valid_true_global_file

@@ -0,0 +1 @@
+true

internal/consent/testdata/TestGetConsentState/golden/valid_true_global_file,_valid_false_source

@@ -0,0 +1 @@
+false

internal/consent/testdata/TestGetConsentState/golden/valid_true_global_file,_valid_true_source

@@ -0,0 +1 @@
+true