@@ -3,6 +3,7 @@ package main
33import (
44 "flag"
55 "fmt"
6+ "io"
67 "io/ioutil"
78 "log"
89 "math"
@@ -13,44 +14,86 @@ import (
1314 "strings"
1415)
1516
16- var topDescribeRE = regexp .MustCompile (`var _ = Describe\("(.+)", func\(.*` )
17+ type options struct {
18+ numChunks int
19+ printChunk int
20+ printDebug bool
21+ writer io.Writer
22+ }
1723
1824func main () {
19- var numChunks , printChunk int
20- flag .IntVar (& numChunks , "chunks" , 1 , "Number of chunks to create focus regexps for" )
21- flag .IntVar (& printChunk , "print-chunk" , 0 , "Chunk to print a regexp for" )
22- flag .Parse ()
23-
24- if printChunk >= numChunks {
25- log .Fatalf ("the chunk to print (%d) must be a smaller number than the number of chunks (%d)" , printChunk , numChunks )
25+ opts := options {
26+ writer : os .Stdout ,
2627 }
28+ flag .IntVar (& opts .numChunks , "chunks" , 1 , "Number of chunks to create focus regexps for" )
29+ flag .IntVar (& opts .printChunk , "print-chunk" , 0 , "Chunk to print a regexp for" )
30+ flag .BoolVar (& opts .printDebug , "print-debug" , false , "Print all spec prefixes in non-regexp format. Use for debugging" )
31+ flag .Parse ()
2732
2833 dir := flag .Arg (0 )
34+ if dir == "" {
35+ exitIfErr (fmt .Errorf ("test directory required as the argument" ))
36+ }
2937
3038 // Clean dir.
3139 var err error
32- if dir , err = filepath .Abs (dir ); err != nil {
33- log .Fatal (err )
34- }
40+ dir , err = filepath .Abs (dir )
41+ exitIfErr (err )
3542 wd , err := os .Getwd ()
43+ exitIfErr (err )
44+ dir , err = filepath .Rel (wd , dir )
45+ exitIfErr (err )
46+
47+ exitIfErr (opts .run (dir ))
48+ }
49+
50+ func exitIfErr (err error ) {
3651 if err != nil {
3752 log .Fatal (err )
3853 }
39- if dir , err = filepath .Rel (wd , dir ); err != nil {
40- log .Fatal (err )
54+ }
55+
56+ func (opts options ) run (dir string ) error {
57+ describes , err := findDescribes (dir )
58+ if err != nil {
59+ return err
60+ }
61+
62+ // Find minimal prefixes for all spec strings so no spec runs are duplicated across chunks.
63+ prefixes := findMinimalWordPrefixes (describes )
64+ sort .Strings (prefixes )
65+
66+ var out string
67+ if opts .printDebug {
68+ out = strings .Join (prefixes , "\n " )
69+ } else {
70+ out , err = createChunkRegexp (opts .numChunks , opts .printChunk , prefixes )
71+ if err != nil {
72+ return err
73+ }
4174 }
4275
76+ fmt .Fprint (opts .writer , out )
77+ return nil
78+ }
79+
80+ // TODO: this is hacky because top-level tests may be defined elsewise.
81+ // A better strategy would be to use the output of `ginkgo -noColor -dryRun`
82+ // like https://github.com/operator-framework/operator-lifecycle-manager/pull/1476 does.
83+ var topDescribeRE = regexp .MustCompile (`var _ = Describe\("(.+)", func\(.*` )
84+
85+ func findDescribes (dir string ) ([]string , error ) {
4386 // Find all Ginkgo specs in dir's test files.
4487 // These can be grouped independently.
4588 describeTable := make (map [string ]struct {})
4689 matches , err := filepath .Glob (filepath .Join (dir , "*_test.go" ))
4790 if err != nil {
48- log . Fatal ( err )
91+ return nil , err
4992 }
5093 for _ , match := range matches {
5194 b , err := ioutil .ReadFile (match )
5295 if err != nil {
53- log . Fatal ( err )
96+ return nil , err
5497 }
5598 specNames := topDescribeRE .FindAllSubmatch (b , - 1 )
5699 if len (specNames ) == 0 {
@@ -76,25 +119,107 @@ func main() {
76119 describes [i ] = describeKey
77120 i ++
78121 }
79- sort .Strings (describes )
122+ return describes , nil
123+ }
124+
125+ func createChunkRegexp (numChunks , printChunk int , specs []string ) (string , error ) {
126+
127+ if printChunk >= numChunks {
128+ return "" , fmt .Errorf ("the chunk to print (%d) must be a smaller number than the number of chunks (%d)" , printChunk , numChunks )
129+ }
80130
131+ numSpecs := len (specs )
132+ if numSpecs < numChunks {
133+ return "" , fmt .Errorf ("have more desired chunks (%d) than specs (%d)" , numChunks , numSpecs )
134+ }
135+
136+ // Create chunks of size ceil(number of specs/number of chunks) in alphanumeric order.
137+ // This is deterministic on inputs.
81138 chunks := make ([][]string , numChunks )
82- interval := int (math .Ceil (float64 (len ( describes ) ) / float64 (numChunks )))
139+ interval := int (math .Ceil (float64 (numSpecs ) / float64 (numChunks )))
83140 currIdx := 0
84141 for chunkIdx := 0 ; chunkIdx < numChunks ; chunkIdx ++ {
85- nextIdx := int (math .Min (float64 (currIdx + interval ), float64 (len ( describes ) )))
86- chunks [chunkIdx ] = describes [currIdx :nextIdx ]
142+ nextIdx := int (math .Min (float64 (currIdx + interval ), float64 (numSpecs )))
143+ chunks [chunkIdx ] = specs [currIdx :nextIdx ]
87144 currIdx = nextIdx
88145 }
89146
90- sb := strings.Builder {}
91- sb .WriteString ("(" )
92- sb .WriteString (chunks [printChunk ][0 ])
93- for _ , test := range chunks [printChunk ][1 :] {
94- sb .WriteString ("|" )
95- sb .WriteString (test )
147+ chunk := chunks [printChunk ]
148+ if len (chunk ) == 0 {
149+ // This is a panic because the caller may skip this error, resulting in missed test specs.
150+ panic (fmt .Sprintf ("bug: chunk %d has no elements" , printChunk ))
151+ }
152+
153+ // Write out the regexp to focus chunk specs via `ginkgo -focus <re>`.
154+ var reStr string
155+ if len (chunk ) == 1 {
156+ reStr = fmt .Sprintf ("%s .*" , chunk [0 ])
157+ } else {
158+ sb := strings.Builder {}
159+ sb .WriteString (chunk [0 ])
160+ for _ , test := range chunk [1 :] {
161+ sb .WriteString ("|" )
162+ sb .WriteString (test )
163+ }
164+ reStr = fmt .Sprintf ("(%s) .*" , sb .String ())
165+ }
166+
167+ return reStr , nil
168+ }
169+
170+ func findMinimalWordPrefixes (specs []string ) (prefixes []string ) {
171+
172+ // Create a word trie of all spec strings.
173+ t := make (wordTrie )
174+ for _ , spec := range specs {
175+ t .push (spec )
176+ }
177+
178+ // Now find the first branch point for each path in the trie by DFS.
179+ for word , node := range t {
180+ var prefixElements []string
181+ next:
182+ if word != "" {
183+ prefixElements = append (prefixElements , word )
184+ }
185+ if len (node .children ) == 1 {
186+ for nextWord , nextNode := range node .children {
187+ word , node = nextWord , nextNode
188+ }
189+ goto next
190+ }
191+ // TODO: this might need to be joined by "\s+"
192+ // in case multiple spaces were used in the spec name.
193+ prefixes = append (prefixes , strings .Join (prefixElements , " " ))
96194 }
97- sb .WriteString (").*" )
98195
99- fmt .Println (sb .String ())
196+ return prefixes
197+ }
198+
199+ // wordTrie is a trie of word nodes, instead of individual characters.
200+ type wordTrie map [string ]* wordTrieNode
201+
202+ type wordTrieNode struct {
203+ word string
204+ children map [string ]* wordTrieNode
205+ }
206+
207+ // push creates s branch of the trie from each word in s.
208+ func (t wordTrie ) push (s string ) {
209+ split := strings .Split (s , " " )
210+
211+ curr := & wordTrieNode {word : "" , children : t }
212+ for _ , sp := range split {
213+ if sp = strings .TrimSpace (sp ); sp == "" {
214+ continue
215+ }
216+ next , hasNext := curr .children [sp ]
217+ if ! hasNext {
218+ next = & wordTrieNode {word : sp , children : make (map [string ]* wordTrieNode )}
219+ curr .children [sp ] = next
220+ }
221+ curr = next
222+ }
223+ // Add termination node so "foo" and "foo bar" have a branching point of "foo".
224+ curr .children ["" ] = & wordTrieNode {}
100225}
0 commit comments