From 27298fe5a9f87796607f594e3bc48ffc00bc7c52 Mon Sep 17 00:00:00 2001
From: Raul Jordan <raul@prysmaticlabs.com>
Date: Mon, 15 Jul 2024 11:11:25 -0500
Subject: [PATCH] test pass

---
 staker/challenge-cache/cache.go      |  8 +++++++-
 staker/challenge-cache/cache_test.go | 24 ++++++------------------
 2 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/staker/challenge-cache/cache.go b/staker/challenge-cache/cache.go
index c2627247a1..ed4fad6450 100644
--- a/staker/challenge-cache/cache.go
+++ b/staker/challenge-cache/cache.go
@@ -177,7 +177,10 @@ func (c *Cache) Prune(ctx context.Context, messageNumber uint64) error {
 	messageNumPattern := fmt.Sprintf(`%s-(\d+)-`, messageNumberPrefix)
 	pattern := regexp.MustCompile(messageNumPattern)
 	pathsToDelete := make([]string, 0)
-	if err := filepath.Walk(c.baseDir, func(path string, info os.FileInfo, err error) error {
+	if err := filepath.WalkDir(c.baseDir, func(path string, info os.DirEntry, err error) error {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
 		if err != nil {
 			return err
 		}
@@ -201,6 +204,9 @@ func (c *Cache) Prune(ctx context.Context, messageNumber uint64) error {
 	// We delete separately from collecting the paths, as deleting while walking
 	// a dir can cause issues with the filepath.Walk function.
 	for _, path := range pathsToDelete {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
 		if err := os.RemoveAll(path); err != nil {
 			return fmt.Errorf("could not prune directory with path %s: %w", path, err)
 		}
diff --git a/staker/challenge-cache/cache_test.go b/staker/challenge-cache/cache_test.go
index c0d380fb12..af0a058f78 100644
--- a/staker/challenge-cache/cache_test.go
+++ b/staker/challenge-cache/cache_test.go
@@ -18,16 +18,12 @@ import (
 var _ HistoryCommitmentCacher = (*Cache)(nil)
 
 func TestCache(t *testing.T) {
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	basePath := t.TempDir()
 	if err := os.MkdirAll(basePath, os.ModePerm); err != nil {
 		t.Fatal(err)
 	}
-	t.Cleanup(func() {
-		if err := os.RemoveAll(basePath); err != nil {
-			t.Fatal(err)
-		}
-	})
 	cache, err := New(basePath)
 	if err != nil {
 		t.Fatal(err)
@@ -75,16 +71,12 @@ func TestCache(t *testing.T) {
 }
 
 func TestPrune(t *testing.T) {
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	basePath := t.TempDir()
 	if err := os.MkdirAll(basePath, os.ModePerm); err != nil {
 		t.Fatal(err)
 	}
-	t.Cleanup(func() {
-		if err := os.RemoveAll(basePath); err != nil {
-			t.Fatal(err)
-		}
-	})
 	cache, err := New(basePath)
 	if err != nil {
 		t.Fatal(err)
@@ -429,17 +421,13 @@ func Test_determineFilePath(t *testing.T) {
 }
 
 func BenchmarkCache_Read_32Mb(b *testing.B) {
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	b.StopTimer()
 	basePath := os.TempDir()
 	if err := os.MkdirAll(basePath, os.ModePerm); err != nil {
 		b.Fatal(err)
 	}
-	b.Cleanup(func() {
-		if err := os.RemoveAll(basePath); err != nil {
-			b.Fatal(err)
-		}
-	})
 	cache, err := New(basePath)
 	if err != nil {
 		b.Fatal(err)