From f719484616266c3448fe9672cfdc3a276e9ad7b2 Mon Sep 17 00:00:00 2001
From: Elliott Sales de Andrade <quantum.analyst@gmail.com>
Date: May 15 2019 07:13:05 +0000
Subject: Add option for querying main and test simultaneously.


---

diff --git a/cmd/golist/golist.go b/cmd/golist/golist.go
index b96b8c1..1ce6bf0 100644
--- a/cmd/golist/golist.go
+++ b/cmd/golist/golist.go
@@ -61,7 +61,11 @@ func main() {
 		},
 		cli.BoolFlag{
 			Name:  "tests",
-			Usage: "Apply the listing options over tests",
+			Usage: "Apply the listing options over tests only",
+		},
+		cli.BoolFlag{
+			Name:  "with-tests",
+			Usage: "Apply the listing options over both the main package and the tests",
 		},
 		cli.BoolFlag{
 			Name:  "show-main",
@@ -147,7 +151,11 @@ func main() {
 		}
 
 		if c.Bool("provided") {
-			pkgs, err := collector.BuildPackageTree(c.Bool("show-main"), c.Bool("tests"))
+			if c.Bool("with-tests") && c.Bool("tests") {
+				return fmt.Errorf("Both --with-tests and --tests cannot be set at the same time")
+			}
+
+			pkgs, err := collector.BuildPackageTree(c.Bool("show-main"), c.Bool("with-tests"), c.Bool("tests"))
 			if err != nil {
 				return err
 			}
@@ -162,7 +170,11 @@ func main() {
 		}
 
 		if c.Bool("imported") {
-			pkgs, err := collector.CollectProjectDeps(c.Bool("all-deps"), c.Bool("skip-self"), c.Bool("tests"))
+			if c.Bool("with-tests") && c.Bool("tests") {
+				return fmt.Errorf("Both --with-tests and --tests cannot be set at the same time")
+			}
+
+			pkgs, err := collector.CollectProjectDeps(c.Bool("all-deps"), c.Bool("skip-self"), c.Bool("with-tests"), c.Bool("tests"))
 			if err != nil {
 				return err
 			}
diff --git a/pkg/util/util.go b/pkg/util/util.go
index 44ea6ce..70f9b2b 100644
--- a/pkg/util/util.go
+++ b/pkg/util/util.go
@@ -208,7 +208,7 @@ func (p *PackageInfoCollector) BuildArtifact() (*ProjectData, error) {
 	var err error
 
 	// Get provided packages
-	if data.Packages, err = p.BuildPackageTree(false, false); err != nil {
+	if data.Packages, err = p.BuildPackageTree(false, false, false); err != nil {
 		return nil, err
 	}
 	sort.Strings(data.Packages)
@@ -285,13 +285,15 @@ func (p *PackageInfoCollector) CollectInstalledResources() ([]string, error) {
 	return resources, nil
 }
 
-func (p *PackageInfoCollector) CollectProjectDeps(standard bool, skipSelf bool, tests bool) ([]string, error) {
+func (p *PackageInfoCollector) CollectProjectDeps(standard, skipSelf, withTests, testsOnly bool) ([]string, error) {
 	imports := make(map[string]struct{})
 
 	for _, info := range p.packageInfos {
 		var pkgImports []string
-		if tests {
+		if testsOnly {
 			pkgImports = info.TestImports
+		} else if withTests {
+			pkgImports = append(info.Imports, info.TestImports...)
 		} else {
 			pkgImports = info.Imports
 		}
@@ -331,16 +333,24 @@ func (p *PackageInfoCollector) CollectProjectDeps(standard bool, skipSelf bool, 
 	return pkgs, nil
 }
 
-func (p *PackageInfoCollector) BuildPackageTree(includeMain bool, tests bool) ([]string, error) {
+func (p *PackageInfoCollector) BuildPackageTree(includeMain, withTests, testsOnly bool) ([]string, error) {
 	// TODO(jchaloup): strip all main package unless explicitely requested
 	var entryPoints []string
-	if tests {
+	if testsOnly {
 		for p, pkgInfo := range p.packageInfos {
 			if len(pkgInfo.TestGoFiles) > 0 {
 				entryPoints = append(entryPoints, p)
 			}
 		}
 	} else {
+		if withTests {
+			for p, pkgInfo := range p.packageInfos {
+				if len(pkgInfo.TestGoFiles) > 0 {
+					entryPoints = append(entryPoints, p)
+				}
+			}
+		}
+
 		for pkgName, pkgInfo := range p.packageInfos {
 			// check package name of each file
 			var nonMainFiles []string