diff --git a/lib/mergeset/encoding.go b/lib/mergeset/encoding.go
index 9eab5ebd1f..b048d2970a 100644
--- a/lib/mergeset/encoding.go
+++ b/lib/mergeset/encoding.go
@@ -2,6 +2,7 @@ package mergeset
 
 import (
 	"fmt"
+	"os"
 	"sort"
 	"strings"
 	"sync"
@@ -146,7 +147,9 @@ func (ib *inmemoryBlock) MarshalUnsortedData(sb *storageBlock, firstItemDst, com
 	return ib.marshalData(sb, firstItemDst, commonPrefixDst, compressLevel)
 }
 
-var isInTest bool
+var isInTest = func() bool {
+	return strings.HasSuffix(os.Args[0], ".test")
+}()
 
 // MarshalUnsortedData marshals sorted items from ib to sb.
 //
diff --git a/lib/mergeset/table_search_test.go b/lib/mergeset/table_search_test.go
index 3673ac8ab2..3d5912b793 100644
--- a/lib/mergeset/table_search_test.go
+++ b/lib/mergeset/table_search_test.go
@@ -11,7 +11,6 @@ import (
 )
 
 func TestMain(m *testing.M) {
-	isInTest = true
 	n := m.Run()
 	os.Exit(n)
 }