Fix some tests.

[lightning.git] / stats.go

diff --git a/stats.go b/stats.go
index a02fef701d6cf1dc671f2699ac46519f3ed0b15a..2f9bd4192c9ea3fc7f3a29cf10db29ec2cdc1bc8 100644 (file)
--- a/stats.go
+++ b/stats.go
@@ -1,8 +1,11 @@
-package main
+// Copyright (C) The Lightning Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package lightning
 
 import (
        "bufio"
-       "encoding/gob"
        "encoding/json"
        "errors"
        "flag"
@@ -12,16 +15,17 @@ import (
        "net/http"
        _ "net/http/pprof"
        "os"
+       "strings"
 
        "git.arvados.org/arvados.git/sdk/go/arvados"
        log "github.com/sirupsen/logrus"
 )
 
-type stats struct {
+type statscmd struct {
        debugUnplaced bool
 }
 
-func (cmd *stats) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
+func (cmd *statscmd) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
        var err error
        defer func() {
                if err != nil {
@@ -43,6 +47,9 @@ func (cmd *stats) RunCommand(prog string, args []string, stdin io.Reader, stdout
                return 0
        } else if err != nil {
                return 2
+       } else if flags.NArg() > 0 {
+               err = fmt.Errorf("errant command line arguments after parsed flags: %v", flags.Args())
+               return 2
        }
 
        if *pprof != "" {
@@ -101,7 +108,7 @@ func (cmd *stats) RunCommand(prog string, args []string, stdin io.Reader, stdout
        }
 
        bufw := bufio.NewWriter(output)
-       err = cmd.doStats(input, bufw)
+       err = cmd.doStats(input, strings.HasSuffix(*inputFilename, ".gz"), bufw)
        if err != nil {
                return 1
        }
@@ -116,7 +123,7 @@ func (cmd *stats) RunCommand(prog string, args []string, stdin io.Reader, stdout
        return 0
 }
 
-func (cmd *stats) doStats(input io.Reader, output io.Writer) error {
+func (cmd *statscmd) doStats(input io.Reader, gz bool, output io.Writer) error {
        var ret struct {
                Genomes          int
                CalledBases      []int64
@@ -131,15 +138,7 @@ func (cmd *stats) doStats(input io.Reader, output io.Writer) error {
        var tagSet [][]byte
        var tagPlacements []int
        tileVariantCalls := map[tileLibRef]int{}
-       dec := gob.NewDecoder(bufio.NewReaderSize(input, 1<<26))
-       for {
-               var ent LibraryEntry
-               err := dec.Decode(&ent)
-               if err == io.EOF {
-                       break
-               } else if err != nil {
-                       return fmt.Errorf("gob decode: %w", err)
-               }
+       err := DecodeLibrary(input, gz, func(ent *LibraryEntry) error {
                ret.Genomes += len(ent.CompactGenomes)
                ret.TileVariants += len(ent.TileVariants)
                if len(ent.TagSet) > 0 {
@@ -186,6 +185,10 @@ func (cmd *stats) doStats(input io.Reader, output io.Writer) error {
                        }
                        ret.CalledBases = append(ret.CalledBases, calledBases)
                }
+               return nil
+       })
+       if err != nil {
+               return err
        }
        for id, p := range tagPlacements {
                for len(ret.TagsPlacedNTimes) <= p {