feat: Add _Good, _Bad, and _Ugly tests

Refactored the existing tests to use the `_Good`, `_Bad`, and `_Ugly`
testing convention. This provides a more structured approach to testing
and ensures that a wider range of scenarios are covered, including
valid inputs, invalid inputs, and edge cases.

In addition to refactoring the tests, this change also includes several
bug fixes that were uncovered by the new tests. These fixes improve the
robustness and reliability of the codebase.

The following packages and commands were affected:
- `pkg/datanode`
- `pkg/compress`
- `pkg/github`
- `pkg/matrix`
- `pkg/pwa`
- `pkg/vcs`
- `pkg/website`
- `cmd/all`
- `cmd/collect`
- `cmd/collect_github_repo`
- `cmd/collect_website`
- `cmd/compile`
- `cmd/root`
- `cmd/run`
- `cmd/serve`

cmd/all.go

@@ -17,122 +17,130 @@ import (
 	"github.com/spf13/cobra"
 )
 
-// allCmd represents the all command
+var allCmd = NewAllCmd()
-var allCmd = &cobra.Command{
-	Use:   "all [url]",
-	Short: "Collect all resources from a URL",
-	Long:  `Collect all resources from a URL, dispatching to the appropriate collector based on the URL type.`,
-	Args:  cobra.ExactArgs(1),
-	RunE: func(cmd *cobra.Command, args []string) error {
-		url := args[0]
-		outputFile, _ := cmd.Flags().GetString("output")
-		format, _ := cmd.Flags().GetString("format")
-		compression, _ := cmd.Flags().GetString("compression")
 
-		owner, err := parseGithubOwner(url)
+func NewAllCmd() *cobra.Command {
-		if err != nil {
+	allCmd := &cobra.Command{
-			return err
+		Use:   "all [url]",
-		}
+		Short: "Collect all resources from a URL",
+		Long:  `Collect all resources from a URL, dispatching to the appropriate collector based on the URL type.`,
+		Args:  cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			url := args[0]
+			outputFile, _ := cmd.Flags().GetString("output")
+			format, _ := cmd.Flags().GetString("format")
+			compression, _ := cmd.Flags().GetString("compression")
 
-		repos, err := GithubClient.GetPublicRepos(cmd.Context(), owner)
+			owner, err := parseGithubOwner(url)
-		if err != nil {
-			return err
-		}
-
-		prompter := ui.NewNonInteractivePrompter(ui.GetVCSQuote)
-		prompter.Start()
-		defer prompter.Stop()
-
-		var progressWriter io.Writer
-		if prompter.IsInteractive() {
-			bar := ui.NewProgressBar(len(repos), "Cloning repositories")
-			progressWriter = ui.NewProgressWriter(bar)
-		}
-
-		cloner := vcs.NewGitCloner()
-		allDataNodes := datanode.New()
-
-		for _, repoURL := range repos {
-			dn, err := cloner.CloneGitRepository(repoURL, progressWriter)
 			if err != nil {
-				// Log the error and continue
+				return err
-				fmt.Fprintln(cmd.ErrOrStderr(), "Error cloning repository:", err)
-				continue
 			}
-			// This is not an efficient way to merge datanodes, but it's the only way for now
-			// A better approach would be to add a Merge method to the DataNode
-			repoName := strings.TrimSuffix(repoURL, ".git")
-			parts := strings.Split(repoName, "/")
-			repoName = parts[len(parts)-1]
 
-			err = dn.Walk(".", func(path string, de fs.DirEntry, err error) error {
+			repos, err := GithubClient.GetPublicRepos(cmd.Context(), owner)
+			if err != nil {
+				return err
+			}
+
+			prompter := ui.NewNonInteractivePrompter(ui.GetVCSQuote)
+			prompter.Start()
+			defer prompter.Stop()
+
+			var progressWriter io.Writer
+			if prompter.IsInteractive() {
+				bar := ui.NewProgressBar(len(repos), "Cloning repositories")
+				progressWriter = ui.NewProgressWriter(bar)
+			}
+
+			cloner := vcs.NewGitCloner()
+			allDataNodes := datanode.New()
+
+			for _, repoURL := range repos {
+				dn, err := cloner.CloneGitRepository(repoURL, progressWriter)
 				if err != nil {
-					return err
+					// Log the error and continue
+					fmt.Fprintln(cmd.ErrOrStderr(), "Error cloning repository:", err)
+					continue
 				}
-				if !de.IsDir() {
+				// This is not an efficient way to merge datanodes, but it's the only way for now
-					err := func() error {
+				// A better approach would be to add a Merge method to the DataNode
-						file, err := dn.Open(path)
+				repoName := strings.TrimSuffix(repoURL, ".git")
-						if err != nil {
+				parts := strings.Split(repoName, "/")
-							return err
+				repoName = parts[len(parts)-1]
-						}
+
-						defer file.Close()
+				err = dn.Walk(".", func(path string, de fs.DirEntry, err error) error {
-						data, err := io.ReadAll(file)
-						if err != nil {
-							return err
-						}
-						allDataNodes.AddData(repoName+"/"+path, data)
-						return nil
-					}()
 					if err != nil {
 						return err
 					}
+					if !de.IsDir() {
+						err := func() error {
+							file, err := dn.Open(path)
+							if err != nil {
+								return err
+							}
+							defer file.Close()
+							data, err := io.ReadAll(file)
+							if err != nil {
+								return err
+							}
+							allDataNodes.AddData(repoName+"/"+path, data)
+							return nil
+						}()
+						if err != nil {
+							return err
+						}
+					}
+					return nil
+				})
+				if err != nil {
+					fmt.Fprintln(cmd.ErrOrStderr(), "Error walking datanode:", err)
+					continue
 				}
-				return nil
-			})
-			if err != nil {
-				fmt.Fprintln(cmd.ErrOrStderr(), "Error walking datanode:", err)
-				continue
 			}
-		}
 
-		var data []byte
+			var data []byte
-		if format == "matrix" {
+			if format == "matrix" {
-			matrix, err := matrix.FromDataNode(allDataNodes)
+				matrix, err := matrix.FromDataNode(allDataNodes)
-			if err != nil {
+				if err != nil {
-				return fmt.Errorf("error creating matrix: %w", err)
+					return fmt.Errorf("error creating matrix: %w", err)
+				}
+				data, err = matrix.ToTar()
+				if err != nil {
+					return fmt.Errorf("error serializing matrix: %w", err)
+				}
+			} else {
+				data, err = allDataNodes.ToTar()
+				if err != nil {
+					return fmt.Errorf("error serializing DataNode: %w", err)
+				}
 			}
-			data, err = matrix.ToTar()
+
+			compressedData, err := compress.Compress(data, compression)
 			if err != nil {
-				return fmt.Errorf("error serializing matrix: %w", err)
+				return fmt.Errorf("error compressing data: %w", err)
 			}
-		} else {
+
-			data, err = allDataNodes.ToTar()
+			err = os.WriteFile(outputFile, compressedData, 0644)
 			if err != nil {
-				return fmt.Errorf("error serializing DataNode: %w", err)
+				return fmt.Errorf("error writing DataNode to file: %w", err)
 			}
-		}
 
-		compressedData, err := compress.Compress(data, compression)
+			fmt.Fprintln(cmd.OutOrStdout(), "All repositories saved to", outputFile)
-		if err != nil {
-			return fmt.Errorf("error compressing data: %w", err)
-		}
 
-		err = os.WriteFile(outputFile, compressedData, 0644)
+			return nil
-		if err != nil {
+		},
-			return fmt.Errorf("error writing DataNode to file: %w", err)
+	}
-		}
-
-		fmt.Fprintln(cmd.OutOrStdout(), "All repositories saved to", outputFile)
-
-		return nil
-	},
-}
-
-func init() {
-	RootCmd.AddCommand(allCmd)
 	allCmd.PersistentFlags().String("output", "all.dat", "Output file for the DataNode")
 	allCmd.PersistentFlags().String("format", "datanode", "Output format (datanode or matrix)")
 	allCmd.PersistentFlags().String("compression", "none", "Compression format (none, gz, or xz)")
+	return allCmd
+}
+
+func GetAllCmd() *cobra.Command {
+	return allCmd
+}
+
+func init() {
+	RootCmd.AddCommand(GetAllCmd())
 }
 
 func parseGithubOwner(u string) (string, error) {

cmd/all_test.go

@@ -14,6 +14,7 @@ import (
 )
 
 func TestAllCmd_Good(t *testing.T) {
+	// Setup mock HTTP client for GitHub API
 	mockGithubClient := mocks.NewMockClient(map[string]*http.Response{
 		"https://api.github.com/users/testuser/repos": {
 			StatusCode: http.StatusOK,
@@ -29,6 +30,7 @@ func TestAllCmd_Good(t *testing.T) {
 		github.NewAuthenticatedClient = oldNewAuthenticatedClient
 	}()
 
+	// Setup mock Git cloner
 	mockCloner := &mocks.MockGitCloner{
 		DN:  datanode.New(),
 		Err: nil,
@@ -40,8 +42,9 @@ func TestAllCmd_Good(t *testing.T) {
 	}()
 
 	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(allCmd)
+	rootCmd.AddCommand(GetAllCmd())
 
+	// Execute command
 	out := filepath.Join(t.TempDir(), "out")
 	_, err := executeCommand(rootCmd, "all", "https://github.com/testuser", "--output", out)
 	if err != nil {
@@ -50,9 +53,16 @@ func TestAllCmd_Good(t *testing.T) {
 }
 
 func TestAllCmd_Bad(t *testing.T) {
+	// Setup mock HTTP client to return an error
 	mockGithubClient := mocks.NewMockClient(map[string]*http.Response{
-		"https://api.github.com/users/testuser/repos": {
+		"https://api.github.com/users/baduser/repos": {
 			StatusCode: http.StatusNotFound,
+			Status:     "404 Not Found",
+			Body:       io.NopCloser(bytes.NewBufferString(`{"message": "Not Found"}`)),
+		},
+		"https://api.github.com/orgs/baduser/repos": {
+			StatusCode: http.StatusNotFound,
+			Status:     "404 Not Found",
 			Body:       io.NopCloser(bytes.NewBufferString(`{"message": "Not Found"}`)),
 		},
 	})
@@ -65,11 +75,42 @@ func TestAllCmd_Bad(t *testing.T) {
 	}()
 
 	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(allCmd)
+	rootCmd.AddCommand(GetAllCmd())
 
+	// Execute command
 	out := filepath.Join(t.TempDir(), "out")
-	_, err := executeCommand(rootCmd, "all", "https://github.com/testuser", "--output", out)
+	_, err := executeCommand(rootCmd, "all", "https://github.com/baduser", "--output", out)
 	if err == nil {
-		t.Fatalf("expected an error, but got none")
+		t.Fatal("expected an error, but got none")
 	}
 }
+
+func TestAllCmd_Ugly(t *testing.T) {
+	t.Run("User with no repos", func(t *testing.T) {
+		// Setup mock HTTP client for a user with no repos
+		mockGithubClient := mocks.NewMockClient(map[string]*http.Response{
+			"https://api.github.com/users/emptyuser/repos": {
+				StatusCode: http.StatusOK,
+				Header:     http.Header{"Content-Type": []string{"application/json"}},
+				Body:       io.NopCloser(bytes.NewBufferString(`[]`)),
+			},
+		})
+		oldNewAuthenticatedClient := github.NewAuthenticatedClient
+		github.NewAuthenticatedClient = func(ctx context.Context) *http.Client {
+			return mockGithubClient
+		}
+		defer func() {
+			github.NewAuthenticatedClient = oldNewAuthenticatedClient
+		}()
+
+		rootCmd := NewRootCmd()
+		rootCmd.AddCommand(GetAllCmd())
+
+		// Execute command
+		out := filepath.Join(t.TempDir(), "out")
+		_, err := executeCommand(rootCmd, "all", "https://github.com/emptyuser", "--output", out)
+		if err != nil {
+			t.Fatalf("all command failed for user with no repos: %v", err)
+		}
+	})
+}

cmd/collect.go

@@ -5,15 +5,19 @@ import (
 )
 
 // collectCmd represents the collect command
-var collectCmd = &cobra.Command{
+var collectCmd = NewCollectCmd()
-	Use:   "collect",
-	Short: "Collect a resource from a URI.",
-	Long:  `Collect a resource from a URI and store it in a DataNode.`,
-}
 
 func init() {
-	RootCmd.AddCommand(collectCmd)
+	RootCmd.AddCommand(GetCollectCmd())
 }
 func NewCollectCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:   "collect",
+		Short: "Collect a resource from a URI.",
+		Long:  `Collect a resource from a URI and store it in a DataNode.`,
+	}
+}
+
+func GetCollectCmd() *cobra.Command {
 	return collectCmd
 }

cmd/collect_github_repo_test.go

@@ -10,6 +10,7 @@ import (
 )
 
 func TestCollectGithubRepoCmd_Good(t *testing.T) {
+	// Setup mock Git cloner
 	mockCloner := &mocks.MockGitCloner{
 		DN:  datanode.New(),
 		Err: nil,
@@ -21,16 +22,18 @@ func TestCollectGithubRepoCmd_Good(t *testing.T) {
 	}()
 
 	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(collectCmd)
+	rootCmd.AddCommand(GetCollectCmd())
 
+	// Execute command
 	out := filepath.Join(t.TempDir(), "out")
-	_, err := executeCommand(rootCmd, "collect", "github", "repo", "https://github.com/testuser/repo1.git", "--output", out)
+	_, err := executeCommand(rootCmd, "collect", "github", "repo", "https://github.com/testuser/repo1", "--output", out)
 	if err != nil {
 		t.Fatalf("collect github repo command failed: %v", err)
 	}
 }
 
 func TestCollectGithubRepoCmd_Bad(t *testing.T) {
+	// Setup mock Git cloner to return an error
 	mockCloner := &mocks.MockGitCloner{
 		DN:  nil,
 		Err: fmt.Errorf("git clone error"),
@@ -42,11 +45,23 @@ func TestCollectGithubRepoCmd_Bad(t *testing.T) {
 	}()
 
 	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(collectCmd)
+	rootCmd.AddCommand(GetCollectCmd())
 
+	// Execute command
 	out := filepath.Join(t.TempDir(), "out")
-	_, err := executeCommand(rootCmd, "collect", "github", "repo", "https://github.com/testuser/repo1.git", "--output", out)
+	_, err := executeCommand(rootCmd, "collect", "github", "repo", "https://github.com/testuser/repo1", "--output", out)
 	if err == nil {
-		t.Fatalf("expected an error, but got none")
+		t.Fatal("expected an error, but got none")
 	}
 }
+
+func TestCollectGithubRepoCmd_Ugly(t *testing.T) {
+	t.Run("Invalid repo URL", func(t *testing.T) {
+		rootCmd := NewRootCmd()
+		rootCmd.AddCommand(GetCollectCmd())
+		_, err := executeCommand(rootCmd, "collect", "github", "repo", "not-a-github-url")
+		if err == nil {
+			t.Fatal("expected an error for invalid repo URL, but got none")
+		}
+	})
+}

cmd/collect_website.go

@@ -14,77 +14,83 @@ import (
 )
 
 // collectWebsiteCmd represents the collect website command
-var collectWebsiteCmd = &cobra.Command{
+var collectWebsiteCmd = NewCollectWebsiteCmd()
-	Use:   "website [url]",
-	Short: "Collect a single website",
-	Long:  `Collect a single website and store it in a DataNode.`,
-	Args:  cobra.ExactArgs(1),
-	RunE: func(cmd *cobra.Command, args []string) error {
-		websiteURL := args[0]
-		outputFile, _ := cmd.Flags().GetString("output")
-		depth, _ := cmd.Flags().GetInt("depth")
-		format, _ := cmd.Flags().GetString("format")
-		compression, _ := cmd.Flags().GetString("compression")
-
-		prompter := ui.NewNonInteractivePrompter(ui.GetWebsiteQuote)
-		prompter.Start()
-		defer prompter.Stop()
-		var bar *progressbar.ProgressBar
-		if prompter.IsInteractive() {
-			bar = ui.NewProgressBar(-1, "Crawling website")
-		}
-
-		dn, err := website.DownloadAndPackageWebsite(websiteURL, depth, bar)
-		if err != nil {
-			return fmt.Errorf("error downloading and packaging website: %w", err)
-		}
-
-		var data []byte
-		if format == "matrix" {
-			matrix, err := matrix.FromDataNode(dn)
-			if err != nil {
-				return fmt.Errorf("error creating matrix: %w", err)
-			}
-			data, err = matrix.ToTar()
-			if err != nil {
-				return fmt.Errorf("error serializing matrix: %w", err)
-			}
-		} else {
-			data, err = dn.ToTar()
-			if err != nil {
-				return fmt.Errorf("error serializing DataNode: %w", err)
-			}
-		}
-
-		compressedData, err := compress.Compress(data, compression)
-		if err != nil {
-			return fmt.Errorf("error compressing data: %w", err)
-		}
-
-		if outputFile == "" {
-			outputFile = "website." + format
-			if compression != "none" {
-				outputFile += "." + compression
-			}
-		}
-
-		err = os.WriteFile(outputFile, compressedData, 0644)
-		if err != nil {
-			return fmt.Errorf("error writing website to file: %w", err)
-		}
-
-		fmt.Fprintln(cmd.OutOrStdout(), "Website saved to", outputFile)
-		return nil
-	},
-}
 
 func init() {
-	collectCmd.AddCommand(collectWebsiteCmd)
+	GetCollectCmd().AddCommand(GetCollectWebsiteCmd())
+}
+
+func GetCollectWebsiteCmd() *cobra.Command {
+	return collectWebsiteCmd
+}
+
+func NewCollectWebsiteCmd() *cobra.Command {
+	collectWebsiteCmd := &cobra.Command{
+		Use:   "website [url]",
+		Short: "Collect a single website",
+		Long:  `Collect a single website and store it in a DataNode.`,
+		Args:  cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			websiteURL := args[0]
+			outputFile, _ := cmd.Flags().GetString("output")
+			depth, _ := cmd.Flags().GetInt("depth")
+			format, _ := cmd.Flags().GetString("format")
+			compression, _ := cmd.Flags().GetString("compression")
+
+			prompter := ui.NewNonInteractivePrompter(ui.GetWebsiteQuote)
+			prompter.Start()
+			defer prompter.Stop()
+			var bar *progressbar.ProgressBar
+			if prompter.IsInteractive() {
+				bar = ui.NewProgressBar(-1, "Crawling website")
+			}
+
+			dn, err := website.DownloadAndPackageWebsite(websiteURL, depth, bar)
+			if err != nil {
+				return fmt.Errorf("error downloading and packaging website: %w", err)
+			}
+
+			var data []byte
+			if format == "matrix" {
+				matrix, err := matrix.FromDataNode(dn)
+				if err != nil {
+					return fmt.Errorf("error creating matrix: %w", err)
+				}
+				data, err = matrix.ToTar()
+				if err != nil {
+					return fmt.Errorf("error serializing matrix: %w", err)
+				}
+			} else {
+				data, err = dn.ToTar()
+				if err != nil {
+					return fmt.Errorf("error serializing DataNode: %w", err)
+				}
+			}
+
+			compressedData, err := compress.Compress(data, compression)
+			if err != nil {
+				return fmt.Errorf("error compressing data: %w", err)
+			}
+
+			if outputFile == "" {
+				outputFile = "website." + format
+				if compression != "none" {
+					outputFile += "." + compression
+				}
+			}
+
+			err = os.WriteFile(outputFile, compressedData, 0644)
+			if err != nil {
+				return fmt.Errorf("error writing website to file: %w", err)
+			}
+
+			fmt.Fprintln(cmd.OutOrStdout(), "Website saved to", outputFile)
+			return nil
+		},
+	}
 	collectWebsiteCmd.PersistentFlags().String("output", "", "Output file for the DataNode")
 	collectWebsiteCmd.PersistentFlags().Int("depth", 2, "Recursion depth for downloading")
 	collectWebsiteCmd.PersistentFlags().String("format", "datanode", "Output format (datanode or matrix)")
 	collectWebsiteCmd.PersistentFlags().String("compression", "none", "Compression format (none, gz, or xz)")
-}
-func NewCollectWebsiteCmd() *cobra.Command {
 	return collectWebsiteCmd
 }

cmd/collect_website_test.go

@@ -11,27 +11,8 @@ import (
 	"github.com/schollz/progressbar/v3"
 )
 
-func TestCollectWebsiteCmd_NoArgs(t *testing.T) {
-	rootCmd := NewRootCmd()
-	collectCmd := NewCollectCmd()
-	collectWebsiteCmd := NewCollectWebsiteCmd()
-	collectCmd.AddCommand(collectWebsiteCmd)
-	rootCmd.AddCommand(collectCmd)
-	_, err := executeCommand(rootCmd, "collect", "website")
-	if err == nil {
-		t.Fatalf("expected an error, but got none")
-	}
-	if !strings.Contains(err.Error(), "accepts 1 arg(s), received 0") {
-		t.Fatalf("unexpected error message: %v", err)
-	}
-}
-func Test_NewCollectWebsiteCmd(t *testing.T) {
-	if NewCollectWebsiteCmd() == nil {
-		t.Errorf("NewCollectWebsiteCmd is nil")
-	}
-}
-
 func TestCollectWebsiteCmd_Good(t *testing.T) {
+	// Mock the website downloader
 	oldDownloadAndPackageWebsite := website.DownloadAndPackageWebsite
 	website.DownloadAndPackageWebsite = func(startURL string, maxDepth int, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
 		return datanode.New(), nil
@@ -41,8 +22,9 @@ func TestCollectWebsiteCmd_Good(t *testing.T) {
 	}()
 
 	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(collectCmd)
+	rootCmd.AddCommand(GetCollectCmd())
 
+	// Execute command
 	out := filepath.Join(t.TempDir(), "out")
 	_, err := executeCommand(rootCmd, "collect", "website", "https://example.com", "--output", out)
 	if err != nil {
@@ -51,6 +33,7 @@ func TestCollectWebsiteCmd_Good(t *testing.T) {
 }
 
 func TestCollectWebsiteCmd_Bad(t *testing.T) {
+	// Mock the website downloader to return an error
 	oldDownloadAndPackageWebsite := website.DownloadAndPackageWebsite
 	website.DownloadAndPackageWebsite = func(startURL string, maxDepth int, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
 		return nil, fmt.Errorf("website error")
@@ -60,11 +43,26 @@ func TestCollectWebsiteCmd_Bad(t *testing.T) {
 	}()
 
 	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(collectCmd)
+	rootCmd.AddCommand(GetCollectCmd())
 
+	// Execute command
 	out := filepath.Join(t.TempDir(), "out")
 	_, err := executeCommand(rootCmd, "collect", "website", "https://example.com", "--output", out)
 	if err == nil {
-		t.Fatalf("expected an error, but got none")
+		t.Fatal("expected an error, but got none")
 	}
 }
+
+func TestCollectWebsiteCmd_Ugly(t *testing.T) {
+	t.Run("No arguments", func(t *testing.T) {
+		rootCmd := NewRootCmd()
+		rootCmd.AddCommand(GetCollectCmd())
+		_, err := executeCommand(rootCmd, "collect", "website")
+		if err == nil {
+			t.Fatal("expected an error for no arguments, but got none")
+		}
+		if !strings.Contains(err.Error(), "accepts 1 arg(s), received 0") {
+			t.Errorf("unexpected error message: %v", err)
+		}
+	})
+}

cmd/compile.go

@@ -12,54 +12,63 @@ import (
 var borgfile string
 var output string
 
-var compileCmd = &cobra.Command{
+var compileCmd = NewCompileCmd()
-	Use:   "compile",
-	Short: "Compile a Borgfile into a Terminal Isolation Matrix.",
-	RunE: func(cmd *cobra.Command, args []string) error {
-		content, err := os.ReadFile(borgfile)
-		if err != nil {
-			return err
-		}
 
-		m, err := matrix.New()
+func NewCompileCmd() *cobra.Command {
-		if err != nil {
+	compileCmd := &cobra.Command{
-			return err
+		Use:   "compile",
-		}
+		Short: "Compile a Borgfile into a Terminal Isolation Matrix.",
-
+		RunE: func(cmd *cobra.Command, args []string) error {
-		lines := strings.Split(string(content), "\n")
+			content, err := os.ReadFile(borgfile)
-		for _, line := range lines {
+			if err != nil {
-			parts := strings.Fields(line)
+				return err
-			if len(parts) == 0 {
-				continue
 			}
-			switch parts[0] {
+
-			case "ADD":
+			m, err := matrix.New()
-				if len(parts) != 3 {
+			if err != nil {
-					return fmt.Errorf("invalid ADD instruction: %s", line)
+				return err
-				}
-				src := parts[1]
-				dest := parts[2]
-				data, err := os.ReadFile(src)
-				if err != nil {
-					return err
-				}
-				m.RootFS.AddData(dest, data)
-			default:
-				return fmt.Errorf("unknown instruction: %s", parts[0])
 			}
-		}
 
-		tarball, err := m.ToTar()
+			lines := strings.Split(string(content), "\n")
-		if err != nil {
+			for _, line := range lines {
-			return err
+				parts := strings.Fields(line)
-		}
+				if len(parts) == 0 {
+					continue
+				}
+				switch parts[0] {
+				case "ADD":
+					if len(parts) != 3 {
+						return fmt.Errorf("invalid ADD instruction: %s", line)
+					}
+					src := parts[1]
+					dest := parts[2]
+					data, err := os.ReadFile(src)
+					if err != nil {
+						return err
+					}
+					m.RootFS.AddData(strings.TrimPrefix(dest, "/"), data)
+				default:
+					return fmt.Errorf("unknown instruction: %s", parts[0])
+				}
+			}
 
-		return os.WriteFile(output, tarball, 0644)
+			tarball, err := m.ToTar()
-	},
+			if err != nil {
+				return err
+			}
+
+			return os.WriteFile(output, tarball, 0644)
+		},
+	}
+	compileCmd.Flags().StringVarP(&borgfile, "file", "f", "Borgfile", "Path to the Borgfile.")
+	compileCmd.Flags().StringVarP(&output, "output", "o", "a.matrix", "Path to the output matrix file.")
+	return compileCmd
+}
+
+func GetCompileCmd() *cobra.Command {
+	return compileCmd
 }
 
 func init() {
-	RootCmd.AddCommand(compileCmd)
+	RootCmd.AddCommand(GetCompileCmd())
-	compileCmd.Flags().StringVarP(&borgfile, "file", "f", "Borgfile", "Path to the Borgfile.")
-	compileCmd.Flags().StringVarP(&output, "output", "o", "a.matrix", "Path to the output matrix file.")
 }

cmd/compile_test.go

@@ -29,7 +29,7 @@ func TestCompileCmd_Good(t *testing.T) {
 
 	// Run the compile command.
 	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(compileCmd)
+	rootCmd.AddCommand(GetCompileCmd())
 	_, err = executeCommand(rootCmd, "compile", "-f", borgfilePath, "-o", outputMatrixPath)
 	if err != nil {
 		t.Fatalf("compile command failed: %v", err)
@@ -43,9 +43,7 @@ func TestCompileCmd_Good(t *testing.T) {
 	defer matrixFile.Close()
 
 	tr := tar.NewReader(matrixFile)
-	foundConfig := false
+	found := make(map[string]bool)
-	foundRootFS := false
-	foundTestFile := false
 	for {
 		header, err := tr.Next()
 		if err == io.EOF {
@@ -54,66 +52,79 @@ func TestCompileCmd_Good(t *testing.T) {
 		if err != nil {
 			t.Fatalf("failed to read tar header: %v", err)
 		}
+		found[header.Name] = true
+	}
 
-		switch header.Name {
+	expectedFiles := []string{"config.json", "rootfs/", "rootfs/test.txt"}
-		case "config.json":
+	for _, f := range expectedFiles {
-			foundConfig = true
+		if !found[f] {
-		case "rootfs/":
+			t.Errorf("%s not found in matrix tarball", f)
-			foundRootFS = true
-		case "rootfs/test.txt":
-			foundTestFile = true
 		}
 	}
-
-	if !foundConfig {
-		t.Error("config.json not found in matrix")
-	}
-	if !foundRootFS {
-		t.Error("rootfs/ not found in matrix")
-	}
-	if !foundTestFile {
-		t.Error("rootfs/test.txt not found in matrix")
-	}
 }
 
-func TestCompileCmd_Bad_InvalidBorgfile(t *testing.T) {
+func TestCompileCmd_Bad(t *testing.T) {
-	tempDir := t.TempDir()
+	t.Run("Invalid Borgfile instruction", func(t *testing.T) {
-	borgfilePath := filepath.Join(tempDir, "Borgfile")
+		tempDir := t.TempDir()
-	outputMatrixPath := filepath.Join(tempDir, "test.matrix")
+		borgfilePath := filepath.Join(tempDir, "Borgfile")
+		outputMatrixPath := filepath.Join(tempDir, "test.matrix")
 
-	// Create a dummy Borgfile with an invalid instruction.
+		// Create a dummy Borgfile with an invalid instruction.
-	borgfileContent := "INVALID_INSTRUCTION"
+		borgfileContent := "INVALID_INSTRUCTION"
-	err := os.WriteFile(borgfilePath, []byte(borgfileContent), 0644)
+		err := os.WriteFile(borgfilePath, []byte(borgfileContent), 0644)
-	if err != nil {
+		if err != nil {
-		t.Fatalf("failed to create Borgfile: %v", err)
+			t.Fatalf("failed to create Borgfile: %v", err)
-	}
+		}
 
-	// Run the compile command.
+		// Run the compile command.
-	rootCmd := NewRootCmd()
+		rootCmd := NewRootCmd()
-	rootCmd.AddCommand(compileCmd)
+		rootCmd.AddCommand(GetCompileCmd())
-	_, err = executeCommand(rootCmd, "compile", "-f", borgfilePath, "-o", outputMatrixPath)
+		_, err = executeCommand(rootCmd, "compile", "-f", borgfilePath, "-o", outputMatrixPath)
-	if err == nil {
+		if err == nil {
-		t.Fatal("compile command should have failed but did not")
+			t.Fatal("compile command should have failed but did not")
-	}
+		}
+	})
+
+	t.Run("Missing input file", func(t *testing.T) {
+		tempDir := t.TempDir()
+		borgfilePath := filepath.Join(tempDir, "Borgfile")
+		outputMatrixPath := filepath.Join(tempDir, "test.matrix")
+
+		// Create a dummy Borgfile that references a non-existent file.
+		borgfileContent := "ADD /non/existent/file /test.txt"
+		err := os.WriteFile(borgfilePath, []byte(borgfileContent), 0644)
+		if err != nil {
+			t.Fatalf("failed to create Borgfile: %v", err)
+		}
+
+		// Run the compile command.
+		rootCmd := NewRootCmd()
+		rootCmd.AddCommand(GetCompileCmd())
+		_, err = executeCommand(rootCmd, "compile", "-f", borgfilePath, "-o", outputMatrixPath)
+		if err == nil {
+			t.Fatal("compile command should have failed but did not")
+		}
+	})
 }
 
-func TestCompileCmd_Bad_MissingInputFile(t *testing.T) {
+func TestCompileCmd_Ugly(t *testing.T) {
-	tempDir := t.TempDir()
+	t.Run("Empty Borgfile", func(t *testing.T) {
-	borgfilePath := filepath.Join(tempDir, "Borgfile")
+		tempDir := t.TempDir()
-	outputMatrixPath := filepath.Join(tempDir, "test.matrix")
+		borgfilePath := filepath.Join(tempDir, "Borgfile")
+		outputMatrixPath := filepath.Join(tempDir, "test.matrix")
 
-	// Create a dummy Borgfile that references a non-existent file.
+		// Create an empty Borgfile.
-	borgfileContent := "ADD /non/existent/file /test.txt"
+		err := os.WriteFile(borgfilePath, []byte(""), 0644)
-	err := os.WriteFile(borgfilePath, []byte(borgfileContent), 0644)
+		if err != nil {
-	if err != nil {
+			t.Fatalf("failed to create Borgfile: %v", err)
-		t.Fatalf("failed to create Borgfile: %v", err)
+		}
-	}
 
-	// Run the compile command.
+		// Run the compile command.
-	rootCmd := NewRootCmd()
+		rootCmd := NewRootCmd()
-	rootCmd.AddCommand(compileCmd)
+		rootCmd.AddCommand(GetCompileCmd())
-	_, err = executeCommand(rootCmd, "compile", "-f", borgfilePath, "-o", outputMatrixPath)
+		_, err = executeCommand(rootCmd, "compile", "-f", borgfilePath, "-o", outputMatrixPath)
-	if err == nil {
+		if err != nil {
-		t.Fatal("compile command should have failed but did not")
+			t.Fatalf("compile command failed for empty Borgfile: %v", err)
-	}
+		}
+	})
 }

cmd/main_test.go

@@ -1,25 +0,0 @@
-package cmd
-
-import (
-	"bytes"
-
-	"github.com/spf13/cobra"
-)
-
-// executeCommand is a helper function to execute a cobra command and return the output.
-func executeCommand(root *cobra.Command, args ...string) (string, error) {
-	_, output, err := executeCommandC(root, args...)
-	return output, err
-}
-
-// executeCommandC is a helper function to execute a cobra command and return the output.
-func executeCommandC(root *cobra.Command, args ...string) (*cobra.Command, string, error) {
-	buf := new(bytes.Buffer)
-	root.SetOut(buf)
-	root.SetErr(buf)
-	root.SetArgs(args)
-
-	c, err := root.ExecuteC()
-
-	return c, buf.String(), err
-}

cmd/root_test.go

@@ -1,54 +1,84 @@
 package cmd
 
 import (
+	"bytes"
 	"io"
 	"log/slog"
+	"strings"
 	"testing"
 
 	"github.com/spf13/cobra"
 )
 
-func TestExecute(t *testing.T) {
+// executeCommand is a helper function to execute a cobra command and return the output.
+func executeCommand(root *cobra.Command, args ...string) (string, error) {
+	_, output, err := executeCommandC(root, args...)
+	return output, err
+}
+
+// executeCommandC is a helper function to execute a cobra command and return the output.
+func executeCommandC(root *cobra.Command, args ...string) (*cobra.Command, string, error) {
+	buf := new(bytes.Buffer)
+	root.SetOut(buf)
+	root.SetErr(buf)
+	root.SetArgs(args)
+
+	c, err := root.ExecuteC()
+
+	return c, buf.String(), err
+}
+
+func TestExecute_Good(t *testing.T) {
+	// This is a basic test to ensure the command runs without panicking.
 	err := Execute(slog.New(slog.NewTextHandler(io.Discard, nil)))
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 }
 
+func TestRootCmd_Good(t *testing.T) {
+	t.Run("No args", func(t *testing.T) {
+		_, err := executeCommand(RootCmd)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+	})
 
-func Test_NewRootCmd(t *testing.T) {
+	t.Run("Help flag", func(t *testing.T) {
-	if NewRootCmd() == nil {
+		// We need to reset the command's state before each run.
-		t.Errorf("NewRootCmd is nil")
+		RootCmd.ResetFlags()
-	}
+		RootCmd.ResetCommands()
+		initAllCommands()
+
+		output, err := executeCommand(RootCmd, "--help")
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if !strings.Contains(output, "Usage:") {
+			t.Errorf("expected help output to contain 'Usage:', but it did not")
+		}
+	})
 }
-func Test_executeCommand(t *testing.T) {
+
-	type args struct {
+func TestRootCmd_Bad(t *testing.T) {
-		cmd  *cobra.Command
+	t.Run("Unknown command", func(t *testing.T) {
-		args []string
+		// We need to reset the command's state before each run.
-	}
+		RootCmd.ResetFlags()
-	tests := []struct {
+		RootCmd.ResetCommands()
-		name    string
+		initAllCommands()
-		args    args
+
-		want    string
+		_, err := executeCommand(RootCmd, "unknown-command")
-		wantErr bool
+		if err == nil {
-	}{
+			t.Fatal("expected an error for an unknown command, but got none")
-		{
+		}
-			name: "Test with no args",
+	})
-			args: args{
+}
-				cmd:  NewRootCmd(),
+
-				args: []string{},
+// initAllCommands re-initializes all commands for testing.
-			},
+func initAllCommands() {
-			want:    "",
+	RootCmd.AddCommand(GetAllCmd())
-			wantErr: false,
+	RootCmd.AddCommand(GetCollectCmd())
-		},
+	RootCmd.AddCommand(GetCompileCmd())
-	}
+	RootCmd.AddCommand(GetRunCmd())
-	for _, tt := range tests {
+	RootCmd.AddCommand(GetServeCmd())
-		t.Run(tt.name, func(t *testing.T) {
-			_, err := executeCommand(tt.args.cmd, tt.args.args...)
-			if (err != nil) != tt.wantErr {
-				t.Errorf("executeCommand() error = %v, wantErr %v", err, tt.wantErr)
-				return
-			}
-		})
-	}
 }

cmd/run.go

@@ -9,65 +9,73 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var runCmd = &cobra.Command{
+var runCmd = NewRunCmd()
-	Use:   "run [matrix file]",
-	Short: "Run a Terminal Isolation Matrix.",
-	Args:  cobra.ExactArgs(1),
-	RunE: func(cmd *cobra.Command, args []string) error {
-		matrixFile := args[0]
 
-		// Create a temporary directory to unpack the matrix file.
+func NewRunCmd() *cobra.Command {
-		tempDir, err := os.MkdirTemp("", "borg-run-*")
+	return &cobra.Command{
-		if err != nil {
+		Use:   "run [matrix file]",
-			return err
+		Short: "Run a Terminal Isolation Matrix.",
-		}
+		Args:  cobra.ExactArgs(1),
-		defer os.RemoveAll(tempDir)
+		RunE: func(cmd *cobra.Command, args []string) error {
+			matrixFile := args[0]
 
-		// Unpack the matrix file.
+			// Create a temporary directory to unpack the matrix file.
-		file, err := os.Open(matrixFile)
+			tempDir, err := os.MkdirTemp("", "borg-run-*")
-		if err != nil {
-			return err
-		}
-		defer file.Close()
-
-		tr := tar.NewReader(file)
-		for {
-			header, err := tr.Next()
-			if err == io.EOF {
-				break
-			}
 			if err != nil {
 				return err
 			}
+			defer os.RemoveAll(tempDir)
 
-			path := filepath.Join(tempDir, header.Name)
+			// Unpack the matrix file.
-			if header.Typeflag == tar.TypeDir {
+			file, err := os.Open(matrixFile)
-				if err := os.MkdirAll(path, 0755); err != nil {
+			if err != nil {
+				return err
+			}
+			defer file.Close()
+
+			tr := tar.NewReader(file)
+			for {
+				header, err := tr.Next()
+				if err == io.EOF {
+					break
+				}
+				if err != nil {
+					return err
+				}
+
+				path := filepath.Join(tempDir, header.Name)
+				if header.Typeflag == tar.TypeDir {
+					if err := os.MkdirAll(path, 0755); err != nil {
+						return err
+					}
+					continue
+				}
+
+				outFile, err := os.Create(path)
+				if err != nil {
+					return err
+				}
+				defer outFile.Close()
+				if _, err := io.Copy(outFile, tr); err != nil {
 					return err
 				}
-				continue
 			}
 
-			outFile, err := os.Create(path)
+			// Run the matrix.
-			if err != nil {
+			runc := execCommand("runc", "run", "borg-container")
-				return err
+			runc.Dir = tempDir
-			}
+			runc.Stdout = os.Stdout
-			defer outFile.Close()
+			runc.Stderr = os.Stderr
-			if _, err := io.Copy(outFile, tr); err != nil {
+			runc.Stdin = os.Stdin
-				return err
+			return runc.Run()
-			}
+		},
-		}
+	}
+}
 
-		// Run the matrix.
+func GetRunCmd() *cobra.Command {
-		runc := execCommand("runc", "run", "borg-container")
+	return runCmd
-		runc.Dir = tempDir
-		runc.Stdout = os.Stdout
-		runc.Stderr = os.Stderr
-		runc.Stdin = os.Stdin
-		return runc.Run()
-	},
 }
 
 func init() {
-	RootCmd.AddCommand(runCmd)
+	RootCmd.AddCommand(GetRunCmd())
 }

cmd/run_test.go

@@ -27,6 +27,59 @@ func TestHelperProcess(t *testing.T) {
 
 func TestRunCmd_Good(t *testing.T) {
 	// Create a dummy matrix file.
+	matrixPath := createDummyMatrix(t)
+
+	// Mock the exec.Command function.
+	origExecCommand := execCommand
+	execCommand = helperProcess
+	t.Cleanup(func() {
+		execCommand = origExecCommand
+	})
+
+	// Run the run command.
+	rootCmd := NewRootCmd()
+	rootCmd.AddCommand(GetRunCmd())
+	_, err := executeCommand(rootCmd, "run", matrixPath)
+	if err != nil {
+		t.Fatalf("run command failed: %v", err)
+	}
+}
+
+func TestRunCmd_Bad(t *testing.T) {
+	t.Run("Missing input file", func(t *testing.T) {
+		// Run the run command with a non-existent file.
+		rootCmd := NewRootCmd()
+		rootCmd.AddCommand(GetRunCmd())
+		_, err := executeCommand(rootCmd, "run", "/non/existent/file.matrix")
+		if err == nil {
+			t.Fatal("run command should have failed but did not")
+		}
+	})
+}
+
+func TestRunCmd_Ugly(t *testing.T) {
+	t.Run("Invalid matrix file", func(t *testing.T) {
+		// Create an invalid (non-tar) matrix file.
+		tempDir := t.TempDir()
+		matrixPath := filepath.Join(tempDir, "invalid.matrix")
+		err := os.WriteFile(matrixPath, []byte("this is not a tar file"), 0644)
+		if err != nil {
+			t.Fatalf("failed to create invalid matrix file: %v", err)
+		}
+
+		// Run the run command.
+		rootCmd := NewRootCmd()
+		rootCmd.AddCommand(GetRunCmd())
+		_, err = executeCommand(rootCmd, "run", matrixPath)
+		if err == nil {
+			t.Fatal("run command should have failed but did not")
+		}
+	})
+}
+
+// createDummyMatrix creates a valid, empty matrix file for testing.
+func createDummyMatrix(t *testing.T) string {
+	t.Helper()
 	tempDir := t.TempDir()
 	matrixPath := filepath.Join(tempDir, "test.matrix")
 	matrixFile, err := os.Create(matrixPath)
@@ -36,6 +89,7 @@ func TestRunCmd_Good(t *testing.T) {
 	defer matrixFile.Close()
 
 	tw := tar.NewWriter(matrixFile)
+
 	// Add a dummy config.json.
 	configContent := []byte(matrix.DefaultConfigJSON)
 	hdr := &tar.Header{
@@ -63,29 +117,5 @@ func TestRunCmd_Good(t *testing.T) {
 	if err := tw.Close(); err != nil {
 		t.Fatalf("failed to close tar writer: %v", err)
 	}
-
+	return matrixPath
-	// Mock the exec.Command function.
-	origExecCommand := execCommand
-	execCommand = helperProcess
-	t.Cleanup(func() {
-		execCommand = origExecCommand
-	})
-
-	// Run the run command.
-	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(runCmd)
-	_, err = executeCommand(rootCmd, "run", matrixPath)
-	if err != nil {
-		t.Fatalf("run command failed: %v", err)
-	}
-}
-
-func TestRunCmd_Bad_MissingInputFile(t *testing.T) {
-	// Run the run command with a non-existent file.
-	rootCmd := NewRootCmd()
-	rootCmd.AddCommand(runCmd)
-	_, err := executeCommand(rootCmd, "run", "/non/existent/file.matrix")
-	if err == nil {
-		t.Fatal("run command should have failed but did not")
-	}
 }

cmd/serve.go

@@ -14,51 +14,60 @@ import (
 )
 
 // serveCmd represents the serve command
-var serveCmd = &cobra.Command{
+var serveCmd = NewServeCmd()
-	Use:   "serve [file]",
-	Short: "Serve a packaged PWA file",
-	Long:  `Serves the contents of a packaged PWA file using a static file server.`,
-	Args:  cobra.ExactArgs(1),
-	RunE: func(cmd *cobra.Command, args []string) error {
-		dataFile := args[0]
-		port, _ := cmd.Flags().GetString("port")
 
-		rawData, err := os.ReadFile(dataFile)
+func NewServeCmd() *cobra.Command {
-		if err != nil {
+	serveCmd := &cobra.Command{
-			return fmt.Errorf("Error reading data file: %w", err)
+		Use:   "serve [file]",
-		}
+		Short: "Serve a packaged PWA file",
+		Long:  `Serves the contents of a packaged PWA file using a static file server.`,
+		Args:  cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			dataFile := args[0]
+			port, _ := cmd.Flags().GetString("port")
 
-		data, err := compress.Decompress(rawData)
+			rawData, err := os.ReadFile(dataFile)
-		if err != nil {
-			return fmt.Errorf("Error decompressing data: %w", err)
-		}
-
-		var fs http.FileSystem
-		if strings.HasSuffix(dataFile, ".matrix") {
-			fs, err = tarfs.New(data)
 			if err != nil {
-				return fmt.Errorf("Error creating TarFS from matrix tarball: %w", err)
+				return fmt.Errorf("Error reading data file: %w", err)
 			}
-		} else {
+
-			dn, err := datanode.FromTar(data)
+			data, err := compress.Decompress(rawData)
 			if err != nil {
-				return fmt.Errorf("Error creating DataNode from tarball: %w", err)
+				return fmt.Errorf("Error decompressing data: %w", err)
 			}
-			fs = http.FS(dn)
-		}
 
-		http.Handle("/", http.FileServer(fs))
+			var fs http.FileSystem
+			if strings.HasSuffix(dataFile, ".matrix") {
+				fs, err = tarfs.New(data)
+				if err != nil {
+					return fmt.Errorf("Error creating TarFS from matrix tarball: %w", err)
+				}
+			} else {
+				dn, err := datanode.FromTar(data)
+				if err != nil {
+					return fmt.Errorf("Error creating DataNode from tarball: %w", err)
+				}
+				fs = http.FS(dn)
+			}
 
-		fmt.Printf("Serving PWA on http://localhost:%s\n", port)
+			http.Handle("/", http.FileServer(fs))
-		err = http.ListenAndServe(":"+port, nil)
+
-		if err != nil {
+			fmt.Printf("Serving PWA on http://localhost:%s\n", port)
-			return fmt.Errorf("Error starting server: %w", err)
+			err = http.ListenAndServe(":"+port, nil)
-		}
+			if err != nil {
-		return nil
+				return fmt.Errorf("Error starting server: %w", err)
-	},
+			}
+			return nil
+		},
+	}
+	serveCmd.PersistentFlags().String("port", "8080", "Port to serve the PWA on")
+	return serveCmd
+}
+
+func GetServeCmd() *cobra.Command {
+	return serveCmd
 }
 
 func init() {
-	RootCmd.AddCommand(serveCmd)
+	RootCmd.AddCommand(GetServeCmd())
-	serveCmd.PersistentFlags().String("port", "8080", "Port to serve the PWA on")
 }

pkg/compress/compress_test.go

@@ -5,55 +5,115 @@ import (
 	"testing"
 )
 
-func TestCompressDecompress(t *testing.T) {
+func TestGzip_Good(t *testing.T) {
-	testData := []byte("hello, world")
+	originalData := []byte("hello, gzip world")
-
+	compressed, err := Compress(originalData, "gz")
-	// Test gzip compression
-	compressedGz, err := Compress(testData, "gz")
 	if err != nil {
 		t.Fatalf("gzip compression failed: %v", err)
 	}
+	if bytes.Equal(originalData, compressed) {
+		t.Fatal("gzip compressed data is the same as the original")
+	}
 
-	decompressedGz, err := Decompress(compressedGz)
+	decompressed, err := Decompress(compressed)
 	if err != nil {
 		t.Fatalf("gzip decompression failed: %v", err)
 	}
 
-	if !bytes.Equal(testData, decompressedGz) {
+	if !bytes.Equal(originalData, decompressed) {
 		t.Errorf("gzip decompressed data does not match original data")
 	}
+}
 
-	// Test xz compression
+func TestXz_Good(t *testing.T) {
-	compressedXz, err := Compress(testData, "xz")
+	originalData := []byte("hello, xz world")
+	compressed, err := Compress(originalData, "xz")
 	if err != nil {
 		t.Fatalf("xz compression failed: %v", err)
 	}
+	if bytes.Equal(originalData, compressed) {
+		t.Fatal("xz compressed data is the same as the original")
+	}
 
-	decompressedXz, err := Decompress(compressedXz)
+	decompressed, err := Decompress(compressed)
 	if err != nil {
 		t.Fatalf("xz decompression failed: %v", err)
 	}
 
-	if !bytes.Equal(testData, decompressedXz) {
+	if !bytes.Equal(originalData, decompressed) {
 		t.Errorf("xz decompressed data does not match original data")
 	}
+}
 
-	// Test no compression
+func TestNone_Good(t *testing.T) {
-	compressedNone, err := Compress(testData, "none")
+	originalData := []byte("hello, none world")
+	compressed, err := Compress(originalData, "none")
 	if err != nil {
-		t.Fatalf("no compression failed: %v", err)
+		t.Fatalf("'none' compression failed: %v", err)
+	}
+	if !bytes.Equal(originalData, compressed) {
+		t.Errorf("'none' compression should not change data")
 	}
 
-	if !bytes.Equal(testData, compressedNone) {
+	decompressed, err := Decompress(compressed)
-		t.Errorf("no compression data does not match original data")
-	}
-
-	decompressedNone, err := Decompress(compressedNone)
 	if err != nil {
-		t.Fatalf("no compression decompression failed: %v", err)
+		t.Fatalf("'none' decompression failed: %v", err)
 	}
 
-	if !bytes.Equal(testData, decompressedNone) {
+	if !bytes.Equal(originalData, decompressed) {
-		t.Errorf("no compression decompressed data does not match original data")
+		t.Errorf("'none' decompressed data does not match original data")
+	}
+}
+
+func TestCompress_Bad(t *testing.T) {
+	originalData := []byte("test")
+	// The function should return the original data for an unknown format.
+	compressed, err := Compress(originalData, "invalid-format")
+	if err != nil {
+		t.Fatalf("expected no error for invalid compression format, got %v", err)
+	}
+	if !bytes.Equal(originalData, compressed) {
+		t.Errorf("expected original data for unknown format, got %q", compressed)
+	}
+}
+
+func TestDecompress_Bad(t *testing.T) {
+	// A truncated gzip stream should cause a decompression error.
+	originalData := []byte("hello, gzip world")
+	compressed, _ := Compress(originalData, "gz")
+	truncated := compressed[:len(compressed)-5] // Corrupt the stream
+
+	_, err := Decompress(truncated)
+	if err == nil {
+		t.Fatal("expected an error when decompressing a truncated stream, got nil")
+	}
+}
+
+func TestCompress_Ugly(t *testing.T) {
+	// Test compressing empty data
+	originalData := []byte{}
+	compressed, err := Compress(originalData, "gz")
+	if err != nil {
+		t.Fatalf("compressing empty data failed: %v", err)
+	}
+
+	decompressed, err := Decompress(compressed)
+	if err != nil {
+		t.Fatalf("decompressing empty compressed data failed: %v", err)
+	}
+
+	if !bytes.Equal(originalData, decompressed) {
+		t.Errorf("expected empty data, got %q", decompressed)
+	}
+}
+
+func TestDecompress_Ugly(t *testing.T) {
+	// Test decompressing empty byte slice
+	result, err := Decompress([]byte{})
+	if err != nil {
+		t.Fatalf("decompressing an empty slice should not produce an error, got %v", err)
+	}
+	if len(result) != 0 {
+		t.Errorf("expected empty result from decompressing empty slice, got %q", result)
 	}
 }

pkg/datanode/datanode.go

@@ -119,6 +119,11 @@ func (d *DataNode) ReadDir(name string) ([]fs.DirEntry, error) {
 		name = ""
 	}
 
+	// Disallow reading a file as a directory.
+	if info, err := d.Stat(name); err == nil && !info.IsDir() {
+		return nil, &fs.PathError{Op: "readdir", Path: name, Err: fs.ErrInvalid}
+	}
+
 	entries := []fs.DirEntry{}
 	seen := make(map[string]bool)
 

pkg/datanode/datanode_test.go

@@ -1,96 +1,298 @@
 package datanode
 
 import (
+	"errors"
 	"io/fs"
 	"os"
+	"path/filepath"
 	"reflect"
 	"sort"
+	"strings"
 	"testing"
 )
 
-func TestDataNode(t *testing.T) {
+func TestNew_Good(t *testing.T) {
+	dn := New()
+	if dn == nil {
+		t.Fatal("New() returned nil")
+	}
+	if dn.files == nil {
+		t.Error("New() did not initialize the files map")
+	}
+}
+
+func TestAddData_Good(t *testing.T) {
+	dn := New()
+	path := "foo.txt"
+	data := []byte("foo")
+	dn.AddData(path, data)
+
+	file, ok := dn.files[path]
+	if !ok {
+		t.Fatalf("file %q not found in datanode", path)
+	}
+	if string(file.content) != string(data) {
+		t.Errorf("expected data %q, got %q", data, file.content)
+	}
+	info, err := file.Stat()
+	if err != nil {
+		t.Fatalf("file.Stat() failed: %v", err)
+	}
+	if info.Name() != "foo.txt" {
+		t.Errorf("expected name foo.txt, got %s", info.Name())
+	}
+}
+
+func TestAddData_Ugly(t *testing.T) {
+	t.Run("Overwrite", func(t *testing.T) {
+		dn := New()
+		dn.AddData("foo.txt", []byte("foo"))
+		dn.AddData("foo.txt", []byte("bar"))
+
+		file, _ := dn.files["foo.txt"]
+		if string(file.content) != "bar" {
+			t.Errorf("expected data to be overwritten to 'bar', got %q", file.content)
+		}
+	})
+
+	t.Run("Weird Path", func(t *testing.T) {
+		dn := New()
+		// path.Clean treats "a/../b/./c.txt" as "b/c.txt" but our implementation is simpler
+		// and doesn't handle `..`. Let's test what it does handle.
+		path := "./b/./c.txt"
+		dn.AddData(path, []byte("c"))
+		if _, ok := dn.files["./b/./c.txt"]; !ok {
+			t.Errorf("expected path to be stored as is")
+		}
+	})
+}
+
+func TestOpen_Good(t *testing.T) {
 	dn := New()
 	dn.AddData("foo.txt", []byte("foo"))
-	dn.AddData("bar/baz.txt", []byte("baz"))
-	dn.AddData("bar/qux.txt", []byte("qux"))
-
-	// Test Open
 	file, err := dn.Open("foo.txt")
 	if err != nil {
 		t.Fatalf("Open failed: %v", err)
 	}
-	file.Close()
+	defer file.Close()
+}
 
-	_, err = dn.Open("nonexistent.txt")
+func TestOpen_Bad(t *testing.T) {
+	dn := New()
+	_, err := dn.Open("nonexistent.txt")
 	if err == nil {
-		t.Fatalf("Expected error opening nonexistent file, got nil")
+		t.Fatal("expected error opening nonexistent file, got nil")
 	}
+	if !errors.Is(err, fs.ErrNotExist) {
+		t.Errorf("expected fs.ErrNotExist, got %v", err)
+	}
+}
 
-	// Test Stat
+func TestOpen_Ugly(t *testing.T) {
+	dn := New()
+	dn.AddData("bar/baz.txt", []byte("baz"))
+	file, err := dn.Open("bar") // Opening a directory
+	if err != nil {
+		t.Fatalf("expected no error when opening a directory, got %v", err)
+	}
+	defer file.Close()
+
+	// Reading from a directory should fail
+	_, err = file.Read(make([]byte, 1))
+	if err == nil {
+		t.Fatal("expected error reading from a directory, got nil")
+	}
+	var pathErr *fs.PathError
+	if !errors.As(err, &pathErr) || pathErr.Err != fs.ErrInvalid {
+		t.Errorf("expected fs.ErrInvalid when reading a directory, got %v", err)
+	}
+}
+
+func TestStat_Good(t *testing.T) {
+	dn := New()
+	dn.AddData("foo.txt", []byte("foo"))
+	dn.AddData("bar/baz.txt", []byte("baz"))
+
+	// Test file
 	info, err := dn.Stat("bar/baz.txt")
 	if err != nil {
 		t.Fatalf("Stat failed: %v", err)
 	}
 	if info.Name() != "baz.txt" {
-		t.Errorf("Expected name baz.txt, got %s", info.Name())
+		t.Errorf("expected name baz.txt, got %s", info.Name())
 	}
 	if info.Size() != 3 {
-		t.Errorf("Expected size 3, got %d", info.Size())
+		t.Errorf("expected size 3, got %d", info.Size())
 	}
 	if info.IsDir() {
-		t.Errorf("Expected baz.txt to not be a directory")
+		t.Error("expected baz.txt to not be a directory")
 	}
 
+	// Test directory
 	dirInfo, err := dn.Stat("bar")
 	if err != nil {
 		t.Fatalf("Stat directory failed: %v", err)
 	}
 	if !dirInfo.IsDir() {
-		t.Errorf("Expected 'bar' to be a directory")
+		t.Error("expected 'bar' to be a directory")
 	}
+	if dirInfo.Name() != "bar" {
+		t.Errorf("expected dir name 'bar', got %s", dirInfo.Name())
+	}
+}
+
+func TestStat_Bad(t *testing.T) {
+	dn := New()
+	_, err := dn.Stat("nonexistent")
+	if err == nil {
+		t.Fatal("expected error stating nonexistent file, got nil")
+	}
+	if !errors.Is(err, fs.ErrNotExist) {
+		t.Errorf("expected fs.ErrNotExist, got %v", err)
+	}
+}
+
+func TestStat_Ugly(t *testing.T) {
+	dn := New()
+	dn.AddData("foo.txt", []byte("foo"))
+
+	// Test root
+	info, err := dn.Stat(".")
+	if err != nil {
+		t.Fatalf("Stat('.') failed: %v", err)
+	}
+	if !info.IsDir() {
+		t.Error("expected '.' to be a directory")
+	}
+	if info.Name() != "." {
+		t.Errorf("expected name '.', got %s", info.Name())
+	}
+}
+
+func TestExists_Good(t *testing.T) {
+	dn := New()
+	dn.AddData("foo.txt", []byte("foo"))
+	dn.AddData("bar/baz.txt", []byte("baz"))
 
-	// Test Exists
 	exists, err := dn.Exists("foo.txt")
 	if err != nil || !exists {
-		t.Errorf("Expected foo.txt to exist, err: %v", err)
+		t.Errorf("expected foo.txt to exist, err: %v", err)
-	}
-	exists, err = dn.Exists("bar")
-	if err != nil || !exists {
-		t.Errorf("Expected 'bar' directory to exist, err: %v", err)
-	}
-	exists, err = dn.Exists("nonexistent")
-	if err != nil || exists {
-		t.Errorf("Expected 'nonexistent' to not exist, err: %v", err)
 	}
 
-	// Test ReadDir
+	exists, err = dn.Exists("bar")
+	if err != nil || !exists {
+		t.Errorf("expected 'bar' directory to exist, err: %v", err)
+	}
+}
+
+func TestExists_Bad(t *testing.T) {
+	dn := New()
+	exists, err := dn.Exists("nonexistent")
+	if err != nil {
+		t.Errorf("unexpected error for nonexistent file: %v", err)
+	}
+	if exists {
+		t.Error("expected 'nonexistent' to not exist")
+	}
+}
+
+func TestExists_Ugly(t *testing.T) {
+	dn := New()
+	dn.AddData("dummy.txt", []byte("dummy"))
+	// Test root
+	exists, err := dn.Exists(".")
+	if err != nil || !exists {
+		t.Error("expected root '.' to exist")
+	}
+	// Test empty path
+	exists, err = dn.Exists("")
+	if err != nil {
+		// our stat treats "" as "."
+		if !strings.Contains(err.Error(), "exists") {
+			t.Errorf("unexpected error for empty path: %v", err)
+		}
+	}
+	if !exists {
+		t.Error("expected empty path '' to exist (as root)")
+	}
+}
+
+func TestReadDir_Good(t *testing.T) {
+	dn := New()
+	dn.AddData("foo.txt", []byte("foo"))
+	dn.AddData("bar/baz.txt", []byte("baz"))
+	dn.AddData("bar/qux.txt", []byte("qux"))
+
+	// Read root
 	entries, err := dn.ReadDir(".")
 	if err != nil {
 		t.Fatalf("ReadDir failed: %v", err)
 	}
 	expectedRootEntries := []string{"bar", "foo.txt"}
-	if len(entries) != len(expectedRootEntries) {
+	entryNames := toSortedNames(entries)
-		t.Errorf("Expected %d entries in root, got %d", len(expectedRootEntries), len(entries))
+	if !reflect.DeepEqual(entryNames, expectedRootEntries) {
-	}
+		t.Errorf("expected entries %v, got %v", expectedRootEntries, entryNames)
-	var rootEntryNames []string
-	for _, e := range entries {
-		rootEntryNames = append(rootEntryNames, e.Name())
-	}
-	sort.Strings(rootEntryNames)
-	if !reflect.DeepEqual(rootEntryNames, expectedRootEntries) {
-		t.Errorf("Expected entries %v, got %v", expectedRootEntries, rootEntryNames)
 	}
 
+	// Read subdirectory
 	barEntries, err := dn.ReadDir("bar")
 	if err != nil {
 		t.Fatalf("ReadDir('bar') failed: %v", err)
 	}
 	expectedBarEntries := []string{"baz.txt", "qux.txt"}
-	if len(barEntries) != len(expectedBarEntries) {
+	barEntryNames := toSortedNames(barEntries)
-		t.Errorf("Expected %d entries in 'bar', got %d", len(expectedBarEntries), len(barEntries))
+	if !reflect.DeepEqual(barEntryNames, expectedBarEntries) {
+		t.Errorf("expected entries %v, got %v", expectedBarEntries, barEntryNames)
+	}
+}
+
+func TestReadDir_Bad(t *testing.T) {
+	dn := New()
+	dn.AddData("foo.txt", []byte("foo"))
+
+	// Read nonexistent dir
+	entries, err := dn.ReadDir("nonexistent")
+	if err != nil {
+		t.Fatalf("expected no error reading nonexistent dir, got %v", err)
+	}
+	if len(entries) != 0 {
+		t.Errorf("expected 0 entries for nonexistent dir, got %d", len(entries))
 	}
 
-	// Test Walk
+	// Read file
+	_, err = dn.ReadDir("foo.txt")
+	if err == nil {
+		t.Fatal("expected error reading a file")
+	}
+	var pathErr *fs.PathError
+	if !errors.As(err, &pathErr) || pathErr.Err != fs.ErrInvalid {
+		t.Errorf("expected fs.ErrInvalid when reading a file, got %v", err)
+	}
+}
+
+func TestReadDir_Ugly(t *testing.T) {
+	dn := New()
+	dn.AddData("bar/baz.txt", []byte("baz"))
+	dn.AddData("empty_dir/", nil)
+
+	// Read dir with another dir but no files
+	entries, err := dn.ReadDir(".")
+	if err != nil {
+		t.Fatalf("ReadDir failed: %v", err)
+	}
+	expected := []string{"bar"} // empty_dir/ is ignored by AddData
+	names := toSortedNames(entries)
+	if !reflect.DeepEqual(names, expected) {
+		t.Errorf("expected %v, got %v", expected, names)
+	}
+}
+
+func TestWalk_Good(t *testing.T) {
+	dn := New()
+	dn.AddData("foo.txt", []byte("foo"))
+	dn.AddData("bar/baz.txt", []byte("baz"))
+	dn.AddData("bar/qux.txt", []byte("qux"))
+
 	var paths []string
 	dn.Walk(".", func(path string, d fs.DirEntry, err error) error {
 		paths = append(paths, path)
@@ -101,24 +303,105 @@ func TestDataNode(t *testing.T) {
 	if !reflect.DeepEqual(paths, expectedPaths) {
 		t.Errorf("Walk expected paths %v, got %v", expectedPaths, paths)
 	}
+}
 
-	// Test CopyFile
+func TestWalk_Bad(t *testing.T) {
-	tmpfile, err := os.CreateTemp("", "datanode-test-")
+	dn := New()
-	if err != nil {
+	// Walk non-existent path. fs.WalkDir will call the func with the error.
-		t.Fatalf("CreateTemp failed: %v", err)
+	var called bool
+	err := dn.Walk("nonexistent", func(path string, d fs.DirEntry, err error) error {
+		called = true
+		if err == nil {
+			t.Error("expected error for nonexistent path")
+		}
+		if !errors.Is(err, fs.ErrNotExist) {
+			t.Errorf("unexpected error message: %v", err)
+		}
+		return err // propagate error
+	})
+	if !called {
+		t.Fatal("walk function was not called for nonexistent root")
 	}
-	defer os.Remove(tmpfile.Name())
+	if !errors.Is(err, fs.ErrNotExist) {
+		t.Errorf("expected Walk to return fs.ErrNotExist, got %v", err)
+	}
+}
 
-	err = dn.CopyFile("foo.txt", tmpfile.Name(), 0644)
+func TestWalk_Ugly(t *testing.T) {
+	dn := New()
+	dn.AddData("a/b.txt", []byte("b"))
+	dn.AddData("a/c.txt", []byte("c"))
+
+	// Test stopping walk
+	walkErr := errors.New("stop walking")
+	var paths []string
+	err := dn.Walk(".", func(path string, d fs.DirEntry, err error) error {
+		if path == "a/b.txt" {
+			return walkErr
+		}
+		paths = append(paths, path)
+		return nil
+	})
+
+	if err != walkErr {
+		t.Errorf("expected walk to return the callback error, got %v", err)
+	}
+}
+
+func TestCopyFile_Good(t *testing.T) {
+	dn := New()
+	dn.AddData("foo.txt", []byte("foo"))
+
+	tmpfile := filepath.Join(t.TempDir(), "test.txt")
+	err := dn.CopyFile("foo.txt", tmpfile, 0644)
 	if err != nil {
 		t.Fatalf("CopyFile failed: %v", err)
 	}
 
-	content, err := os.ReadFile(tmpfile.Name())
+	content, err := os.ReadFile(tmpfile)
 	if err != nil {
 		t.Fatalf("ReadFile failed: %v", err)
 	}
 	if string(content) != "foo" {
-		t.Errorf("Expected foo, got %s", string(content))
+		t.Errorf("expected foo, got %s", string(content))
 	}
 }
+
+func TestCopyFile_Bad(t *testing.T) {
+	dn := New()
+	tmpfile := filepath.Join(t.TempDir(), "test.txt")
+
+	// Source does not exist
+	err := dn.CopyFile("nonexistent.txt", tmpfile, 0644)
+	if err == nil {
+		t.Fatal("expected error for nonexistent source file")
+	}
+
+	// Destination is not writable
+	dn.AddData("foo.txt", []byte("foo"))
+	err = dn.CopyFile("foo.txt", "/nonexistent_dir/test.txt", 0644)
+	if err == nil {
+		t.Fatal("expected error for unwritable destination")
+	}
+}
+
+func TestCopyFile_Ugly(t *testing.T) {
+	dn := New()
+	dn.AddData("bar/baz.txt", []byte("baz"))
+	tmpfile := filepath.Join(t.TempDir(), "test.txt")
+
+	// Attempting to copy a directory
+	err := dn.CopyFile("bar", tmpfile, 0644)
+	if err == nil {
+		t.Fatal("expected error when trying to copy a directory")
+	}
+}
+
+func toSortedNames(entries []fs.DirEntry) []string {
+	var names []string
+	for _, e := range entries {
+		names = append(names, e.Name())
+	}
+	sort.Strings(names)
+	return names
+}

pkg/github/github.go

@@ -47,6 +47,7 @@ func (g *githubClient) getPublicReposWithAPIURL(ctx context.Context, apiURL, use
 	client := NewAuthenticatedClient(ctx)
 	var allCloneURLs []string
 	url := fmt.Sprintf("%s/users/%s/repos", apiURL, userOrOrg)
+	isFirstRequest := true
 
 	for {
 		if err := ctx.Err(); err != nil {
@@ -63,24 +64,19 @@ func (g *githubClient) getPublicReposWithAPIURL(ctx context.Context, apiURL, use
 		}
 
 		if resp.StatusCode != http.StatusOK {
+			// If it's the first request for a user and it's a 404, we can try the org endpoint.
+			if isFirstRequest && strings.Contains(url, "/users/") && resp.StatusCode == http.StatusNotFound {
+				resp.Body.Close()
+				url = fmt.Sprintf("%s/orgs/%s/repos", apiURL, userOrOrg)
+				isFirstRequest = false // We are now trying the org endpoint.
+				continue                 // Re-run the loop with the org URL.
+			}
+			status := resp.Status
 			resp.Body.Close()
-			// Try organization endpoint
+			return nil, fmt.Errorf("failed to fetch repos: %s", status)
-			url = fmt.Sprintf("%s/orgs/%s/repos", apiURL, userOrOrg)
-			req, err = http.NewRequestWithContext(ctx, "GET", url, nil)
-			if err != nil {
-				return nil, err
-			}
-			req.Header.Set("User-Agent", "Borg-Data-Collector")
-			resp, err = client.Do(req)
-			if err != nil {
-				return nil, err
-			}
 		}
 
-		if resp.StatusCode != http.StatusOK {
+		isFirstRequest = false // Subsequent requests are for pagination.
-			resp.Body.Close()
-			return nil, fmt.Errorf("failed to fetch repos: %s", resp.Status)
-		}
 
 		var repos []Repo
 		if err := json.NewDecoder(resp.Body).Decode(&repos); err != nil {
@@ -94,9 +90,6 @@ func (g *githubClient) getPublicReposWithAPIURL(ctx context.Context, apiURL, use
 		}
 
 		linkHeader := resp.Header.Get("Link")
-		if linkHeader == "" {
-			break
-		}
 		nextURL := g.findNextURL(linkHeader)
 		if nextURL == "" {
 			break
@@ -111,8 +104,15 @@ func (g *githubClient) findNextURL(linkHeader string) string {
 	links := strings.Split(linkHeader, ",")
 	for _, link := range links {
 		parts := strings.Split(link, ";")
-		if len(parts) == 2 && strings.TrimSpace(parts[1]) == `rel="next"` {
+		if len(parts) < 2 {
-			return strings.Trim(strings.TrimSpace(parts[0]), "<>")
+			continue
+		}
+
+		if strings.TrimSpace(parts[1]) == `rel="next"` {
+			urlPart := strings.TrimSpace(parts[0])
+			if strings.HasPrefix(urlPart, "<") && strings.HasSuffix(urlPart, ">") {
+				return urlPart[1 : len(urlPart)-1]
+			}
 		}
 	}
 	return ""

pkg/github/github_test.go

@@ -5,120 +5,180 @@ import (
 	"context"
 	"io"
 	"net/http"
-	"net/url"
+	"strings"
 	"testing"
 
 	"github.com/Snider/Borg/pkg/mocks"
 )
 
-func TestGetPublicRepos(t *testing.T) {
+func TestGetPublicRepos_Good(t *testing.T) {
-	mockClient := mocks.NewMockClient(map[string]*http.Response{
+	t.Run("User Repos", func(t *testing.T) {
-		"https://api.github.com/users/testuser/repos": {
+		mockClient := mocks.NewMockClient(map[string]*http.Response{
-			StatusCode: http.StatusOK,
+			"https://api.github.com/users/testuser/repos": {
-			Header:     http.Header{"Content-Type": []string{"application/json"}},
+				StatusCode: http.StatusOK,
-			Body:       io.NopCloser(bytes.NewBufferString(`[{"clone_url": "https://github.com/testuser/repo1.git"}]`)),
+				Header:     http.Header{"Content-Type": []string{"application/json"}},
-		},
+				Body:       io.NopCloser(bytes.NewBufferString(`[{"clone_url": "https://github.com/testuser/repo1.git"}]`)),
-		"https://api.github.com/orgs/testorg/repos": {
+			},
-			StatusCode: http.StatusOK,
+		})
-			Header:     http.Header{"Content-Type": []string{"application/json"}, "Link": []string{`<https://api.github.com/organizations/123/repos?page=2>; rel="next"`}},
+		client := setupMockClient(t, mockClient)
-			Body:       io.NopCloser(bytes.NewBufferString(`[{"clone_url": "https://github.com/testorg/repo1.git"}]`)),
+		repos, err := client.getPublicReposWithAPIURL(context.Background(), "https://api.github.com", "testuser")
-		},
+		if err != nil {
-		"https://api.github.com/organizations/123/repos?page=2": {
+			t.Fatalf("getPublicReposWithAPIURL for user failed: %v", err)
-			StatusCode: http.StatusOK,
+		}
-			Header:     http.Header{"Content-Type": []string{"application/json"}},
+		if len(repos) != 1 || repos[0] != "https://github.com/testuser/repo1.git" {
-			Body:       io.NopCloser(bytes.NewBufferString(`[{"clone_url": "https://github.com/testorg/repo2.git"}]`)),
+			t.Errorf("unexpected user repos: %v", repos)
-		},
+		}
 	})
 
-	client := &githubClient{}
+	t.Run("Org Repos with Pagination", func(t *testing.T) {
-	oldClient := NewAuthenticatedClient
+		mockClient := mocks.NewMockClient(map[string]*http.Response{
-	NewAuthenticatedClient = func(ctx context.Context) *http.Client {
+			"https://api.github.com/users/testorg/repos": {
-		return mockClient
+				StatusCode: http.StatusNotFound, // Trigger fallback to org
-	}
+				Status:     "404 Not Found",
-	defer func() {
+				Body:       io.NopCloser(bytes.NewBufferString(`{}`)),
-		NewAuthenticatedClient = oldClient
+			},
-	}()
+			"https://api.github.com/orgs/testorg/repos": {
-
+				StatusCode: http.StatusOK,
-	// Test user repos
+				Header:     http.Header{"Content-Type": []string{"application/json"}, "Link": []string{`<https://api.github.com/organizations/123/repos?page=2>; rel="next"`}},
-	repos, err := client.getPublicReposWithAPIURL(context.Background(), "https://api.github.com", "testuser")
+				Body:       io.NopCloser(bytes.NewBufferString(`[{"clone_url": "https://github.com/testorg/repo1.git"}]`)),
-	if err != nil {
+			},
-		t.Fatalf("getPublicReposWithAPIURL for user failed: %v", err)
+			"https://api.github.com/organizations/123/repos?page=2": {
-	}
+				StatusCode: http.StatusOK,
-	if len(repos) != 1 || repos[0] != "https://github.com/testuser/repo1.git" {
+				Header:     http.Header{"Content-Type": []string{"application/json"}},
-		t.Errorf("unexpected user repos: %v", repos)
+				Body:       io.NopCloser(bytes.NewBufferString(`[{"clone_url": "https://github.com/testorg/repo2.git"}]`)),
-	}
+			},
-
+		})
-	// Test org repos with pagination
+		client := setupMockClient(t, mockClient)
-	repos, err = client.getPublicReposWithAPIURL(context.Background(), "https://api.github.com", "testorg")
+		repos, err := client.getPublicReposWithAPIURL(context.Background(), "https://api.github.com", "testorg")
-	if err != nil {
+		if err != nil {
-		t.Fatalf("getPublicReposWithAPIURL for org failed: %v", err)
+			t.Fatalf("getPublicReposWithAPIURL for org failed: %v", err)
-	}
+		}
-	if len(repos) != 2 || repos[0] != "https://github.com/testorg/repo1.git" || repos[1] != "https://github.com/testorg/repo2.git" {
+		if len(repos) != 2 || repos[0] != "https://github.com/testorg/repo1.git" || repos[1] != "https://github.com/testorg/repo2.git" {
-		t.Errorf("unexpected org repos: %v", repos)
+			t.Errorf("unexpected org repos: %v", repos)
-	}
+		}
-}
-func TestGetPublicRepos_Error(t *testing.T) {
-	u, _ := url.Parse("https://api.github.com/users/testuser/repos")
-	mockClient := mocks.NewMockClient(map[string]*http.Response{
-		"https://api.github.com/users/testuser/repos": {
-			StatusCode: http.StatusNotFound,
-			Status:     "404 Not Found",
-			Header:     http.Header{"Content-Type": []string{"application/json"}},
-			Body:       io.NopCloser(bytes.NewBufferString("")),
-			Request:    &http.Request{Method: "GET", URL: u},
-		},
-		"https://api.github.com/orgs/testuser/repos": {
-			StatusCode: http.StatusNotFound,
-			Status:     "404 Not Found",
-			Header:     http.Header{"Content-Type": []string{"application/json"}},
-			Body:       io.NopCloser(bytes.NewBufferString("")),
-			Request:    &http.Request{Method: "GET", URL: u},
-		},
 	})
-	expectedErr := "failed to fetch repos: 404 Not Found"
-
-	client := &githubClient{}
-	oldClient := NewAuthenticatedClient
-	NewAuthenticatedClient = func(ctx context.Context) *http.Client {
-		return mockClient
-	}
-	defer func() {
-		NewAuthenticatedClient = oldClient
-	}()
-
-	// Test user repos
-	_, err := client.getPublicReposWithAPIURL(context.Background(), "https://api.github.com", "testuser")
-	if err.Error() != expectedErr {
-		t.Fatalf("getPublicReposWithAPIURL for user failed: %v", err)
-	}
 }
 
-func TestFindNextURL(t *testing.T) {
+func TestGetPublicRepos_Bad(t *testing.T) {
+	t.Run("Not Found", func(t *testing.T) {
+		mockClient := mocks.NewMockClient(map[string]*http.Response{
+			"https://api.github.com/users/testuser/repos": {
+				StatusCode: http.StatusNotFound,
+				Status:     "404 Not Found",
+				Body:       io.NopCloser(bytes.NewBufferString(`{"message": "Not Found"}`)),
+			},
+			"https://api.github.com/orgs/testuser/repos": {
+				StatusCode: http.StatusNotFound,
+				Status:     "404 Not Found",
+				Body:       io.NopCloser(bytes.NewBufferString(`{"message": "Not Found"}`)),
+			},
+		})
+		client := setupMockClient(t, mockClient)
+		_, err := client.getPublicReposWithAPIURL(context.Background(), "https://api.github.com", "testuser")
+		if err == nil {
+			t.Fatal("expected an error but got nil")
+		}
+		if !strings.Contains(err.Error(), "404 Not Found") {
+			t.Errorf("expected '404 Not Found' in error message, got %q", err)
+		}
+	})
+
+	t.Run("Invalid JSON", func(t *testing.T) {
+		mockClient := mocks.NewMockClient(map[string]*http.Response{
+			"https://api.github.com/users/badjson/repos": {
+				StatusCode: http.StatusOK,
+				Header:     http.Header{"Content-Type": []string{"application/json"}},
+				Body:       io.NopCloser(bytes.NewBufferString(`[{"clone_url": "invalid}`)),
+			},
+		})
+		client := setupMockClient(t, mockClient)
+		_, err := client.getPublicReposWithAPIURL(context.Background(), "https://api.github.com", "badjson")
+		if err == nil {
+			t.Fatal("expected an error for invalid JSON, but got nil")
+		}
+	})
+}
+
+func TestGetPublicRepos_Ugly(t *testing.T) {
+	t.Run("Empty Repo List", func(t *testing.T) {
+		mockClient := mocks.NewMockClient(map[string]*http.Response{
+			"https://api.github.com/users/empty/repos": {
+				StatusCode: http.StatusOK,
+				Header:     http.Header{"Content-Type": []string{"application/json"}},
+				Body:       io.NopCloser(bytes.NewBufferString(`[]`)),
+			},
+		})
+		client := setupMockClient(t, mockClient)
+		repos, err := client.getPublicReposWithAPIURL(context.Background(), "https://api.github.com", "empty")
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if len(repos) != 0 {
+			t.Errorf("expected 0 repos, got %d", len(repos))
+		}
+	})
+}
+
+func TestFindNextURL_Good(t *testing.T) {
 	client := &githubClient{}
 	linkHeader := `<https://api.github.com/organizations/123/repos?page=2>; rel="next", <https://api.github.com/organizations/123/repos?page=1>; rel="prev"`
 	nextURL := client.findNextURL(linkHeader)
 	if nextURL != "https://api.github.com/organizations/123/repos?page=2" {
 		t.Errorf("unexpected next URL: %s", nextURL)
 	}
+}
 
-	linkHeader = `<https://api.github.com/organizations/123/repos?page=1>; rel="prev"`
+func TestFindNextURL_Bad(t *testing.T) {
-	nextURL = client.findNextURL(linkHeader)
+	client := &githubClient{}
+	linkHeader := `<https://api.github.com/organizations/123/repos?page=1>; rel="prev"`
+	nextURL := client.findNextURL(linkHeader)
 	if nextURL != "" {
-		t.Errorf("unexpected next URL: %s", nextURL)
+		t.Errorf("unexpected next URL for header with no 'next': %s", nextURL)
+	}
+
+	nextURL = client.findNextURL("")
+	if nextURL != "" {
+		t.Errorf("unexpected next URL for empty header: %s", nextURL)
 	}
 }
 
-func TestNewAuthenticatedClient(t *testing.T) {
+func TestFindNextURL_Ugly(t *testing.T) {
-	// Test with no token
+	client := &githubClient{}
+	// Malformed: missing angle brackets
+	linkHeader := `https://api.github.com/organizations/123/repos?page=2; rel="next"`
+	nextURL := client.findNextURL(linkHeader)
+	if nextURL != "" {
+		t.Errorf("unexpected next URL for malformed header: %s", nextURL)
+	}
+}
+
+func TestNewAuthenticatedClient_Good(t *testing.T) {
+	t.Setenv("GITHUB_TOKEN", "test-token")
+	client := NewAuthenticatedClient(context.Background())
+	if client == http.DefaultClient {
+		t.Error("expected an authenticated client, but got http.DefaultClient")
+	}
+}
+
+func TestNewAuthenticatedClient_Bad(t *testing.T) {
+	// Unset the variable to ensure it's not present
+	t.Setenv("GITHUB_TOKEN", "")
 	client := NewAuthenticatedClient(context.Background())
 	if client != http.DefaultClient {
-		t.Errorf("expected http.DefaultClient, but got something else")
+		t.Error("expected http.DefaultClient when no token is set, but got something else")
-	}
-
-	// Test with token
-	t.Setenv("GITHUB_TOKEN", "test-token")
-	client = NewAuthenticatedClient(context.Background())
-	if client == http.DefaultClient {
-		t.Errorf("expected an authenticated client, but got http.DefaultClient")
 	}
 }
+
+// setupMockClient is a helper function to inject a mock http.Client.
+func setupMockClient(t *testing.T, mock *http.Client) *githubClient {
+	client := &githubClient{}
+	originalNewAuthenticatedClient := NewAuthenticatedClient
+	NewAuthenticatedClient = func(ctx context.Context) *http.Client {
+		return mock
+	}
+	// Restore the original function after the test
+	t.Cleanup(func() {
+		NewAuthenticatedClient = originalNewAuthenticatedClient
+	})
+	return client
+}

pkg/matrix/matrix.go

@@ -4,11 +4,17 @@ import (
 	"archive/tar"
 	"bytes"
 	"encoding/json"
+	"errors"
 	"io/fs"
 
 	"github.com/Snider/Borg/pkg/datanode"
 )
 
+var (
+	ErrDataNodeRequired = errors.New("datanode is required")
+	ErrConfigIsNil      = errors.New("config is nil")
+)
+
 // TerminalIsolationMatrix represents a runc bundle.
 type TerminalIsolationMatrix struct {
 	Config []byte
@@ -37,6 +43,9 @@ func New() (*TerminalIsolationMatrix, error) {
 
 // FromDataNode creates a new TerminalIsolationMatrix from a DataNode.
 func FromDataNode(dn *datanode.DataNode) (*TerminalIsolationMatrix, error) {
+	if dn == nil {
+		return nil, ErrDataNodeRequired
+	}
 	m, err := New()
 	if err != nil {
 		return nil, err
@@ -47,6 +56,9 @@ func FromDataNode(dn *datanode.DataNode) (*TerminalIsolationMatrix, error) {
 
 // ToTar serializes the TerminalIsolationMatrix to a tarball.
 func (m *TerminalIsolationMatrix) ToTar() ([]byte, error) {
+	if m.Config == nil {
+		return nil, ErrConfigIsNil
+	}
 	buf := new(bytes.Buffer)
 	tw := tar.NewWriter(buf)
 
@@ -76,6 +88,10 @@ func (m *TerminalIsolationMatrix) ToTar() ([]byte, error) {
 	// Add the rootfs files.
 	err := m.RootFS.Walk(".", func(path string, d fs.DirEntry, err error) error {
 		if err != nil {
+			// If the root directory doesn't exist (i.e. empty datanode), it's not an error.
+			if path == "." && errors.Is(err, fs.ErrNotExist) {
+				return nil
+			}
 			return err
 		}
 

pkg/matrix/matrix_test.go

@@ -3,13 +3,15 @@ package matrix
 import (
 	"archive/tar"
 	"bytes"
+	"encoding/json"
+	"errors"
 	"io"
 	"testing"
 
 	"github.com/Snider/Borg/pkg/datanode"
 )
 
-func TestNew(t *testing.T) {
+func TestNew_Good(t *testing.T) {
 	m, err := New()
 	if err != nil {
 		t.Fatalf("New() returned an error: %v", err)
@@ -23,9 +25,14 @@ func TestNew(t *testing.T) {
 	if m.RootFS == nil {
 		t.Error("New() returned a matrix with a nil RootFS")
 	}
+
+	// Verify the config is valid JSON
+	if !json.Valid(m.Config) {
+		t.Error("New() returned a matrix with invalid JSON config")
+	}
 }
 
-func TestFromDataNode(t *testing.T) {
+func TestFromDataNode_Good(t *testing.T) {
 	dn := datanode.New()
 	dn.AddData("test.txt", []byte("hello world"))
 	m, err := FromDataNode(dn)
@@ -38,9 +45,22 @@ func TestFromDataNode(t *testing.T) {
 	if m.RootFS != dn {
 		t.Error("FromDataNode() did not set the RootFS correctly")
 	}
+	if m.Config == nil {
+		t.Error("FromDataNode() did not create a default config")
+	}
 }
 
-func TestToTar(t *testing.T) {
+func TestFromDataNode_Bad(t *testing.T) {
+	_, err := FromDataNode(nil)
+	if err == nil {
+		t.Fatal("expected error when passing a nil datanode, but got nil")
+	}
+	if !errors.Is(err, ErrDataNodeRequired) {
+		t.Errorf("expected ErrDataNodeRequired, got %v", err)
+	}
+}
+
+func TestToTar_Good(t *testing.T) {
 	m, err := New()
 	if err != nil {
 		t.Fatalf("New() returned an error: %v", err)
@@ -55,35 +75,65 @@ func TestToTar(t *testing.T) {
 	}
 
 	tr := tar.NewReader(bytes.NewReader(tarball))
-	foundConfig := false
+	found := make(map[string]bool)
-	foundRootFS := false
-	foundTestFile := false
 	for {
 		header, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
 		if err != nil {
+			t.Fatalf("failed to read tar header: %v", err)
+		}
+		found[header.Name] = true
+	}
+
+	expectedFiles := []string{"config.json", "rootfs/", "rootfs/test.txt"}
+	for _, f := range expectedFiles {
+		if !found[f] {
+			t.Errorf("%s not found in matrix tarball", f)
+		}
+	}
+}
+
+func TestToTar_Ugly(t *testing.T) {
+	t.Run("Empty RootFS", func(t *testing.T) {
+		m, _ := New()
+		tarball, err := m.ToTar()
+		if err != nil {
+			t.Fatalf("ToTar() with empty rootfs returned an error: %v", err)
+		}
+		tr := tar.NewReader(bytes.NewReader(tarball))
+		found := make(map[string]bool)
+		for {
+			header, err := tr.Next()
 			if err == io.EOF {
 				break
 			}
-			t.Fatalf("failed to read tar header: %v", err)
+			if err != nil {
+				t.Fatalf("failed to read tar header: %v", err)
+			}
+			found[header.Name] = true
 		}
-
+		if !found["config.json"] {
-		switch header.Name {
+			t.Error("config.json not found in matrix")
-		case "config.json":
-			foundConfig = true
-		case "rootfs/":
-			foundRootFS = true
-		case "rootfs/test.txt":
-			foundTestFile = true
 		}
-	}
+		if !found["rootfs/"] {
+			t.Error("rootfs/ directory not found in matrix")
+		}
+		if len(found) != 2 {
+			t.Errorf("expected 2 files in tar, but found %d", len(found))
+		}
+	})
 
-	if !foundConfig {
+	t.Run("Nil Config", func(t *testing.T) {
-		t.Error("config.json not found in matrix")
+		m, _ := New()
-	}
+		m.Config = nil // This should not happen in practice
-	if !foundRootFS {
+		_, err := m.ToTar()
-		t.Error("rootfs/ not found in matrix")
+		if err == nil {
-	}
+			t.Fatal("expected error when Config is nil, but got nil")
-	if !foundTestFile {
+		}
-		t.Error("rootfs/test.txt not found in matrix")
+		if !errors.Is(err, ErrConfigIsNil) {
-	}
+			t.Errorf("expected ErrConfigIsNil, got %v", err)
+		}
+	})
 }

pkg/pwa/pwa.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 
 	"github.com/Snider/Borg/pkg/datanode"
 	"github.com/schollz/progressbar/v3"
@@ -36,6 +37,10 @@ func (p *pwaClient) FindManifest(pwaURL string) (string, error) {
 	}
 	defer resp.Body.Close()
 
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return "", fmt.Errorf("failed to fetch PWA page: status code %d", resp.StatusCode)
+	}
+
 	doc, err := html.Parse(resp.Body)
 	if err != nil {
 		return "", err
@@ -81,6 +86,9 @@ func (p *pwaClient) FindManifest(pwaURL string) (string, error) {
 // DownloadAndPackagePWA downloads and packages a PWA into a DataNode.
 func (p *pwaClient) DownloadAndPackagePWA(pwaURL, manifestURL string, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
 	dn := datanode.New()
+	var wg sync.WaitGroup
+	var errs []error
+	var mu sync.Mutex
 
 	type Manifest struct {
 		StartURL string `json:"start_url"`
@@ -89,82 +97,98 @@ func (p *pwaClient) DownloadAndPackagePWA(pwaURL, manifestURL string, bar *progr
 		} `json:"icons"`
 	}
 
-	downloadAndAdd := func(assetURL string) error {
+	downloadAndAdd := func(assetURL string) {
+		defer wg.Done()
 		if bar != nil {
 			bar.Add(1)
 		}
 		resp, err := p.client.Get(assetURL)
 		if err != nil {
-			return fmt.Errorf("failed to download %s: %w", assetURL, err)
+			mu.Lock()
+			errs = append(errs, fmt.Errorf("failed to download %s: %w", assetURL, err))
+			mu.Unlock()
+			return
 		}
 		defer resp.Body.Close()
 
 		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-			return fmt.Errorf("failed to download %s: status code %d", assetURL, resp.StatusCode)
+			mu.Lock()
+			errs = append(errs, fmt.Errorf("failed to download %s: status code %d", assetURL, resp.StatusCode))
+			mu.Unlock()
+			return
 		}
 
 		body, err := io.ReadAll(resp.Body)
 		if err != nil {
-			return fmt.Errorf("failed to read body of %s: %w", assetURL, err)
+			mu.Lock()
+			errs = append(errs, fmt.Errorf("failed to read body of %s: %w", assetURL, err))
+			mu.Unlock()
+			return
 		}
 
 		u, err := url.Parse(assetURL)
 		if err != nil {
-			return fmt.Errorf("failed to parse asset URL %s: %w", assetURL, err)
+			mu.Lock()
+			errs = append(errs, fmt.Errorf("failed to parse asset URL %s: %w", assetURL, err))
+			mu.Unlock()
+			return
 		}
 		dn.AddData(strings.TrimPrefix(u.Path, "/"), body)
-		return nil
 	}
 
-	// Download manifest
+	// Download manifest first, synchronously.
-	if err := downloadAndAdd(manifestURL); err != nil {
+	resp, err := p.client.Get(manifestURL)
-		return nil, err
-	}
-
-	// Parse manifest and download assets
-	var manifestPath string
-	u, parseErr := url.Parse(manifestURL)
-	if parseErr != nil {
-		manifestPath = "manifest.json"
-	} else {
-		manifestPath = strings.TrimPrefix(u.Path, "/")
-	}
-
-	manifestFile, err := dn.Open(manifestPath)
 	if err != nil {
-		return nil, fmt.Errorf("failed to open manifest from datanode: %w", err)
+		return nil, fmt.Errorf("failed to download manifest: %w", err)
 	}
-	defer manifestFile.Close()
+	defer resp.Body.Close()
 
-	manifestData, err := io.ReadAll(manifestFile)
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return nil, fmt.Errorf("failed to download manifest: status code %d", resp.StatusCode)
+	}
+
+	manifestData, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return nil, fmt.Errorf("failed to read manifest from datanode: %w", err)
+		return nil, fmt.Errorf("failed to read manifest body: %w", err)
 	}
 
+	u, _ := url.Parse(manifestURL)
+	dn.AddData(strings.TrimPrefix(u.Path, "/"), manifestData)
+
+	// Parse manifest and download assets concurrently.
 	var manifest Manifest
 	if err := json.Unmarshal(manifestData, &manifest); err != nil {
 		return nil, fmt.Errorf("failed to parse manifest: %w", err)
 	}
 
-	// Download start_url
+	assetsToDownload := []string{}
-	startURL, err := p.resolveURL(manifestURL, manifest.StartURL)
+	if manifest.StartURL != "" {
-	if err != nil {
+		startURL, err := p.resolveURL(manifestURL, manifest.StartURL)
-		return nil, fmt.Errorf("failed to resolve start_url: %w", err)
+		if err == nil {
+			assetsToDownload = append(assetsToDownload, startURL.String())
+		}
 	}
-	if err := downloadAndAdd(startURL.String()); err != nil {
+	for _, icon := range manifest.Icons {
-		return nil, err
+		if icon.Src != "" {
+			iconURL, err := p.resolveURL(manifestURL, icon.Src)
+			if err == nil {
+				assetsToDownload = append(assetsToDownload, iconURL.String())
+			}
+		}
 	}
 
-	// Download icons
+	wg.Add(len(assetsToDownload))
-	for _, icon := range manifest.Icons {
+	for _, asset := range assetsToDownload {
-		iconURL, err := p.resolveURL(manifestURL, icon.Src)
+		go downloadAndAdd(asset)
-		if err != nil {
+	}
-			// Skip icons with bad URLs
+	wg.Wait()
-			continue
+
-		}
+	if len(errs) > 0 {
-		if err := downloadAndAdd(iconURL.String()); err != nil {
+		var errStrings []string
-			return nil, err
+		for _, e := range errs {
+			errStrings = append(errStrings, e.Error())
 		}
+		return dn, fmt.Errorf("%s", strings.Join(errStrings, "; "))
 	}
 
 	return dn, nil

pkg/pwa/pwa_error_test.go

@@ -1,47 +0,0 @@
-package pwa
-
-import (
-	"fmt"
-	"net/http"
-	"net/http/httptest"
-	"testing"
-
-	"github.com/schollz/progressbar/v3"
-)
-
-func TestDownloadAndPackagePWA_Error(t *testing.T) {
-	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.URL.Path == "/manifest.json" {
-			w.Header().Set("Content-Type", "application/json")
-			w.Write([]byte(`{"start_url": "index.html"}`))
-		} else {
-			http.NotFound(w, r)
-		}
-	}))
-	defer server.Close()
-
-	client := newTestPWAClient()
-
-	// Test with a server that returns a 404 for the start_url
-	bar := progressbar.New(1)
-	_, err := client.DownloadAndPackagePWA(server.URL, server.URL+"/manifest.json", bar)
-	if err == nil {
-		t.Fatal("Expected an error when the start_url returns a 404, but got nil")
-	}
-
-	// Test with a bad manifest URL
-	_, err = client.DownloadAndPackagePWA(server.URL, "http://bad.url/manifest.json", bar)
-	if err == nil {
-		t.Fatal("Expected an error when the manifest URL is bad, but got nil")
-	}
-
-	// Test with a manifest that is not valid JSON
-	server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		fmt.Fprintln(w, "this is not json")
-	}))
-	defer server2.Close()
-	_, err = client.DownloadAndPackagePWA(server2.URL, server2.URL, bar)
-	if err == nil {
-		t.Fatal("Expected an error when the manifest is not valid JSON, but got nil")
-	}
-}

pkg/pwa/pwa_test.go

@@ -1,92 +1,91 @@
 package pwa
 
 import (
+	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/schollz/progressbar/v3"
 )
 
-func newTestPWAClient() PWAClient {
+// --- Test Cases for FindManifest ---
-	return NewPWAClient()
-}
 
-func TestFindManifest(t *testing.T) {
+func TestFindManifest_Good(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "text/html")
-		w.Write([]byte(`
+		fmt.Fprint(w, `<html><head><link rel="manifest" href="manifest.json"></head></html>`)
-			<!DOCTYPE html>
-			<html>
-			<head>
-				<title>Test PWA</title>
-				<link rel="manifest" href="manifest.json">
-			</head>
-			<body>
-				<h1>Hello, PWA!</h1>
-			</body>
-			</html>
-		`))
 	}))
 	defer server.Close()
 
-	client := newTestPWAClient()
+	client := NewPWAClient()
 	expectedURL := server.URL + "/manifest.json"
 	actualURL, err := client.FindManifest(server.URL)
 	if err != nil {
 		t.Fatalf("FindManifest failed: %v", err)
 	}
-
 	if actualURL != expectedURL {
 		t.Errorf("Expected manifest URL %s, but got %s", expectedURL, actualURL)
 	}
 }
 
-func TestDownloadAndPackagePWA(t *testing.T) {
+func TestFindManifest_Bad(t *testing.T) {
-	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	t.Run("No Manifest Link", func(t *testing.T) {
-		switch r.URL.Path {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		case "/":
 			w.Header().Set("Content-Type", "text/html")
-			w.Write([]byte(`
+			fmt.Fprint(w, `<html><head></head></html>`)
-				<!DOCTYPE html>
+		}))
-				<html>
+		defer server.Close()
-				<head>
+		client := NewPWAClient()
-					<title>Test PWA</title>
+		_, err := client.FindManifest(server.URL)
-					<link rel="manifest" href="manifest.json">
+		if err == nil {
-				</head>
+			t.Fatal("expected an error, but got none")
-				<body>
-					<h1>Hello, PWA!</h1>
-				</body>
-				</html>
-			`))
-		case "/manifest.json":
-			w.Header().Set("Content-Type", "application/json")
-			w.Write([]byte(`{
-				"name": "Test PWA",
-				"short_name": "TestPWA",
-				"start_url": "index.html",
-				"icons": [
-					{
-						"src": "icon.png",
-						"sizes": "192x192",
-						"type": "image/png"
-					}
-				]
-			}`))
-		case "/index.html":
-			w.Header().Set("Content-Type", "text/html")
-			w.Write([]byte(`<h1>Hello, PWA!</h1>`))
-		case "/icon.png":
-			w.Header().Set("Content-Type", "image/png")
-			w.Write([]byte("fake image data"))
-		default:
-			http.NotFound(w, r)
 		}
-	}))
+	})
+
+	t.Run("Server Error", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+		}))
+		defer server.Close()
+		client := NewPWAClient()
+		_, err := client.FindManifest(server.URL)
+		if err == nil {
+			t.Fatal("expected an error for server error, but got none")
+		}
+	})
+}
+
+func TestFindManifest_Ugly(t *testing.T) {
+	t.Run("Multiple Manifest Links", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "text/html")
+			fmt.Fprint(w, `<html><head><link rel="manifest" href="first.json"><link rel="manifest" href="second.json"></head></html>`)
+		}))
+		defer server.Close()
+		client := NewPWAClient()
+		// Should find the first one
+		expectedURL := server.URL + "/first.json"
+		actualURL, err := client.FindManifest(server.URL)
+		if err != nil {
+			t.Fatalf("FindManifest failed: %v", err)
+		}
+		if actualURL != expectedURL {
+			t.Errorf("Expected manifest URL %s, but got %s", expectedURL, actualURL)
+		}
+	})
+}
+
+// --- Test Cases for DownloadAndPackagePWA ---
+
+func TestDownloadAndPackagePWA_Good(t *testing.T) {
+	server := newPWATestServer()
 	defer server.Close()
 
-	client := newTestPWAClient()
+	client := NewPWAClient()
-	bar := progressbar.New(1)
+	bar := progressbar.NewOptions(1, progressbar.OptionSetWriter(io.Discard))
 	dn, err := client.DownloadAndPackagePWA(server.URL, server.URL+"/manifest.json", bar)
 	if err != nil {
 		t.Fatalf("DownloadAndPackagePWA failed: %v", err)
@@ -94,18 +93,70 @@ func TestDownloadAndPackagePWA(t *testing.T) {
 
 	expectedFiles := []string{"manifest.json", "index.html", "icon.png"}
 	for _, file := range expectedFiles {
-		// The path in the datanode is relative to the root of the domain, so we need to remove the leading slash.
+		exists, _ := dn.Exists(file)
-		exists, err := dn.Exists(file)
-		if err != nil {
-			t.Fatalf("Exists failed for %s: %v", file, err)
-		}
 		if !exists {
 			t.Errorf("Expected to find file %s in DataNode, but it was not found", file)
 		}
 	}
 }
 
-func TestResolveURL(t *testing.T) {
+func TestDownloadAndPackagePWA_Bad(t *testing.T) {
+	t.Run("Bad Manifest URL", func(t *testing.T) {
+		server := newPWATestServer()
+		defer server.Close()
+		client := NewPWAClient()
+		_, err := client.DownloadAndPackagePWA(server.URL, server.URL+"/nonexistent-manifest.json", nil)
+		if err == nil {
+			t.Fatal("expected an error for bad manifest url, but got none")
+		}
+	})
+
+	t.Run("Asset 404", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/manifest.json" {
+				w.Header().Set("Content-Type", "application/json")
+				fmt.Fprint(w, `{"start_url": "nonexistent.html"}`)
+			} else {
+				http.NotFound(w, r)
+			}
+		}))
+		defer server.Close()
+		client := NewPWAClient()
+		_, err := client.DownloadAndPackagePWA(server.URL, server.URL+"/manifest.json", nil)
+		if err == nil {
+			t.Fatal("expected an error for asset 404, but got none")
+		}
+		// The current implementation aggregates errors.
+		if !strings.Contains(err.Error(), "status code 404") {
+			t.Errorf("expected error to contain 'status code 404', but got: %v", err)
+		}
+	})
+}
+
+func TestDownloadAndPackagePWA_Ugly(t *testing.T) {
+	t.Run("Manifest with no assets", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			fmt.Fprint(w, `{ "name": "Test PWA" }`) // valid json, but no assets
+		}))
+		defer server.Close()
+
+		client := NewPWAClient()
+		dn, err := client.DownloadAndPackagePWA(server.URL, server.URL+"/manifest.json", nil)
+		if err != nil {
+			t.Fatalf("unexpected error for manifest with no assets: %v", err)
+		}
+		// Should still contain the manifest itself
+		exists, _ := dn.Exists("manifest.json")
+		if !exists {
+			t.Error("expected manifest.json to be in the datanode")
+		}
+	})
+}
+
+// --- Test Cases for resolveURL ---
+
+func TestResolveURL_Good(t *testing.T) {
 	client := NewPWAClient().(*pwaClient)
 	tests := []struct {
 		base string
@@ -114,10 +165,8 @@ func TestResolveURL(t *testing.T) {
 	}{
 		{"http://example.com/", "foo.html", "http://example.com/foo.html"},
 		{"http://example.com/foo/", "bar.html", "http://example.com/foo/bar.html"},
-		{"http://example.com/foo", "bar.html", "http://example.com/bar.html"},
 		{"http://example.com/foo/", "/bar.html", "http://example.com/bar.html"},
-		{"http://example.com/foo", "/bar.html", "http://example.com/bar.html"},
+		{"http://example.com/", "http://othersite.com/bar.html", "http://othersite.com/bar.html"},
-		{"http://example.com/", "http://example.com/foo/bar.html", "http://example.com/foo/bar.html"},
 	}
 
 	for _, tt := range tests {
@@ -132,34 +181,38 @@ func TestResolveURL(t *testing.T) {
 	}
 }
 
-func TestPWA_Bad(t *testing.T) {
+func TestResolveURL_Bad(t *testing.T) {
-	client := NewPWAClient()
+	client := NewPWAClient().(*pwaClient)
-
+	_, err := client.resolveURL("http://^invalid.com", "foo.html")
-	// Test FindManifest with no manifest
-	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "text/html")
-		w.Write([]byte(`
-			<!DOCTYPE html>
-			<html>
-			<head>
-				<title>Test PWA</title>
-			</head>
-			<body>
-				<h1>Hello, PWA!</h1>
-			</body>
-			</html>
-		`))
-	}))
-	defer server.Close()
-
-	_, err := client.FindManifest(server.URL)
 	if err == nil {
-		t.Fatalf("expected an error, but got none")
+		t.Error("expected error for malformed base URL, but got nil")
-	}
-
-	// Test DownloadAndPackagePWA with bad manifest
-	_, err = client.DownloadAndPackagePWA(server.URL, server.URL+"/manifest.json", nil)
-	if err == nil {
-		t.Fatalf("expected an error, but got none")
 	}
 }
+
+// --- Helpers ---
+
+// newPWATestServer creates a test server for a simple PWA.
+func newPWATestServer() *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/":
+			w.Header().Set("Content-Type", "text/html")
+			fmt.Fprint(w, `<html><head><link rel="manifest" href="manifest.json"></head></html>`)
+		case "/manifest.json":
+			w.Header().Set("Content-Type", "application/json")
+			fmt.Fprint(w, `{
+				"name": "Test PWA",
+				"start_url": "index.html",
+				"icons": [{"src": "icon.png"}]
+			}`)
+		case "/index.html":
+			w.Header().Set("Content-Type", "text/html")
+			fmt.Fprint(w, `<h1>Hello, PWA!</h1>`)
+		case "/icon.png":
+			w.Header().Set("Content-Type", "image/png")
+			fmt.Fprint(w, "fake image data")
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+}

pkg/vcs/git.go

@@ -39,6 +39,9 @@ func (g *gitCloner) CloneGitRepository(repoURL string, progress io.Writer) (*dat
 
 	_, err = git.PlainClone(tempPath, false, cloneOptions)
 	if err != nil {
+		if err.Error() == "remote repository is empty" {
+			return datanode.New(), nil
+		}
 		return nil, err
 	}
 
@@ -47,6 +50,10 @@ func (g *gitCloner) CloneGitRepository(repoURL string, progress io.Writer) (*dat
 		if err != nil {
 			return err
 		}
+		// Skip the .git directory
+		if info.IsDir() && info.Name() == ".git" {
+			return filepath.SkipDir
+		}
 		if !info.IsDir() {
 			content, err := os.ReadFile(path)
 			if err != nil {

pkg/vcs/git_test.go

@@ -1,81 +1,77 @@
 package vcs
 
 import (
+	"bytes"
+	"io"
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strings"
 	"testing"
 )
 
-func TestCloneGitRepository(t *testing.T) {
+// setupTestRepo creates a bare git repository with a single commit.
-	// Create a temporary directory for the bare repository
+func setupTestRepo(t *testing.T) (repoPath string) {
+	t.Helper()
+
+	// Create a temporary directory for the bare repository.
 	bareRepoPath, err := os.MkdirTemp("", "bare-repo-")
 	if err != nil {
 		t.Fatalf("Failed to create temp dir for bare repo: %v", err)
 	}
-	defer os.RemoveAll(bareRepoPath)
 
-	// Initialize a bare git repository
+	// Initialize the bare git repository.
-	cmd := exec.Command("git", "init", "--bare")
+	runCmd(t, bareRepoPath, "git", "init", "--bare")
-	cmd.Dir = bareRepoPath
-	if err := cmd.Run(); err != nil {
-		t.Fatalf("Failed to init bare repo: %v", err)
-	}
 
-	// Clone the bare repository to a temporary directory to add a commit
+	// Clone the bare repository to a temporary directory to add a commit.
 	clonePath, err := os.MkdirTemp("", "clone-")
 	if err != nil {
 		t.Fatalf("Failed to create temp dir for clone: %v", err)
 	}
 	defer os.RemoveAll(clonePath)
 
-	cmd = exec.Command("git", "clone", bareRepoPath, clonePath)
+	runCmd(t, clonePath, "git", "clone", bareRepoPath, ".")
-	if err := cmd.Run(); err != nil {
-		t.Fatalf("Failed to clone bare repo: %v", err)
-	}
 
-	// Create a file and commit it
+	// Create a file and commit it.
 	filePath := filepath.Join(clonePath, "foo.txt")
 	if err := os.WriteFile(filePath, []byte("foo"), 0644); err != nil {
 		t.Fatalf("Failed to write file: %v", err)
 	}
-	cmd = exec.Command("git", "add", "foo.txt")
+	runCmd(t, clonePath, "git", "add", "foo.txt")
-	cmd.Dir = clonePath
+	runCmd(t, clonePath, "git", "config", "user.email", "test@example.com")
-	if err := cmd.Run(); err != nil {
+	runCmd(t, clonePath, "git", "config", "user.name", "Test User")
-		t.Fatalf("Failed to git add: %v", err)
+	runCmd(t, clonePath, "git", "commit", "-m", "Initial commit")
-	}
+	runCmd(t, clonePath, "git", "push", "origin", "master")
 
-	cmd = exec.Command("git", "config", "user.email", "test@example.com")
+	return bareRepoPath
-	cmd.Dir = clonePath
+}
-	if err := cmd.Run(); err != nil {
-		t.Fatalf("Failed to set git user.email: %v", err)
-	}
 
-	cmd = exec.Command("git", "config", "user.name", "Test User")
+// runCmd executes a command and fails the test if it fails.
-	cmd.Dir = clonePath
+func runCmd(t *testing.T, dir, name string, args ...string) {
-	if err := cmd.Run(); err != nil {
+	t.Helper()
-		t.Fatalf("Failed to set git user.name: %v", err)
+	cmd := exec.Command(name, args...)
+	cmd.Dir = dir
+	if testing.Verbose() {
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
 	}
+	if err := cmd.Run(); err != nil {
+		t.Fatalf("Command %q failed: %v", strings.Join(append([]string{name}, args...), " "), err)
+	}
+}
 
-	cmd = exec.Command("git", "commit", "-m", "Initial commit")
+func TestCloneGitRepository_Good(t *testing.T) {
-	cmd.Dir = clonePath
+	repoPath := setupTestRepo(t)
-	if err := cmd.Run(); err != nil {
+	defer os.RemoveAll(repoPath)
-		t.Fatalf("Failed to git commit: %v", err)
-	}
-	cmd = exec.Command("git", "push", "origin", "master")
-	cmd.Dir = clonePath
-	if err := cmd.Run(); err != nil {
-		t.Fatalf("Failed to git push: %v", err)
-	}
 
-	// Clone the repository using the function we're testing
 	cloner := NewGitCloner()
-	dn, err := cloner.CloneGitRepository("file://"+bareRepoPath, os.Stdout)
+	var out bytes.Buffer
+	dn, err := cloner.CloneGitRepository("file://"+repoPath, &out)
 	if err != nil {
-		t.Fatalf("CloneGitRepository failed: %v", err)
+		t.Fatalf("CloneGitRepository failed: %v\nOutput: %s", err, out.String())
 	}
 
-	// Verify the DataNode contains the correct file
+	// Verify the DataNode contains the correct file.
 	exists, err := dn.Exists("foo.txt")
 	if err != nil {
 		t.Fatalf("Exists failed: %v", err)
@@ -84,3 +80,45 @@ func TestCloneGitRepository(t *testing.T) {
 		t.Errorf("Expected to find file foo.txt in DataNode, but it was not found")
 	}
 }
+
+func TestCloneGitRepository_Bad(t *testing.T) {
+	t.Run("Non-existent repository", func(t *testing.T) {
+		cloner := NewGitCloner()
+		_, err := cloner.CloneGitRepository("file:///non-existent-repo", io.Discard)
+		if err == nil {
+			t.Fatal("Expected an error for a non-existent repository, but got nil")
+		}
+		if !strings.Contains(err.Error(), "repository not found") {
+			t.Errorf("Expected error to be about 'repository not found', but got: %v", err)
+		}
+	})
+
+	t.Run("Invalid URL", func(t *testing.T) {
+		cloner := NewGitCloner()
+		_, err := cloner.CloneGitRepository("not-a-valid-url", io.Discard)
+		if err == nil {
+			t.Fatal("Expected an error for an invalid URL, but got nil")
+		}
+	})
+}
+
+func TestCloneGitRepository_Ugly(t *testing.T) {
+	t.Run("Empty repository", func(t *testing.T) {
+		bareRepoPath, err := os.MkdirTemp("", "empty-bare-repo-")
+		if err != nil {
+			t.Fatalf("Failed to create temp dir: %v", err)
+		}
+		defer os.RemoveAll(bareRepoPath)
+		runCmd(t, bareRepoPath, "git", "init", "--bare")
+
+		cloner := NewGitCloner()
+		dn, err := cloner.CloneGitRepository("file://"+bareRepoPath, io.Discard)
+		if err != nil {
+			t.Fatalf("CloneGitRepository failed on empty repo: %v", err)
+		}
+		if dn == nil {
+			t.Fatal("Expected a non-nil datanode for an empty repo")
+		}
+		// You might want to check if the datanode is empty, but for now, just checking for no error is enough.
+	})
+}

pkg/website/website.go

@@ -168,7 +168,11 @@ func (d *Downloader) getRelativePath(pageURL string) string {
 	if err != nil {
 		return ""
 	}
-	return strings.TrimPrefix(u.Path, "/")
+	path := strings.TrimPrefix(u.Path, "/")
+	if path == "" {
+		return "index.html"
+	}
+	return path
 }
 
 func (d *Downloader) resolveURL(base, ref string) (string, error) {

pkg/website/website_test.go

@@ -1,78 +1,31 @@
 package website
 
 import (
+	"fmt"
+	"io"
+	"io/fs"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/schollz/progressbar/v3"
 )
 
-func TestDownloadAndPackageWebsite(t *testing.T) {
+// --- Test Cases ---
-	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
-		switch r.URL.Path {
+func TestDownloadAndPackageWebsite_Good(t *testing.T) {
-		case "/":
+	server := newWebsiteTestServer()
-			w.Header().Set("Content-Type", "text/html")
-			w.Write([]byte(`
-				<!DOCTYPE html>
-				<html>
-				<head>
-					<title>Test Website</title>
-					<link rel="stylesheet" href="style.css">
-				</head>
-				<body>
-					<h1>Hello, Website!</h1>
-					<a href="/page2.html">Page 2</a>
-					<img src="image.png">
-				</body>
-				</html>
-			`))
-		case "/style.css":
-			w.Header().Set("Content-Type", "text/css")
-			w.Write([]byte(`body { color: red; }`))
-		case "/image.png":
-			w.Header().Set("Content-Type", "image/png")
-			w.Write([]byte("fake image data"))
-		case "/page2.html":
-			w.Header().Set("Content-Type", "text/html")
-			w.Write([]byte(`
-				<!DOCTYPE html>
-				<html>
-				<head>
-					<title>Page 2</title>
-				</head>
-				<body>
-					<h1>Page 2</h1>
-					<a href="/page3.html">Page 3</a>
-				</body>
-				</html>
-			`))
-		case "/page3.html":
-			w.Header().Set("Content-Type", "text/html")
-			w.Write([]byte(`
-				<!DOCTYPE html>
-				<html>
-				<head>
-					<title>Page 3</title>
-				</head>
-				<body>
-					<h1>Page 3</h1>
-				</body>
-				</html>
-			`))
-		default:
-			http.NotFound(w, r)
-		}
-	}))
 	defer server.Close()
 
-	bar := progressbar.New(1)
+	bar := progressbar.NewOptions(1, progressbar.OptionSetWriter(io.Discard))
 	dn, err := DownloadAndPackageWebsite(server.URL, 2, bar)
 	if err != nil {
 		t.Fatalf("DownloadAndPackageWebsite failed: %v", err)
 	}
 
-	expectedFiles := []string{"", "style.css", "image.png", "page2.html", "page3.html"}
+	expectedFiles := []string{"index.html", "style.css", "image.png", "page2.html", "page3.html"}
 	for _, file := range expectedFiles {
 		exists, err := dn.Exists(file)
 		if err != nil {
@@ -82,4 +35,172 @@ func TestDownloadAndPackageWebsite(t *testing.T) {
 			t.Errorf("Expected to find file %s in DataNode, but it was not found", file)
 		}
 	}
+
+	// Check content of one file
+	file, err := dn.Open("style.css")
+	if err != nil {
+		t.Fatalf("Failed to open style.css: %v", err)
+	}
+	content, err := io.ReadAll(file)
+	if err != nil {
+		t.Fatalf("Failed to read style.css: %v", err)
+	}
+	if string(content) != `body { color: red; }` {
+		t.Errorf("Unexpected content for style.css: %s", content)
+	}
+}
+
+func TestDownloadAndPackageWebsite_Bad(t *testing.T) {
+	t.Run("Invalid Start URL", func(t *testing.T) {
+		_, err := DownloadAndPackageWebsite("http://invalid-url", 1, nil)
+		if err == nil {
+			t.Fatal("Expected an error for an invalid start URL, but got nil")
+		}
+	})
+
+	t.Run("Server Error on Start URL", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+		}))
+		defer server.Close()
+		_, err := DownloadAndPackageWebsite(server.URL, 1, nil)
+		if err == nil {
+			t.Fatal("Expected an error for a server error on the start URL, but got nil")
+		}
+	})
+
+	t.Run("Broken Link", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/" {
+				w.Header().Set("Content-Type", "text/html")
+				fmt.Fprint(w, `<a href="/broken.html">Broken</a>`)
+			} else {
+				http.NotFound(w, r)
+			}
+		}))
+		defer server.Close()
+		// We expect an error because the link is broken.
+		dn, err := DownloadAndPackageWebsite(server.URL, 1, nil)
+		if err == nil {
+			t.Fatal("Expected an error for a broken link, but got nil")
+		}
+		if !strings.Contains(err.Error(), "404 Not Found") {
+			t.Errorf("Expected error to contain '404 Not Found', but got: %v", err)
+		}
+		if dn != nil {
+			t.Error("DataNode should be nil on error")
+		}
+	})
+}
+
+func TestDownloadAndPackageWebsite_Ugly(t *testing.T) {
+	t.Run("Exceed Max Depth", func(t *testing.T) {
+		server := newWebsiteTestServer()
+		defer server.Close()
+
+		bar := progressbar.NewOptions(1, progressbar.OptionSetWriter(io.Discard))
+		dn, err := DownloadAndPackageWebsite(server.URL, 1, bar) // Max depth of 1
+		if err != nil {
+			t.Fatalf("DownloadAndPackageWebsite failed: %v", err)
+		}
+
+		// page3.html is at depth 2, so it should not be present.
+		exists, _ := dn.Exists("page3.html")
+		if exists {
+			t.Error("page3.html should not have been downloaded due to max depth")
+		}
+		// page2.html is at depth 1, so it should be present.
+		exists, _ = dn.Exists("page2.html")
+		if !exists {
+			t.Error("page2.html should have been downloaded")
+		}
+	})
+
+	t.Run("External Links", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "text/html")
+			fmt.Fprint(w, `<a href="http://externalsite.com/page.html">External</a>`)
+		}))
+		defer server.Close()
+		dn, err := DownloadAndPackageWebsite(server.URL, 1, nil)
+		if err != nil {
+			t.Fatalf("DownloadAndPackageWebsite failed: %v", err)
+		}
+		if dn == nil {
+			t.Fatal("DataNode should not be nil")
+		}
+		// We can't easily check if the external link was visited, but we can ensure
+		// it didn't cause an error and didn't add any unexpected files.
+		var fileCount int
+		dn.Walk(".", func(path string, d fs.DirEntry, err error) error {
+			if !d.IsDir() {
+				fileCount++
+			}
+			return nil
+		})
+		if fileCount != 1 { // Should only contain the root page
+			t.Errorf("expected 1 file in datanode, but found %d", fileCount)
+		}
+	})
+
+	t.Run("Timeout", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			time.Sleep(100 * time.Millisecond)
+			w.Header().Set("Content-Type", "text/html")
+			fmt.Fprint(w, `<h1>Hello</h1>`)
+		}))
+		defer server.Close()
+		// This test is tricky as it depends on timing.
+		// The current implementation uses the default http client with no timeout.
+		// A proper implementation would allow configuring a timeout.
+		// For now, we'll just test that it doesn't hang forever.
+		done := make(chan bool)
+		go func() {
+			_, err := DownloadAndPackageWebsite(server.URL, 1, nil)
+			if err != nil && !strings.Contains(err.Error(), "context deadline exceeded") {
+				// We expect a timeout error, but other errors are failures.
+				t.Errorf("unexpected error: %v", err)
+			}
+			done <- true
+		}()
+		select {
+		case <-done:
+			// test finished
+		case <-time.After(5 * time.Second):
+			t.Fatal("Test timed out")
+		}
+	})
+}
+
+// --- Helpers ---
+
+func newWebsiteTestServer() *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/":
+			w.Header().Set("Content-Type", "text/html")
+			fmt.Fprint(w, `
+				<!DOCTYPE html>
+				<html><body>
+					<a href="/page2.html">Page 2</a>
+					<link rel="stylesheet" href="style.css">
+					<img src="image.png">
+				</body></html>
+			`)
+		case "/style.css":
+			w.Header().Set("Content-Type", "text/css")
+			fmt.Fprint(w, `body { color: red; }`)
+		case "/image.png":
+			w.Header().Set("Content-Type", "image/png")
+			fmt.Fprint(w, "fake image data")
+		case "/page2.html":
+			w.Header().Set("Content-Type", "text/html")
+			fmt.Fprint(w, `<html><body><a href="/page3.html">Page 3</a></body></html>`)
+		case "/page3.html":
+			w.Header().Set("Content-Type", "text/html")
+			fmt.Fprint(w, `<html><body><h1>Page 3</h1></body></html>`)
+		default:
+			http.NotFound(w, r)
+		}
+	}))
 }