fix

store/store.go

@@ -519,7 +519,8 @@ func (s *Store[H]) tryAdvanceHead(ctx context.Context, headers ...H) {
 		if headers[i].Height() != currHeight+1 {
 			break
 		}
-		s.writeHead.Store(&headers[i])
+		newHead := headers[i]
+		s.writeHead.Store(&newHead)
 		currHeight++
 	}
 
@@ -528,10 +529,11 @@ func (s *Store[H]) tryAdvanceHead(ctx context.Context, headers ...H) {
 
 	// advance based on already written headers.
 	for {
-		newHead, err := s.GetByHeight(ctx, currHeight+1)
+		h, err := s.GetByHeight(ctx, currHeight+1)
 		if err != nil {
 			break
 		}
+		newHead := h
 		s.writeHead.Store(&newHead)
 		currHeight++
 	}

store/store_test.go

@@ -186,10 +186,11 @@ func TestStore_Append(t *testing.T) {
 	}
 
 	// wait for batch to be written.
-	time.Sleep(100 * time.Millisecond)
+	time.Sleep(2 * time.Second)
 
 	head, err = store.Head(ctx)
 	assert.NoError(t, err)
+	assert.Equal(t, head.Height(), headers[len(headers)-1].Height())
 	assert.Equal(t, head.Hash(), headers[len(headers)-1].Hash())
 }
 
@@ -222,6 +223,7 @@ func TestStore_Append_stableHeadWhenGaps(t *testing.T) {
 		// head is advanced to the last known header.
 		head, err := store.Head(ctx)
 		require.NoError(t, err)
+		assert.Equal(t, head.Height(), wantHead.Height())
 		assert.Equal(t, head.Hash(), wantHead.Hash())
 
 		// check that store height is aligned with the head.
@@ -253,6 +255,7 @@ func TestStore_Append_stableHeadWhenGaps(t *testing.T) {
 		// after appending missing headers we're on the latest header.
 		head, err := store.Head(ctx)
 		require.NoError(t, err)
+		assert.Equal(t, head.Height(), latestHead.Height())
 		assert.Equal(t, head.Hash(), latestHead.Hash())
 
 		// check that store height is aligned with the head.