Fix tests to mock partitionGroup

Author: danielblando

pkg/compactor/compactor_paritioning_test.go

@@ -1299,7 +1299,24 @@ func TestPartitionCompactor_ShouldCompactOnlyShardsOwnedByTheInstanceOnShardingE
 
 			// Get all of the unique group hashes so that they can be used to ensure all groups were compacted
 			groupHashes[groupHash]++
-			bucketClient.MockGet(userID+"/partitioned-groups/"+fmt.Sprint(groupHash)+".json", "", nil)
+
+			// Create mock partitioned group info for the new validation check
+			partitionedGroupInfo := PartitionedGroupInfo{
+				PartitionedGroupID: groupHash,
+				PartitionCount:     1,
+				Partitions: []Partition{
+					{
+						PartitionID: 0,
+						Blocks:      []ulid.ULID{ulid.MustParse(blockID)},
+					},
+				},
+				RangeStart:   blockTimes["startTime"],
+				RangeEnd:     blockTimes["endTime"],
+				CreationTime: time.Now().Unix(),
+				Version:      PartitionedGroupInfoVersion1,
+			}
+			partitionedGroupInfoContent, _ := json.Marshal(partitionedGroupInfo)
+			bucketClient.MockGet(userID+"/partitioned-groups/"+fmt.Sprint(groupHash)+".json", string(partitionedGroupInfoContent), nil)
 			bucketClient.MockUpload(userID+"/partitioned-groups/"+fmt.Sprint(groupHash)+".json", nil)
 		}
 

pkg/compactor/partition_compaction_planner_test.go

@@ -289,9 +289,24 @@ func TestPartitionCompactionPlanner_Plan(t *testing.T) {
 				Version:            PartitionVisitMarkerVersion1,
 			}
 			visitMarkerFileContent, _ := json.Marshal(visitMarker)
+			// Mock partition group info for race condition fix
+			partitionedGroupInfo := PartitionedGroupInfo{
+				PartitionedGroupID: partitionedGroupID,
+				PartitionCount:     1,
+				Partitions: []Partition{
+					{PartitionID: partitionID, Blocks: []ulid.ULID{}},
+				},
+				RangeStart:   0,
+				RangeEnd:     2 * time.Hour.Milliseconds(),
+				CreationTime: time.Now().Unix(),
+				Version:      PartitionedGroupInfoVersion1,
+			}
+			partitionedGroupContent, _ := json.Marshal(partitionedGroupInfo)
+			partitionedGroupFile := GetPartitionedGroupFile(partitionedGroupID)
+
 			bkt.MockGet(visitMarkerFile, string(visitMarkerFileContent), nil)
+			bkt.MockGet(partitionedGroupFile, string(partitionedGroupContent), nil)
 			bkt.MockUpload(mock.Anything, nil)
-			bkt.MockGet(mock.Anything, "", nil)
 
 			registerer := prometheus.NewPedanticRegistry()
 
@@ -316,9 +331,10 @@ func TestPartitionCompactionPlanner_Plan(t *testing.T) {
 			)
 			actual, err := p.Plan(context.Background(), testData.blocks, nil, &cortextsdb.CortexMetaExtensions{
 				PartitionInfo: &cortextsdb.PartitionInfo{
-					PartitionCount:     1,
-					PartitionID:        partitionID,
-					PartitionedGroupID: partitionedGroupID,
+					PartitionCount:               1,
+					PartitionID:                  partitionID,
+					PartitionedGroupID:           partitionedGroupID,
+					PartitionedGroupCreationTime: partitionedGroupInfo.CreationTime,
 				},
 			})