Merge pull request #45 from brotherlogic/load_staging_startup

Adds ability to load data on startup. This closes #43

datacollector.go

@@ -53,6 +53,7 @@ type Server struct {
 	readConfig *pb.ReadConfig
 	flushTime  time.Duration
 	saveTime   time.Duration
+	loadTime   time.Duration
 }
 
 // Init builds the server
@@ -64,6 +65,7 @@ func Init() *Server {
 		&pb.ReadConfig{},
 		0,
 		0,
+		0,
 	}
 	return s
 }
@@ -124,6 +126,7 @@ func (s *Server) GetState() []*pbg.State {
 		&pbg.State{Key: "collected", Value: int64(len(s.config.Data))},
 		&pbg.State{Key: "flush_time", TimeDuration: s.flushTime.Nanoseconds()},
 		&pbg.State{Key: "save_time", TimeDuration: s.saveTime.Nanoseconds()},
+		&pbg.State{Key: "load_time", TimeDuration: s.loadTime.Nanoseconds()},
 	}
 }
 
@@ -151,6 +154,12 @@ func main() {
 	server := Init()
 	server.PrepServer()
 	server.Register = server
+
+	err := server.loadData("/media/scratch/datacollector/")
+	if err != nil {
+		panic(err)
+	}
+
 	server.RegisterServer("datacollector", false)
 	server.RegisterRepeatingTask(server.collect, "collect", time.Minute*5)
 	server.RegisterRepeatingTask(server.flushToStaging, "flush_to_staging", time.Minute*30)

datacollectorintegration_test.go

@@ -0,0 +1,46 @@
+package main
+
+import (
+	"io/ioutil"
+	"testing"
+	"time"
+
+	"golang.org/x/net/context"
+
+	pb "github.com/brotherlogic/datacollector/proto"
+	pbg "github.com/brotherlogic/goserver/proto"
+	"github.com/golang/protobuf/proto"
+)
+
+func TestSaveAndLoadData(t *testing.T) {
+	s := InitTestServer()
+	tstamp := time.Now().Add(time.Hour * -1).Unix()
+	s.config.Data = append(s.config.Data, &pb.DataSet{JobName: "madeup", Identifier: "madeup", Staging: []*pb.Reading{&pb.Reading{Timestamp: tstamp, Measure: &pbg.State{Key: "blah", Value: int64(20)}, Collapsed: true}}, Readings: []*pb.Reading{&pb.Reading{Timestamp: tstamp, Measure: &pbg.State{Key: "blah", Value: int64(12)}}}})
+	data, file := s.saveData(context.Background())
+
+	if len(file) == 0 {
+		t.Errorf("No filename specified")
+	}
+
+	if len(data.Data) == 0 || len(data.Data[0].Readings) > 0 {
+		t.Errorf("Readings have not been stripped")
+	}
+
+	by, _ := proto.Marshal(data)
+	err := ioutil.WriteFile(file, by, 0644)
+
+	if err != nil {
+		t.Fatalf("Unable to write data: %v", err)
+	}
+
+	s2 := InitTestServer()
+	err = s2.loadData("")
+	if err != nil {
+		t.Fatalf("Unable to read data: %v", err)
+	}
+
+	if len(s2.config.Data) != 1 {
+		t.Errorf("Reading data has failed")
+	}
+
+}

datacollectorutils.go

@@ -3,6 +3,7 @@ package main
 import (
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"time"
 
 	"golang.org/x/net/context"
@@ -110,3 +111,19 @@ func (s *Server) saveData(ctx context.Context) (*pb.Config, string) {
 
 	return saveCopy, fmt.Sprintf("%v%v%v", time.Now().Year(), time.Now().Month(), time.Now().Day())
 }
+
+func (s *Server) loadData(dir string) error {
+	t := time.Now()
+
+	data, err := ioutil.ReadFile(dir + fmt.Sprintf("%v%v%v", time.Now().Year(), time.Now().Month(), time.Now().Day()))
+	if err != nil {
+		return err
+	}
+
+	loadCopy := &pb.Config{}
+	proto.Unmarshal(data, loadCopy)
+	s.config = loadCopy
+
+	s.loadTime = time.Now().Sub(t)
+	return nil
+}

datacollectorutils_test.go

@@ -96,18 +96,10 @@ func TestRunCollapse(t *testing.T) {
 	}
 }
 
-func TestSaveData(t *testing.T) {
+func TestLoadDataBadRead(t *testing.T) {
 	s := InitTestServer()
-	tstamp := time.Now().Add(time.Hour * -1).Unix()
-	s.config.Data = append(s.config.Data, &pb.DataSet{JobName: "madeup", Identifier: "madeup", Staging: []*pb.Reading{&pb.Reading{Timestamp: tstamp, Measure: &pbg.State{Key: "blah", Value: int64(20)}, Collapsed: true}}, Readings: []*pb.Reading{&pb.Reading{Timestamp: tstamp, Measure: &pbg.State{Key: "blah", Value: int64(12)}}}})
-
-	data, file := s.saveData(context.Background())
-
-	if len(file) == 0 {
-		t.Errorf("No filename specified")
-	}
-
-	if len(data.Data) == 0 || len(data.Data[0].Readings) > 0 {
-		t.Errorf("Readings have not been stripped")
+	err := s.loadData("madeupdirectory")
+	if err == nil {
+		t.Errorf("Bad directory did not cause error")
 	}
 }