Merge pull request beego#2351 from amrfaissal/add-safemap-count

Add Count method to BeeMap struct

utils/safemap.go

@@ -61,10 +61,8 @@ func (m *BeeMap) Set(k interface{}, v interface{}) bool {
 func (m *BeeMap) Check(k interface{}) bool {
 	m.lock.RLock()
 	defer m.lock.RUnlock()
-	if _, ok := m.bm[k]; !ok {
-		return false
-	}
-	return true
+	_, ok := m.bm[k]
+	return ok
 }
 
 // Delete the given key and value.
@@ -84,3 +82,10 @@ func (m *BeeMap) Items() map[interface{}]interface{} {
 	}
 	return r
 }
+
+// Count returns the number of items within the map.
+func (m *BeeMap) Count() int {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+	return len(m.bm)
+}

utils/safemap_test.go

@@ -14,25 +14,44 @@
 
 package utils
 
-import (
-	"testing"
-)
-
-func Test_beemap(t *testing.T) {
-	bm := NewBeeMap()
-	if !bm.Set("astaxie", 1) {
-		t.Error("set Error")
+import "testing"
+
+var safeMap *BeeMap
+
+func TestNewBeeMap(t *testing.T) {
+	safeMap = NewBeeMap()
+	if safeMap == nil {
+		t.Fatal("expected to return non-nil BeeMap", "got", safeMap)
+	}
+}
+
+func TestSet(t *testing.T) {
+	if ok := safeMap.Set("astaxie", 1); !ok {
+		t.Error("expected", true, "got", false)
+	}
+}
+
+func TestCheck(t *testing.T) {
+	if exists := safeMap.Check("astaxie"); !exists {
+		t.Error("expected", true, "got", false)
 	}
-	if !bm.Check("astaxie") {
-		t.Error("check err")
+}
+
+func TestGet(t *testing.T) {
+	if val := safeMap.Get("astaxie"); val.(int) != 1 {
+		t.Error("expected value", 1, "got", val)
 	}
+}
 
-	if v := bm.Get("astaxie"); v.(int) != 1 {
-		t.Error("get err")
+func TestDelete(t *testing.T) {
+	safeMap.Delete("astaxie")
+	if exists := safeMap.Check("astaxie"); exists {
+		t.Error("expected element to be deleted")
 	}
+}
 
-	bm.Delete("astaxie")
-	if bm.Check("astaxie") {
-		t.Error("delete err")
+func TestCount(t *testing.T) {
+	if count := safeMap.Count(); count != 0 {
+		t.Error("expected count to be", 0, "got", count)
 	}
 }