Updated database

- Moved functions to be members of the DB type
- Updated tests

database.go

@@ -8,59 +8,61 @@ type entry struct {
 	locations []string
 }
 
-type db struct {
+type DB struct {
 	entries map[PHash]*entry
 }
 
 var (
 	ErrNotFound = errors.New("Image not found")
 )
 
-var DB db
-
-func init() {
-	DB.entries = make(map[PHash]*entry)
+func NewDB() (db *DB) {
+	db = new(DB)
+	db.entries = make(map[PHash]*entry)
+	return
 }
 
-func AddFile(filename string) error {
-	phash, err := HashFile(filename)
-	if err != nil {
-		return err
+func (db *DB) AddFile(filename string) (phash PHash, err error) {
+	phash, err = HashFile(filename)
+	if err == nil {
+		err = db.AddLocation(filename, phash)
 	}
-	return AddLocation(filename, phash)
+	return phash, err
 }
 
-func AddLocation(location string, phash PHash) error {
-	e, found := DB.entries[phash]
+func (db *DB) AddLocation(location string, phash PHash) error {
+	e, found := db.entries[phash]
 	if !found {
 		e = new(entry)
-		DB.entries[phash] = e
+		db.entries[phash] = e
 	}
 	e.locations = append(e.locations, location)
 	return nil
 }
 
-func Find(phash PHash) ([]string, error) {
-	if entry, found := DB.entries[phash]; found {
-		return entry.locations, nil
+func (db *DB) Find(phash PHash) (locations []string, err error) {
+	if entry, found := db.entries[phash]; found {
+		locations = entry.locations
+	} else {
+		err = ErrNotFound
 	}
-	return []string{}, ErrNotFound
+
+	return locations, err
 }
 
-func SearchByFile(filename string, maxDistance uint) ([]string, error) {
-	h, e := HashFile(filename)
-	if e != nil {
-		return nil, e
+func (db *DB) SearchByFile(filename string, maxDistance uint) (matches []string, err error) {
+	h, err := HashFile(filename)
+	if err == nil {
+		matches, err = db.SearchByHash(h, maxDistance)
 	}
-
-	return SearchByHash(h, maxDistance)
+	return matches, err
 }
 
-func SearchByHash(phash PHash, maxDistance uint) ([]string, error) {
+func (db *DB) SearchByHash(phash PHash, maxDistance uint) ([]string, error) {
 	var results []string
 
 	// look for existing entry within maxDistance of the hash
-	for p, e := range DB.entries {
+	for p, e := range db.entries {
 		if p.Distance(phash) <= maxDistance {
 			results = append(results, e.locations...)
 		}

database_test.go

@@ -6,32 +6,60 @@ import (
 )
 
 func TestAddFile(t *testing.T) {
-	if len(DB.entries) > 0 {
-		t.Log("Database should start out with no entries")
-		t.Fail()
-	}
+	db := NewDB()
 
-	err := AddFile("images/ascendingGradient.png")
+	_, err := db.AddFile("images/ascendingGradient.png")
 	if err != nil {
 		t.Logf("Expected no error while adding file.  Got: %v", err)
 		t.Fail()
 	}
 
-	if len(DB.entries) != 1 {
+	if len(db.entries) != 1 {
 		t.Log("Database should contain exactly one entry")
 		t.Fail()
 	}
 
 }
 
+func TestFind(t *testing.T) {
+	db := NewDB()
+	hash1, _ := db.AddFile("images/gopher1.png")
+	hash2, _ := db.AddFile("images/gopher2.png")
+
+	entries, _ := db.Find(hash1)
+	if len(entries) != 2 {
+		t.Logf("Expected to find two entries for hash %v but only got %d", hash1, len(entries))
+		t.Fail()
+	}
+
+	entries, _ = db.Find(hash2)
+	if len(entries) != 2 {
+		t.Logf("Expected to find two entries for hash %v but only got %d", hash2, len(entries))
+		t.Fail()
+	}
+
+	hash3, _ := HashFile("images/ascendingGradient.png")
+	entries, err := db.Find(hash3)
+	if len(entries) > 0 {
+		t.Logf("Expected not to find hash %v but found %d entries", hash3, len(entries))
+		t.Fail()
+	}
+
+	if err != ErrNotFound {
+		t.Logf("Expected %v but got %v", ErrNotFound, err)
+		t.Fail()
+	}
+}
+
 func TestSearchByFile(t *testing.T) {
-	AddFile("images/ascendingGradient.png")
-	AddFile("images/descendingGradient.png")
-	AddFile("images/alternatingGradient.png")
-	AddFile("images/gopher1.png")
-	AddFile("images/gopher2.png")
+	db := NewDB()
+	db.AddFile("images/ascendingGradient.png")
+	db.AddFile("images/descendingGradient.png")
+	db.AddFile("images/alternatingGradient.png")
+	db.AddFile("images/gopher1.png")
+	db.AddFile("images/gopher2.png")
 
-	filenames, err := SearchByFile("images/gopher3.png", 5)
+	filenames, err := db.SearchByFile("images/gopher3.png", 5)
 	if err != nil {
 		t.Logf("Expected no error while searching by file.  Got: %v", err)
 		t.Fail()