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,32 @@ 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()
+         .orElseThrow(() -> new RuntimeException("Failed to get team members"));
+        
+        if (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,45 +227,55 @@ 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
-        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())
-                                .build())
-                        .idempotencyKey(UUID.randomUUID().toString())
-                        .build());
+        try {
+            // Create shift with both start and end time
+            OffsetDateTime now = OffsetDateTime.now();
+            CreateShiftResponse shiftResponse = client.labor()
+                    .shifts()
+                    .create(CreateShiftRequest.builder()
+                            .shift(Shift.builder()
+                                    .locationId(locationId)
+                                    .startAt(now.format(DateTimeFormatter.ISO_INSTANT))
+                                    .teamMemberId(memberId)
+                                    .endAt(now.plusHours(8).format(DateTimeFormatter.ISO_INSTANT)) // Add end time
+                                    .wage(ShiftWage.builder()
+                                            .title("Regular")
+                                            .hourlyRate(Helpers.newTestMoney(1500))
+                                            .build())
+                                    .build())
+                            .idempotencyKey(UUID.randomUUID().toString())
+                            .build());
 
-        if (!shiftResponse.getShift().isPresent()) {
-            throw new RuntimeException("Failed to create shift.");
-        }
-        if (!shiftResponse.getShift().get().getId().isPresent()) {
-            throw new RuntimeException("Shift ID is null.");
+            if (!shiftResponse.getShift().isPresent()) {
+                throw new RuntimeException("Failed to create shift: shift response is empty");
+            }
+            
+            Optional<Shift> shift = shiftResponse.getShift();
+            if (!shift.get().getId().isPresent()) {
+                throw new RuntimeException("Shift ID is null in response: " + shift);
+            }
+            
+            String testShiftId = shift.get().getId().get();
+            System.out.println("Created shift with ID: " + testShiftId);
+            
+            // Add a small delay to ensure the shift is fully created
+            Thread.sleep(1000);
+            
+            DeleteShiftResponse response = client.labor()
+                    .shifts()
+                    .delete(DeleteShiftsRequest.builder().id(testShiftId).build());
+            Assertions.assertNotNull(response);
+            
+        } catch (Exception e) {
+            System.err.println("Test failed with exception: " + e.getMessage());
+            e.printStackTrace();
+            throw e;
         }
-        shiftId = shiftResponse.getShift().get().getId().get();
-        DeleteShiftResponse response = client.labor()
-                .shifts()
-                .delete(DeleteShiftsRequest.builder().id(shiftId).build());
-        Assertions.assertNotNull(response);
     }
 
     @Test
     public void testDeleteBreakType() {
-        // create break type
+        // Create break type
         CreateBreakTypeResponse breakResponse = client.labor()
                 .breakTypes()
                 .create(CreateBreakTypeRequest.builder()
@@ -268,11 +295,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 +310,4 @@ public void testListWorkweekConfigs() {
                 .list(ListWorkweekConfigsRequest.builder().build());
         Assertions.assertFalse(response.getItems().isEmpty());
     }
-}
+}