Fixed Update method bug

Author: beeker1121

pqueue.go

@@ -228,6 +228,11 @@ func (pq *PriorityQueue) Update(item *PriorityItem, newValue []byte) error {
 		return ErrDBClosed
 	}
 
+	// Check if item exists in queue.
+	if item.ID <= pq.levels[item.Priority].head || item.ID > pq.levels[item.Priority].tail {
+		return ErrOutOfBounds
+	}
+
 	item.Value = newValue
 	return pq.db.Put(item.Key, item.Value, nil)
 }

pqueue_test.go

@@ -194,7 +194,7 @@ func TestPriorityQueueDequeueByPriority(t *testing.T) {
 	}
 
 	if pq.Length() != 50 {
-		t.Errorf("Expected queue length of 1, got %d", pq.Length())
+		t.Errorf("Expected queue length of 50, got %d", pq.Length())
 	}
 
 	deqItem, err := pq.DequeueByPriority(3)
@@ -765,6 +765,48 @@ func TestPriorityQueueUpdateObject(t *testing.T) {
 	}
 }
 
+func TestPriorityQueueUpdateOutOfBounds(t *testing.T) {
+	file := fmt.Sprintf("test_db_%d", time.Now().UnixNano())
+	pq, err := OpenPriorityQueue(file, ASC)
+	if err != nil {
+		t.Error(err)
+	}
+	defer pq.Drop()
+
+	for p := 0; p <= 4; p++ {
+		for i := 1; i <= 10; i++ {
+			item := NewPriorityItemString(fmt.Sprintf("value for item %d", i), uint8(p))
+			if err = pq.Enqueue(item); err != nil {
+				t.Error(err)
+			}
+		}
+	}
+
+	if pq.Length() != 50 {
+		t.Errorf("Expected queue length of 50, got %d", pq.Length())
+	}
+
+	deqItem, err := pq.DequeueByPriority(3)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if pq.Length() != 49 {
+		t.Errorf("Expected queue length of 49, got %d", pq.Length())
+	}
+
+	if err = pq.Update(deqItem, []byte(`new value`)); err != ErrOutOfBounds {
+		t.Errorf("Expected to get queue out of bounds error, got %s", err.Error())
+	}
+
+	deqItem.ID++
+	deqItem.Key = idToKey(deqItem.ID)
+
+	if err = pq.Update(deqItem, []byte(`new value`)); err != nil {
+		t.Error(err)
+	}
+}
+
 func TestPriorityQueueHigherPriorityAsc(t *testing.T) {
 	file := fmt.Sprintf("test_db_%d", time.Now().UnixNano())
 	pq, err := OpenPriorityQueue(file, ASC)

queue.go

@@ -153,6 +153,11 @@ func (q *Queue) Update(item *Item, newValue []byte) error {
 		return ErrDBClosed
 	}
 
+	// Check if item exists in queue.
+	if item.ID <= q.head || item.ID > q.tail {
+		return ErrOutOfBounds
+	}
+
 	item.Value = newValue
 	return q.db.Put(item.Key, item.Value, nil)
 }

queue_test.go

@@ -386,6 +386,46 @@ func TestQueueUpdateObject(t *testing.T) {
 	}
 }
 
+func TestQueueUpdateOutOfBounds(t *testing.T) {
+	file := fmt.Sprintf("test_db_%d", time.Now().UnixNano())
+	q, err := OpenQueue(file)
+	if err != nil {
+		t.Error(err)
+	}
+	defer q.Drop()
+
+	for i := 1; i <= 10; i++ {
+		item := NewItemString(fmt.Sprintf("value for item %d", i))
+		if err = q.Enqueue(item); err != nil {
+			t.Error(err)
+		}
+	}
+
+	if q.Length() != 10 {
+		t.Errorf("Expected queue length of 10, got %d", q.Length())
+	}
+
+	deqItem, err := q.Dequeue()
+	if err != nil {
+		t.Error(err)
+	}
+
+	if q.Length() != 9 {
+		t.Errorf("Expected queue length of 9, got %d", q.Length())
+	}
+
+	if err = q.Update(deqItem, []byte(`new value`)); err != ErrOutOfBounds {
+		t.Errorf("Expected to get queue out of bounds error, got %s", err.Error())
+	}
+
+	deqItem.ID++
+	deqItem.Key = idToKey(deqItem.ID)
+
+	if err = q.Update(deqItem, []byte(`new value`)); err != nil {
+		t.Error(err)
+	}
+}
+
 func TestQueueEmpty(t *testing.T) {
 	file := fmt.Sprintf("test_db_%d", time.Now().UnixNano())
 	q, err := OpenQueue(file)

stack.go

@@ -153,6 +153,11 @@ func (s *Stack) Update(item *Item, newValue []byte) error {
 		return ErrDBClosed
 	}
 
+	// Check if item exists in stack.
+	if item.ID > s.head || item.ID <= s.tail {
+		return ErrOutOfBounds
+	}
+
 	item.Value = newValue
 	return s.db.Put(item.Key, item.Value, nil)
 }

stack_test.go

@@ -105,7 +105,7 @@ func TestStackPop(t *testing.T) {
 		t.Errorf("Expected stack length of 10, got %d", s.Length())
 	}
 
-	deqItem, err := s.Pop()
+	popItem, err := s.Pop()
 	if err != nil {
 		t.Error(err)
 	}
@@ -116,8 +116,8 @@ func TestStackPop(t *testing.T) {
 
 	compStr := "value for item 10"
 
-	if deqItem.ToString() != compStr {
-		t.Errorf("Expected string to be '%s', got '%s'", compStr, deqItem.ToString())
+	if popItem.ToString() != compStr {
+		t.Errorf("Expected string to be '%s', got '%s'", compStr, popItem.ToString())
 	}
 }
 
@@ -386,6 +386,46 @@ func TestStackUpdateObject(t *testing.T) {
 	}
 }
 
+func TestStackUpdateOutOfBounds(t *testing.T) {
+	file := fmt.Sprintf("test_db_%d", time.Now().UnixNano())
+	s, err := OpenStack(file)
+	if err != nil {
+		t.Error(err)
+	}
+	defer s.Drop()
+
+	for i := 1; i <= 10; i++ {
+		item := NewItemString(fmt.Sprintf("value for item %d", i))
+		if err = s.Push(item); err != nil {
+			t.Error(err)
+		}
+	}
+
+	if s.Length() != 10 {
+		t.Errorf("Expected stack length of 10, got %d", s.Length())
+	}
+
+	popItem, err := s.Pop()
+	if err != nil {
+		t.Error(err)
+	}
+
+	if s.Length() != 9 {
+		t.Errorf("Expected stack length of 9, got %d", s.Length())
+	}
+
+	if err = s.Update(popItem, []byte(`new value`)); err != ErrOutOfBounds {
+		t.Errorf("Expected to get stack out of bounds error, got %s", err.Error())
+	}
+
+	popItem.ID--
+	popItem.Key = idToKey(popItem.ID)
+
+	if err = s.Update(popItem, []byte(`new value`)); err != nil {
+		t.Error(err)
+	}
+}
+
 func TestStackEmpty(t *testing.T) {
 	file := fmt.Sprintf("test_db_%d", time.Now().UnixNano())
 	s, err := OpenStack(file)