Refactor integration test

Author: Barry-Zou

src/test/java/com/squareup/square/integration/LaborTest.java

@@ -12,11 +12,14 @@
 import com.squareup.square.labor.types.SearchShiftsRequest;
 import com.squareup.square.labor.types.UpdateBreakTypeRequest;
 import com.squareup.square.labor.types.UpdateShiftRequest;
+import com.squareup.square.types.Location;
+import com.squareup.square.types.SearchTeamMembersRequest;
+import com.squareup.square.types.SearchTeamMembersQuery;
+import com.squareup.square.types.SearchTeamMembersFilter;
+import com.squareup.square.types.TeamMemberStatus;
 import com.squareup.square.types.BreakType;
 import com.squareup.square.types.CreateBreakTypeResponse;
 import com.squareup.square.types.CreateShiftResponse;
-import com.squareup.square.types.CreateTeamMemberRequest;
-import com.squareup.square.types.CreateTeamMemberResponse;
 import com.squareup.square.types.DeleteBreakTypeResponse;
 import com.squareup.square.types.DeleteShiftResponse;
 import com.squareup.square.types.GetBreakTypeResponse;
@@ -28,8 +31,11 @@
 import com.squareup.square.types.UpdateBreakTypeResponse;
 import com.squareup.square.types.UpdateShiftResponse;
 import com.squareup.square.types.WorkweekConfig;
+
 import java.time.OffsetDateTime;
 import java.time.format.DateTimeFormatter;
+import java.util.Collections;
+import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
 import org.junit.jupiter.api.AfterEach;
@@ -47,21 +53,31 @@ public class LaborTest {
     @BeforeEach
     public void before() {
         client = TestClientFactory.create();
-        locationId = Helpers.createLocation(client);
-
-        // Create team member for testing
-        CreateTeamMemberResponse teamResponse = client.teamMembers()
-                .create(CreateTeamMemberRequest.builder()
-                        .idempotencyKey(UUID.randomUUID().toString())
-                        .teamMember(TeamMember.builder()
-                                .givenName("Sherlock")
-                                .familyName("Holmes")
-                                .build())
-                        .build());
-        if (!teamResponse.getTeamMember().get().getId().isPresent()) {
-            throw new RuntimeException("Failed to create team member.");
+        
+        // Get first available location
+        List<Location> locations = client.locations().list().getLocations()
+            .orElseThrow(() -> new RuntimeException("No locations available"));
+        if (locations.isEmpty()) {
+            throw new RuntimeException("No locations available for testing");
+        }
+        locationId = locations.get(0).getId().orElseThrow(() -> new RuntimeException("Location ID not present"));
+        
+        // Get first available team member at this location
+        List<TeamMember> teamMembers = client.teamMembers().search(
+            SearchTeamMembersRequest.builder()
+                .query(SearchTeamMembersQuery.builder()
+                    .filter(SearchTeamMembersFilter.builder()
+                        .locationIds(Collections.singletonList(locationId))
+                        .status(TeamMemberStatus.ACTIVE)
+                        .build())
+                    .build())
+                .build()
+        ).getTeamMembers();
+        
+        if (teamMembers == null || teamMembers.isEmpty()) {
+            throw new RuntimeException("No team members available at location " + locationId);
         }
-        memberId = teamResponse.getTeamMember().get().getId().get();
+        memberId = teamMembers.get(0).getId().orElseThrow(() -> new RuntimeException("Team member ID not present"));
 
         // Create break type for testing
         CreateBreakTypeResponse breakResponse = client.labor()
@@ -210,25 +226,14 @@ public void testUpdateShift() {
 
     @Test
     public void testDeleteShift() {
-        // create team member
-        CreateTeamMemberResponse teamMemberResponse = client.teamMembers()
-                .create(CreateTeamMemberRequest.builder()
-                        .idempotencyKey(UUID.randomUUID().toString())
-                        .teamMember(TeamMember.builder()
-                                .givenName("Sherlock")
-                                .familyName("Holmes")
-                                .build())
-                        .build());
-
-        // create shift
+        // Create shift
         CreateShiftResponse shiftResponse = client.labor()
                 .shifts()
                 .create(CreateShiftRequest.builder()
                         .shift(Shift.builder()
                                 .locationId(locationId)
                                 .startAt(OffsetDateTime.now().format(DateTimeFormatter.ISO_INSTANT))
-                                .teamMemberId(
-                                        teamMemberResponse.getTeamMember().get().getId())
+                                .teamMemberId(memberId)
                                 .build())
                         .idempotencyKey(UUID.randomUUID().toString())
                         .build());
@@ -239,16 +244,16 @@ public void testDeleteShift() {
         if (!shiftResponse.getShift().get().getId().isPresent()) {
             throw new RuntimeException("Shift ID is null.");
         }
-        shiftId = shiftResponse.getShift().get().getId().get();
+        String testShiftId = shiftResponse.getShift().get().getId().get();
         DeleteShiftResponse response = client.labor()
                 .shifts()
-                .delete(DeleteShiftsRequest.builder().id(shiftId).build());
+                .delete(DeleteShiftsRequest.builder().id(testShiftId).build());
         Assertions.assertNotNull(response);
     }
 
     @Test
     public void testDeleteBreakType() {
-        // create break type
+        // Create break type
         CreateBreakTypeResponse breakResponse = client.labor()
                 .breakTypes()
                 .create(CreateBreakTypeRequest.builder()
@@ -268,11 +273,11 @@ public void testDeleteBreakType() {
         if (!breakType.get().getId().isPresent()) {
             throw new RuntimeException("Break ID is null.");
         }
-        breakId = breakType.get().getId().get();
+        String testBreakId = breakType.get().getId().get();
 
         DeleteBreakTypeResponse response = client.labor()
                 .breakTypes()
-                .delete(DeleteBreakTypesRequest.builder().id(breakId).build());
+                .delete(DeleteBreakTypesRequest.builder().id(testBreakId).build());
         Assertions.assertNotNull(response);
     }
 
@@ -283,4 +288,4 @@ public void testListWorkweekConfigs() {
                 .list(ListWorkweekConfigsRequest.builder().build());
         Assertions.assertFalse(response.getItems().isEmpty());
     }
-}
+}