github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/cmd/pyroscope/command/ci.go (about)

     1  package command
     2  
     3  import (
     4  	"bufio"
     5  	"fmt"
     6  	"os"
     7  	"path"
     8  	"path/filepath"
     9  	"strings"
    10  	"text/template"
    11  
    12  	"github.com/spf13/cobra"
    13  
    14  	"github.com/pyroscope-io/pyroscope/pkg/cli"
    15  	"github.com/pyroscope-io/pyroscope/pkg/config"
    16  )
    17  
    18  func newCiCmd(cfg *config.CI) *cobra.Command {
    19  	vpr := newViper()
    20  
    21  	var cmd *cobra.Command
    22  	cmd = &cobra.Command{
    23  		Use:    "ci",
    24  		Short:  "",
    25  		Hidden: true,
    26  		RunE: cli.CreateCmdRunFn(cfg, vpr, func(_ *cobra.Command, _ []string) error {
    27  			fmt.Println(cfg)
    28  			printUsageMessage(cmd)
    29  			return nil
    30  		}),
    31  	}
    32  
    33  	cli.PopulateFlagSet(cfg, cmd.Flags(), vpr)
    34  	cmd.AddCommand(newCiInstallCmd(cfg))
    35  
    36  	return cmd
    37  }
    38  
    39  // admin app
    40  func newCiInstallCmd(cfg *config.CI) *cobra.Command {
    41  	vpr := newViper()
    42  
    43  	var cmd *cobra.Command
    44  	cmd = &cobra.Command{
    45  		Use:    "install",
    46  		Short:  "",
    47  		Hidden: true,
    48  		RunE: cli.CreateCmdRunFn(cfg, vpr, func(_ *cobra.Command, arg []string) error {
    49  			// TODO: support passing custom paths
    50  			if len(arg) != 1 {
    51  				printUsageMessage(cmd)
    52  				return nil
    53  			}
    54  
    55  			path := arg[0]
    56  			return install(path, cfg.ApplicationName)
    57  		}),
    58  	}
    59  	cli.PopulateFlagSet(cfg, cmd.Flags(), vpr)
    60  
    61  	return cmd
    62  }
    63  
    64  func isTestFile(p string) bool {
    65  	// TODO: is this enough?
    66  	return strings.HasSuffix(p, "_test.go")
    67  }
    68  
    69  func install(basePath string, appName string) error {
    70  	type TestPackage struct {
    71  		Path        string
    72  		PackageName string
    73  	}
    74  
    75  	// Use a map to not have any duplicates
    76  	testPackages := make(map[string]TestPackage)
    77  
    78  	// Find all test files
    79  	err := filepath.Walk(basePath,
    80  		func(p string, info os.FileInfo, err error) error {
    81  			if err != nil {
    82  				return err
    83  			}
    84  			if info.IsDir() {
    85  				return nil
    86  			}
    87  
    88  			if isTestFile(p) {
    89  				dirName, _ := filepath.Split(p)
    90  				packageName := filepath.Base(dirName)
    91  
    92  				testPackages[dirName] = TestPackage{
    93  					Path:        dirName,
    94  					PackageName: packageName + "_test",
    95  				}
    96  			}
    97  
    98  			return nil
    99  		})
   100  	if err != nil {
   101  		return err
   102  	}
   103  
   104  	t1 := template.New("t1")
   105  	t1, err = t1.Parse(`
   106  // Code generated by pyroscope.
   107  package {{.PackageName }}
   108  
   109  import (
   110  	"github.com/pyroscope-io/client/pyroscope"
   111  )
   112  
   113  func init() {
   114  	pyroscope.Start(pyroscope.Config{
   115  		ProfileTypes:    []pyroscope.ProfileType{pyroscope.ProfileCPU},
   116  		ApplicationName: "{{ .ApplicationName }}",
   117  	})
   118  }
   119  `)
   120  	if err != nil {
   121  		return err
   122  	}
   123  
   124  	// For each package
   125  	// Generate a pyroscope_test.go file
   126  	for _, v := range testPackages {
   127  		testFile := path.Join(v.Path, "pyroscope_test.go")
   128  		output, err := os.Create(testFile)
   129  		if err != nil {
   130  			return err
   131  		}
   132  		w := bufio.NewWriter(output)
   133  
   134  		err = t1.Execute(w, struct {
   135  			PackageName     string
   136  			ApplicationName string
   137  		}{
   138  			PackageName:     v.PackageName,
   139  			ApplicationName: appName,
   140  		})
   141  
   142  		if err != nil {
   143  			return err
   144  		}
   145  
   146  		if err = w.Flush(); err != nil {
   147  			return err
   148  		}
   149  
   150  		fmt.Println("Created file", testFile)
   151  	}
   152  
   153  	return nil
   154  }