.

wcfs/δbtail_test.go

@@ -47,64 +47,89 @@ import (
 )
 
 
+// TreeGenSrv represents connection to running `treegen ...` server.
+type TreeGenSrv struct {
+	argv      []string
+	pysrv     *exec.Cmd      // spawned `treegen ...`
+	pyin      io.WriteCloser // input to pysrv
+	pyoutRaw  io.ReadCloser  // output from pysrv
+	pyout     *bufio.Reader  // buffered ^^^
+}
 
 // TreeSrv represents connection to running `treegen trees` server.
 //
 // Create it with StartTreeSrv(zurl).
 // - Commit(treeTopology) -> tid
 type TreeSrv struct {
+	*TreeGenSrv
 	zurl	  string
-	pysrv     *exec.Cmd      // spawned `treegen trees`
-	pyin      io.WriteCloser // input to pysrv
-	pyoutRaw  io.ReadCloser  // output from pysrv
-	pyout     *bufio.Reader  // buffered ^^^
 	treeRoot  zodb.Oid       // oid of the tree treegen works on
 	head      zodb.Tid       // last made commit
 }
 
-// StartTreeGen spawns `treegen trees` server.
-func StartTreeGen(zurl string) (_ *TreeSrv, err error) {
-	defer xerr.Contextf(&err, "treegen %s: start", zurl)
+// StartTreeGenSrv spawns `treegen ...` server.
+func StartTreeGenSrv(argv ...string) (_ *TreeGenSrv, hello string, err error) {
+	defer xerr.Contextf(&err, "treesrv %v: start", argv)
 
-	// spawn `treegen trees`
-	tg := &TreeSrv{zurl: zurl}
-	tg.pysrv = exec.Command("./testprog/treegen.py", "trees", zurl)
+	// spawn `treegen ...`
+	tg := &TreeGenSrv{argv: argv}
+	tg.pysrv = exec.Command("./testprog/treegen.py", argv...)
 	tg.pyin, err = tg.pysrv.StdinPipe()
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 	tg.pyoutRaw, err = tg.pysrv.StdoutPipe()
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 	tg.pyout = bufio.NewReader(tg.pyoutRaw)
 	tg.pysrv.Stderr = os.Stderr // no redirection
 
 	err = tg.pysrv.Start()
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 
-	// wait for start message; initialize .head and .treeRoot from it
+	// wait for hello message and return it
 	defer func() {
 		if err != nil {
 			tg.Close() // ignore error
 		}
 	}()
 	defer xerr.Context(&err, "handshake")
-	start, err := tg.pyout.ReadString('\n')
+	hello, err = tg.pyout.ReadString('\n')
 	if err != nil {
 		if err == io.EOF {
 			err = io.ErrUnexpectedEOF
 		}
+		return nil, "", err
+	}
+	hello = strings.TrimSuffix(hello, "\n")
+	defer xerr.Contextf(&err, "invalid hello %q", hello)
+
+	return tg, hello, nil
+}
+
+// StartTreeSrv spawns `treegen trees` server.
+func StartTreeSrv(zurl string) (_ *TreeSrv, err error) {
+	defer xerr.Contextf(&err, "treesrv %s: start", zurl)
+
+	tgSrv, hello, err := StartTreeGenSrv("trees", zurl)
+	if err != nil {
 		return nil, err
 	}
-	start = strings.TrimSuffix(start, "\n")
-	defer xerr.Contextf(&err, "invalid start %q", start)
+
+	tg := &TreeSrv{TreeGenSrv: tgSrv, zurl: zurl}
+	defer func() {
+		if err != nil {
+			tgSrv.Close() // ignore error
+		}
+	}()
+	defer xerr.Context(&err, "handshake")
 
 	// tree.srv start @<at> tree=<root>
 	startRe := regexp.MustCompile(`^tree.srv start @([^ ]+) root=([^ ]+)$`)
-	m := startRe.FindStringSubmatch(start)
+	m := startRe.FindStringSubmatch(hello)
 	if m == nil {
 		return nil, fmt.Errorf("unexpected format")
 	}
@@ -121,8 +146,8 @@ func StartTreeGen(zurl string) (_ *TreeSrv, err error) {
 }
 
 // Close shutdowns treegen server.
-func (tg *TreeSrv) Close() (err error) {
-	defer xerr.Contextf(&err, "treegen %s: close", tg.zurl)
+func (tg *TreeGenSrv) Close() (err error) {
+	defer xerr.Contextf(&err, "treegen %v: close", tg.argv)
 
 	err1 := tg.pyin.Close()
 	err2 := tg.pyoutRaw.Close()
@@ -356,7 +381,7 @@ func testΔBTail(t *testing.T, testq chan string) {
 	defer func() {
 		err := os.RemoveAll(work); X(err)
 	}()
-	tg, err := StartTreeGen(work + "/1.fs"); X(err)
+	tg, err := StartTreeSrv(work + "/1.fs"); X(err)
 	defer func() {
 		err := tg.Close(); X(err)
 	}()