Added tests and fixed bucket names.

store/sha_pair.go

@@ -12,7 +12,7 @@ import (
 func (s *BoltStore) UpsertSHAPair(pair gaia.SHAPair) error {
 	return s.db.Update(func(tx *bolt.Tx) error {
 		// Get bucket
-		b := tx.Bucket(userBucket)
+		b := tx.Bucket(shaPairBucket)
 
 		// Marshal user object
 		m, err := json.Marshal(pair)
@@ -26,13 +26,13 @@ func (s *BoltStore) UpsertSHAPair(pair gaia.SHAPair) error {
 }
 
 // GetSHAPair returns a pair of shas for this pipeline run.
-func (s *BoltStore) GetSHAPair(pipelineID int) (ok bool, pair gaia.SHAPair, err error) {
+func (s *BoltStore) GetSHAPair(pipelineID []byte) (ok bool, pair gaia.SHAPair, err error) {
 	return ok, pair, s.db.View(func(tx *bolt.Tx) error {
 		// Get bucket
-		b := tx.Bucket(pipelineBucket)
+		b := tx.Bucket(shaPairBucket)
 
 		// Get pipeline
-		v := b.Get(itob(pipelineID))
+		v := b.Get(pipelineID)
 
 		// Check if we found the pipeline
 		if v == nil {

store/sha_pair_test.go

@@ -1,3 +1,104 @@
 package store
 
-// Thar be some tests here soon.
+import (
+	"bytes"
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/gaia-pipeline/gaia"
+)
+
+func TestGetSHAPAir(t *testing.T) {
+	// Create tmp folder
+	tmp, err := ioutil.TempDir("", "TestGetSHAPAir")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmp)
+
+	store := NewBoltStore()
+	gaia.Cfg.Bolt.Mode = 0600
+	err = store.Init(tmp)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer store.Close()
+
+	pair := gaia.SHAPair{}
+	pair.UniqueID = "unique-id"
+	pair.Original = []byte("original")
+	pair.Worker = []byte("worker")
+	err = store.UpsertSHAPair(pair)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ok, p, err := store.GetSHAPair([]byte(pair.UniqueID))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !ok {
+		t.Fatalf("sha pair not found")
+	}
+
+	if p.UniqueID != pair.UniqueID {
+		t.Fatalf("unique id match error. want %s got %s", pair.UniqueID, p.UniqueID)
+	}
+	if !bytes.Equal(p.Worker, pair.Worker) {
+		t.Fatalf("worker sha match error. want %s got %s", pair.Worker, p.Worker)
+	}
+	if !bytes.Equal(p.Original, pair.Original) {
+		t.Fatalf("original sha match error. want %s got %s", pair.Original, p.Original)
+	}
+}
+
+func TestUpsertSHAPair(t *testing.T) {
+	// Create tmp folder
+	tmp, err := ioutil.TempDir("", "TestUpsertSHAPair")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmp)
+
+	store := NewBoltStore()
+	gaia.Cfg.Bolt.Mode = 0600
+	err = store.Init(tmp)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer store.Close()
+
+	pair := gaia.SHAPair{}
+	pair.UniqueID = "unique-id"
+	pair.Original = []byte("original")
+	pair.Worker = []byte("worker")
+	err = store.UpsertSHAPair(pair)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Test is upsert overwrites existing records.
+	pair.Original = []byte("original2")
+	err = store.UpsertSHAPair(pair)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ok, p, err := store.GetSHAPair([]byte(pair.UniqueID))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !ok {
+		t.Fatalf("sha pair not found")
+	}
+
+	if p.UniqueID != pair.UniqueID {
+		t.Fatalf("unique id match error. want %s got %s", pair.UniqueID, p.UniqueID)
+	}
+	if !bytes.Equal(p.Worker, pair.Worker) {
+		t.Fatalf("worker sha match error. want %s got %s", pair.Worker, p.Worker)
+	}
+	if !bytes.Equal(p.Original, pair.Original) {
+		t.Fatalf("original sha match error. want %s got %s", pair.Original, p.Original)
+	}
+}

store/store.go

@@ -89,7 +89,7 @@ type GaiaStore interface {
 	WorkerDeleteAll() error
 	WorkerGet(id string) (*gaia.Worker, error)
 	UpsertSHAPair(pair gaia.SHAPair) error
-	GetSHAPair(pipelineID int) (bool, gaia.SHAPair, error)
+	GetSHAPair(pipelineID []byte) (bool, gaia.SHAPair, error)
 }
 
 // Compile time interface compliance check for BoltStore. If BoltStore

workers/agent/agent.go

@@ -405,7 +405,7 @@ func (a *Agent) scheduleWork() {
 		}
 
 		if !bytes.Equal(sha256Sum, pipelineSHA256SUM) {
-			if !a.compareSHAs(pipelineRun.PipelineID, sha256Sum, pipelineSHA256SUM) {
+			if !a.compareSHAs([]byte(pipelineRun.UniqueID), sha256Sum, pipelineSHA256SUM) {
 				gaia.Cfg.Logger.Debug("sha mismatch... attempting to re-download the binary")
 				// A possible scenario is that the pipeline has been updated and the old binary still exists here.
 				// Let us try to delete the binary and re-download the pipeline.
@@ -538,7 +538,7 @@ func (a *Agent) scheduleWork() {
 // compareSHAs compares shas of the binaries with possibly stored sha pairs. First it compares the original if they match
 // second it compares the local sha with the new one that the worker possibly rebuilt. If there is no entry,
 // we return false, because we don't know anything about the sha.
-func (a *Agent) compareSHAs(id int, sha256Sum, pipelineSHA256SUM []byte) bool {
+func (a *Agent) compareSHAs(id []byte, sha256Sum, pipelineSHA256SUM []byte) bool {
 	ok, shaPair, err := a.store.GetSHAPair(id)
 	if err != nil {
 		gaia.Cfg.Logger.Error("failed to get sha pair from memdb", "error", err.Error())

workers/agent/agent_test.go

@@ -59,7 +59,7 @@ type mockStore struct {
 func (m *mockStore) UpsertSHAPair(pair gaia.SHAPair) error {
 	return m.err
 }
-func (m *mockStore) GetSHAPair(pipelineID int) (bool, gaia.SHAPair, error) {
+func (m *mockStore) GetSHAPair(pipelineID []byte) (bool, gaia.SHAPair, error) {
 	return m.ok, m.pair, m.err
 }