From 9669878d40da9a4bfc467ff25eea5995c20937e7 Mon Sep 17 00:00:00 2001
From: Victor Vrantchan <vrancean+github@gmail.com>
Date: Sun, 19 Mar 2017 23:01:17 +0000
Subject: [PATCH] save device indexes in a separate bucket

This change makes it easier to range over the device bucket.
Also added some necessary tests for this package.
---
 device/db.go      | 28 ++++++++++++++----
 device/db_test.go | 74 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+), 5 deletions(-)
 create mode 100644 device/db_test.go

diff --git a/device/db.go b/device/db.go
index 31421e7e..13a46752 100644
--- a/device/db.go
+++ b/device/db.go
@@ -11,7 +11,13 @@ import (
 	uuid "github.com/satori/go.uuid"
 )
 
-const DeviceBucket = "mdm.Devices"
+const (
+	DeviceBucket = "mdm.Devices"
+
+	// The deviceIndexBucket index bucket stores serial number and UDID references
+	// to the device uuid.
+	deviceIndexBucket = "mdm.DeviceIdx"
+)
 
 type DB struct {
 	*bolt.DB
@@ -19,7 +25,11 @@ type DB struct {
 
 func NewDB(db *bolt.DB, sub pubsub.Subscriber) (*DB, error) {
 	err := db.Update(func(tx *bolt.Tx) error {
-		_, err := tx.CreateBucketIfNotExists([]byte(DeviceBucket))
+		_, err := tx.CreateBucketIfNotExists([]byte(deviceIndexBucket))
+		if err != nil {
+			return err
+		}
+		_, err = tx.CreateBucketIfNotExists([]byte(DeviceBucket))
 		return err
 	})
 	if err != nil {
@@ -50,18 +60,24 @@ func (db *DB) Save(dev *Device) error {
 	if err != nil {
 		return errors.Wrap(err, "marshalling device")
 	}
+
 	// store an array of indices to reference the UUID, which will be the
 	// key used to store the actual device.
 	indexes := []string{dev.UDID, dev.SerialNumber}
+	idxBucket := tx.Bucket([]byte(deviceIndexBucket))
+	if idxBucket == nil {
+		return fmt.Errorf("bucket %q not found!", deviceIndexBucket)
+	}
 	for _, idx := range indexes {
 		if idx == "" {
 			continue
 		}
 		key := []byte(idx)
-		if err := bkt.Put(key, []byte(dev.UUID)); err != nil {
+		if err := idxBucket.Put(key, []byte(dev.UUID)); err != nil {
 			return errors.Wrap(err, "put device to boltdb")
 		}
 	}
+
 	key := []byte(dev.UUID)
 	if err := bkt.Put(key, devproto); err != nil {
 		return errors.Wrap(err, "put device to boltdb")
@@ -82,7 +98,8 @@ func (db *DB) DeviceByUDID(udid string) (*Device, error) {
 	var dev Device
 	err := db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(DeviceBucket))
-		idx := b.Get([]byte(udid))
+		ib := tx.Bucket([]byte(deviceIndexBucket))
+		idx := ib.Get([]byte(udid))
 		if idx == nil {
 			return &notFound{"Device", fmt.Sprintf("udid %s", udid)}
 		}
@@ -102,7 +119,8 @@ func (db *DB) DeviceBySerial(serial string) (*Device, error) {
 	var dev Device
 	err := db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(DeviceBucket))
-		idx := b.Get([]byte(serial))
+		ib := tx.Bucket([]byte(deviceIndexBucket))
+		idx := ib.Get([]byte(serial))
 		if idx == nil {
 			return &notFound{"Device", fmt.Sprintf("serial %s", serial)}
 		}
diff --git a/device/db_test.go b/device/db_test.go
new file mode 100644
index 00000000..f670bd88
--- /dev/null
+++ b/device/db_test.go
@@ -0,0 +1,74 @@
+package device
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/boltdb/bolt"
+)
+
+func TestSave(t *testing.T) {
+	db := setupDB(t)
+	dev := &Device{
+		UUID:         "a-b-c-d",
+		UDID:         "UDID-FOO-BAR-BAZ",
+		SerialNumber: "foobarbaz",
+		ProductName:  "MacBook",
+	}
+
+	if err := db.Save(dev); err != nil {
+		t.Fatalf("saving device in datastore: %s", err)
+	}
+
+	byUDID, err := db.DeviceByUDID(dev.UDID)
+	if err != nil {
+		t.Fatalf("getting device by UDID: %s", err)
+	}
+
+	bySerial, err := db.DeviceBySerial(dev.SerialNumber)
+	if err != nil {
+		t.Fatalf("getting device by UDID: %s", err)
+	}
+
+	// test helper that verifies that the retrieved device is the same
+	tf := func(haveDev *Device) func(t *testing.T) {
+		return func(t *testing.T) {
+			if have, want := haveDev.UDID, dev.UDID; have != want {
+				t.Errorf("have %s, want %s", have, want)
+			}
+
+			if have, want := haveDev.UUID, dev.UUID; have != want {
+				t.Errorf("have %s, want %s", have, want)
+			}
+
+			if have, want := haveDev.SerialNumber, dev.SerialNumber; have != want {
+				t.Errorf("have %s, want %s", have, want)
+			}
+
+			if have, want := haveDev.ProductName, dev.ProductName; have != want {
+				t.Errorf("have %s, want %s", have, want)
+			}
+		}
+	}
+
+	t.Run("byUDID", tf(byUDID))
+	t.Run("bySerial", tf(bySerial))
+
+}
+
+func setupDB(t *testing.T) *DB {
+	f, _ := ioutil.TempFile("", "bolt-")
+	f.Close()
+	os.Remove(f.Name())
+
+	db, err := bolt.Open(f.Name(), 0777, nil)
+	if err != nil {
+		t.Fatalf("couldn't open bolt, err %s\n", err)
+	}
+	devDB, err := NewDB(db, nil)
+	if err != nil {
+		t.Fatalf("couldn't create device DB, err %s\n", err)
+	}
+	return devDB
+}