cleanup and fix tests

Author: idbentley

cmd/customcode.go

@@ -48,24 +48,6 @@ var registerCustomCodeCmd = &model.ExecutableCommand[RegisterCustomCodeFlags]{
 			Shorthand:   "a",
 			Description: "apply existing custom code patches without running generation",
 		},
-		flag.BooleanFlag{
-			Name:		 "latest-hash",
-			Description: "show the latest commit hash from gen.lock that contains custom code changes",
-		},
-		flag.StringFlag{
-			Name:        "installationURL",
-			Shorthand:   "i",
-			Description: "the language specific installation URL for installation instructions if the SDK is not published to a package manager",
-		},
-		flag.MapFlag{
-			Name:        "installationURLs",
-			Description: "a map from target ID to installation URL for installation instructions if the SDK is not published to a package manager",
-		},
-		flag.StringFlag{
-			Name:        "repo",
-			Shorthand:   "r",
-			Description: "the repository URL for the SDK, if the published (-p) flag isn't used this will be used to generate installation instructions",
-		},
 		flag.EnumFlag{
 			Name:          "output",
 			Shorthand:     "o",
@@ -104,18 +86,12 @@ func registerCustomCode(ctx context.Context, flags RegisterCustomCodeFlags) erro
 		if err != nil {
 			return fmt.Errorf("Could not find workflow file")
 		}
-		for targetName, target := range wf.Targets {
-			fmt.Println("Applying target ", targetName)
+		for _, target := range wf.Targets {
 			registercustomcode.ApplyCustomCodePatch(ctx, target)
 		}
 		return nil
 	}
 
-	// If --latest-hash flag is provided, show the commit hash from gen.lock
-	if flags.LatestHash {
-		return registercustomcode.ShowLatestCommitHash(ctx)
-	}
-
 	// Call the registercustomcode functionality
 	return registercustomcode.RegisterCustomCode(ctx, func(targetName string) error {
 		opts := []run.Opt{

integration/customcode_multitarget_test.go

@@ -32,7 +32,7 @@ func TestMultiTargetCustomCode(t *testing.T) {
 		testMultiTargetIncrementalCustomCode(t, speakeasyBinary)
 	})
 
-	t.Run("ConflictResolutionAcceptOurs", func(t *testing.T) {
+	t.Run("ConflictResolutionAcceptOurs1", func(t *testing.T) {
 		t.Parallel()
 		testMultiTargetCustomCodeConflictResolutionAcceptOurs(t, speakeasyBinary)
 	})
@@ -296,7 +296,7 @@ func testMultiTargetCustomCodeConflictResolutionAcceptOurs(t *testing.T, speakea
 
 	// Step 1: Add custom code to ALL targets
 	modifyLineInFileByPrefix(t, goFilePath, "// The name that needs to be", "\t// custom code in go target")
-	modifyLineInFileByPrefix(t, tsFilePath, "* @deprecated This namespace", "// custom code in typescript target")
+	modifyLineInFileByPrefix(t, tsFilePath, "* @internal", "// custom code in typescript target")
 
 	// Step 2: Register custom code for all targets
 	customCodeCmd := exec.Command(speakeasyBinary, "customcode", "--output", "console")
@@ -593,7 +593,7 @@ func testMultiTargetCustomCodeConflictResolutionAcceptTheirs(t *testing.T, speak
 
 	// Step 1: Add custom code to ALL targets
 	modifyLineInFileByPrefix(t, goFilePath, "// The name that needs to be", "\t// custom code in go target")
-	modifyLineInFileByPrefix(t, tsFilePath, "* @deprecated This namespace", "// custom code in typescript target")
+	modifyLineInFileByPrefix(t, tsFilePath, "* @internal", "// custom code in typescript target")
 
 	// Step 2: Register custom code for all targets
 	customCodeCmd := exec.Command(speakeasyBinary, "customcode", "--output", "console")

integration/customcode_singletarget_test.go

@@ -595,38 +595,52 @@ func ValidateUserID(id int64) bool {
 	gitAddOutput, gitAddErr := gitAddCmd.CombinedOutput()
 	require.NoError(t, gitAddErr, "git add should succeed: %s", string(gitAddOutput))
 
-	// Register custom code
+	// Try to register custom code with new file - should fail
 	customCodeCmd := exec.Command(speakeasyBinary, "customcode", "--output", "console")
 	customCodeCmd.Dir = temp
 	customCodeOutput, customCodeErr := customCodeCmd.CombinedOutput()
-	require.NoError(t, customCodeErr, "customcode command should succeed: %s", string(customCodeOutput))
+	require.Error(t, customCodeErr, "customcode command should fail with new file: %s", string(customCodeOutput))
+	require.Contains(t, string(customCodeOutput), "Cannot register new files through customcode", "Error should mention new files not supported")
 
-	// Verify patch file was created
+	// Commit the new file to make it part of the codebase
+	gitCommitCmd := exec.Command("git", "commit", "-m", "Add helper file")
+	gitCommitCmd.Dir = temp
+	_, err = gitCommitCmd.CombinedOutput()
+	require.NoError(t, err, "git commit should succeed")
+
+	// Now make a modification to an existing file to test normal customcode flow
+	getUserByNamePath := filepath.Join(temp, "models", "operations", "getuserbyname.go")
+	registerCustomCodeByPrefix(t, speakeasyBinary, temp, getUserByNamePath, "// The name that needs to be fetched", "\t// custom code in existing file")
+
+	// Verify patch file was created for the modification
 	patchFile := filepath.Join(temp, ".speakeasy", "patches", "custom-code.diff")
 	_, err = os.Stat(patchFile)
-	require.NoError(t, err, "patch file should exist at %s", patchFile)
+	require.NoError(t, err, "patch file should exist for file modifications")
 
 	// Verify the file exists after registration (before regeneration)
 	_, err = os.Stat(helperFilePath)
 	require.NoError(t, err, "Helper file should exist after registration")
 
 	// Run speakeasy run to regenerate the SDK
-	// This should apply the patch and preserve the new file
+	// This should apply the patch for the existing file modification and preserve the committed helper file
 	runRegeneration(t, speakeasyBinary, temp, true)
 
-	// Verify the new file still exists after regeneration
+	// Verify the helper file still exists after regeneration (it was committed, so should be preserved)
 	_, err = os.Stat(helperFilePath)
 	require.NoError(t, err, "Helper file should exist after regeneration")
 
-	// Verify the file contents are preserved exactly
+	// Verify the helper file contents are preserved exactly
 	verifyCustomCodePresent(t, helperFilePath, "FormatUserID")
 	verifyCustomCodePresent(t, helperFilePath, "ValidateUserID")
 	verifyCustomCodePresent(t, helperFilePath, "package utils")
 
-	// Read the entire file and verify exact content match
+	// Read the entire helper file and verify exact content match
 	actualContent, err := os.ReadFile(helperFilePath)
 	require.NoError(t, err, "Failed to read helper file after regeneration")
 	require.Equal(t, helperFileContent, string(actualContent), "Helper file content should be preserved exactly")
+
+	// Verify the custom code modification was applied to the existing file
+	verifyCustomCodePresent(t, getUserByNamePath, "// custom code in existing file")
 }
 
 // testCustomCodeNewFileDeletion tests that deleting a custom file is properly registered and persisted
@@ -659,44 +673,38 @@ func FormatUserID(id int64) string {
 	_, err = gitAddCmd.CombinedOutput()
 	require.NoError(t, err, "git add should succeed")
 
-	// Register custom code (registers the new file)
-	customCodeCmd := exec.Command(speakeasyBinary, "customcode", "--output", "console")
-	customCodeCmd.Dir = temp
-	customCodeOutput, customCodeErr := customCodeCmd.CombinedOutput()
-	require.NoError(t, customCodeErr, "customcode command should succeed: %s", string(customCodeOutput))
-
-	// Verify patch file was created
-	patchFile := filepath.Join(temp, ".speakeasy", "patches", "custom-code.diff")
-	_, err = os.Stat(patchFile)
-	require.NoError(t, err, "patch file should exist after registering new file")
-
-	// Regenerate and verify the file is preserved
-	runRegeneration(t, speakeasyBinary, temp, true)
-	_, err = os.Stat(helperFilePath)
-	require.NoError(t, err, "Helper file should exist after first regeneration")
-
-	// Commit the regeneration so the file becomes part of HEAD
-	gitCommitCmd := exec.Command("git", "commit", "-am", "regeneration with custom file")
+	// Commit the new file to make it part of the codebase
+	gitCommitCmd := exec.Command("git", "commit", "-m", "Add helper file")
 	gitCommitCmd.Dir = temp
 	_, err = gitCommitCmd.CombinedOutput()
-	require.NoError(t, err, "git commit should succeed after regeneration")
+	require.NoError(t, err, "git commit should succeed")
 
 	// Now delete the file
 	err = os.Remove(helperFilePath)
 	require.NoError(t, err, "Failed to delete helper file")
 
-	// Register the deletion
-	customCodeCmd = exec.Command(speakeasyBinary, "customcode", "--output", "console")
-	customCodeCmd.Dir = temp
-	customCodeOutput, customCodeErr = customCodeCmd.CombinedOutput()
-	require.NoError(t, customCodeErr, "customcode command should succeed after deletion: %s", string(customCodeOutput))
+	// Commit the deletion
+	gitCommitDeleteCmd := exec.Command("git", "commit", "-am", "Delete helper file")
+	gitCommitDeleteCmd.Dir = temp
+	_, err = gitCommitDeleteCmd.CombinedOutput()
+	require.NoError(t, err, "git commit for deletion should succeed")
+
+	// Now make a modification to an existing file to test normal customcode flow
+	getUserByNamePath := filepath.Join(temp, "models", "operations", "getuserbyname.go")
+	registerCustomCodeByPrefix(t, speakeasyBinary, temp, getUserByNamePath, "// The name that needs to be fetched", "\t// custom code for deletion test")
 
-	// Verify patch file was removed (no custom code remaining)
+	// Verify patch file was created for the modification
+	patchFile := filepath.Join(temp, ".speakeasy", "patches", "custom-code.diff")
 	_, err = os.Stat(patchFile)
-	require.True(t, os.IsNotExist(err), "patch file should not exist after deleting the only custom file")
+	require.NoError(t, err, "patch file should exist for file modifications")
 
-	// Regenerate and verify the file remains deleted
+	// Run regeneration to test that the custom code is preserved and deleted file stays deleted
 	runRegeneration(t, speakeasyBinary, temp, true)
+	
+	// Verify the helper file does NOT exist after regeneration (it was deleted and committed)
 	_, err = os.Stat(helperFilePath)
-	require.True(t, os.IsNotExist(err), "Helper file should not exist after regeneration with deletion registered")
+	require.True(t, os.IsNotExist(err), "Helper file should not exist after regeneration (was deleted)")
+
+	// Verify the custom code modification was applied to the existing file
+	verifyCustomCodePresent(t, getUserByNamePath, "// custom code for deletion test")
 }

internal/registercustomcode/registercustomcode.go

@@ -675,6 +675,11 @@ func isLocalPath(location workflow.LocationString) bool {
 
 // Git operations
 func captureCustomCodeDiff(outDir string, excludePaths []string) (string, error) {
+	// First, check for new files and deletions in the staged changes
+	if err := checkForNewFilesOrDeletions(outDir); err != nil {
+		return "", err
+	}
+
 	args := []string{"diff", "HEAD", outDir}
 
 	// Filter excludePaths to only include children of outDir
@@ -699,6 +704,59 @@ func captureCustomCodeDiff(outDir string, excludePaths []string) (string, error)
 	return string(combinedOutput), nil
 }
 
+// checkForNewFilesOrDeletions checks if there are any new files or file deletions
+// in the staged changes and returns an error if found, as these are not supported
+func checkForNewFilesOrDeletions(outDir string) error {
+	// Check for new files (A) and deletions (D) using git diff with filter
+	cmd := exec.Command("git", "diff", "--name-status", "--diff-filter=AD", "HEAD", outDir)
+	output, err := cmd.Output()
+	if err != nil {
+		return fmt.Errorf("failed to check for new files or deletions: %w", err)
+	}
+
+	outputStr := strings.TrimSpace(string(output))
+	if outputStr == "" {
+		// No additions or deletions found
+		return nil
+	}
+
+	lines := strings.Split(outputStr, "\n")
+	var newFiles []string
+	var deletedFiles []string
+
+	for _, line := range lines {
+		if line == "" {
+			continue
+		}
+		
+		// Format is "STATUS\tFILENAME"
+		parts := strings.SplitN(line, "\t", 2)
+		if len(parts) != 2 {
+			continue
+		}
+		
+		status := parts[0]
+		filename := parts[1]
+		
+		if status == "A" {
+			newFiles = append(newFiles, filename)
+		} else if status == "D" {
+			deletedFiles = append(deletedFiles, filename)
+		}
+	}
+
+	// Return error if we found new files or deletions
+	if len(newFiles) > 0 {
+		return fmt.Errorf("Cannot register new files through customcode. New files found: %s", strings.Join(newFiles, ", "))
+	}
+	
+	if len(deletedFiles) > 0 {
+		return fmt.Errorf("Cannot register file deletions through customcode. Deleted files found: %s", strings.Join(deletedFiles, ", "))
+	}
+
+	return nil
+}
+
 func checkForChangesWithExclusions(dir string, excludePaths []string) (bool, error) {
 	args := []string{"diff", "--cached", dir}