diff --git a/airports-flight-service/src/main/java/cz/muni/fi/pa165/service/FlightService.java b/airports-flight-service/src/main/java/cz/muni/fi/pa165/service/FlightService.java
index 1ffa2245701c59ed7e99b3d90a12a085879c944b..90be0f6e01d539f6b822e05a9786f8a610484973 100644
--- a/airports-flight-service/src/main/java/cz/muni/fi/pa165/service/FlightService.java
+++ b/airports-flight-service/src/main/java/cz/muni/fi/pa165/service/FlightService.java
@@ -35,10 +35,16 @@ public class FlightService {
     }
 
     public Flight updateFlight(UUID id,Flight flight){
-        Flight existingFlight = flightRepository.findById(id).get();
-        existingFlight.setDestination(flight.getDestination());
-        existingFlight.setOrigin(flight.getOrigin());
-        existingFlight.setPlaneId(flight.getPlaneId());
+        Flight existingFlight = flightRepository.findById(id).orElseThrow(FlightNotFoundException::new);
+        if(flight.getDestination() != null) {
+            existingFlight.setDestination(flight.getDestination());
+        }
+        if(flight.getOrigin() != null){
+            existingFlight.setOrigin(flight.getOrigin());
+        }
+       if(flight.getPlaneId() != null){
+           existingFlight.setPlaneId(flight.getPlaneId());
+       }
 
         Flight updateFlight =  flightRepository.save(existingFlight);
 
diff --git a/airports-flight-service/src/test/java/cz/muni/fi/pa165/service/FlightServiceTest.java b/airports-flight-service/src/test/java/cz/muni/fi/pa165/service/FlightServiceTest.java
index 8b441700357ac08f0b6609a80ac2d1e57d7cf47f..ca76b51f61c0fa08fbcf9ac6052da49a1fc1eb05 100644
--- a/airports-flight-service/src/test/java/cz/muni/fi/pa165/service/FlightServiceTest.java
+++ b/airports-flight-service/src/test/java/cz/muni/fi/pa165/service/FlightServiceTest.java
@@ -18,6 +18,7 @@ import java.util.Optional;
 import java.util.UUID;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 @ExtendWith(MockitoExtension.class)
 public class FlightServiceTest {
@@ -54,7 +55,7 @@ public class FlightServiceTest {
         Mockito.when(flightRepository.findById(id)).thenReturn(Optional.empty());
 
         // Assert that calling flightService.getFlight(id) should throw FlightNotFoundException
-        Assertions.assertThrows(FlightNotFoundException.class, () -> flightService.getFlight(id));
+        assertThrows(FlightNotFoundException.class, () -> flightService.getFlight(id));
     }
 
 
@@ -75,30 +76,42 @@ public class FlightServiceTest {
     }
 
     @Test
-    void updateFlight_flightFound() {
+    void updateFlight_flightNotFound_throwsFlightNotFoundException() {
         UUID id = UUID.randomUUID(); // Generate a random flight ID
+        Flight update = new Flight(); // Prepare an update object
+
+        // Assert that calling flightService.updateFlight(id, update) should throw FlightNotFoundException
+        assertThrows(FlightNotFoundException.class, () -> flightService.updateFlight(id, update));
+
+        // Verify that stewardRepository.save() was not called
+        Mockito.verify(flightRepository, Mockito.times(0)).save(Mockito.any());
+    }
+
+
+    @Test
+    void updateFlight_flightFound() {
+        UUID id = UUID.randomUUID();
         var destId = UUID.randomUUID();
-        // Prepare an existing flight and an update object
+
         Flight existingFlight = TestDaoFactory.getFlightEntity();
+        existingFlight.setId(id);
+
         Flight update = new Flight();
         update.setDestination(destId.toString());
-        // Prepare the expected updated flight
+
         Flight updatedFlight = TestDaoFactory.getFlightEntity();
+
         updatedFlight.setId(id);
-        updatedFlight.setDestination("New Destination");
+        updatedFlight.setDestination(destId.toString());
 
-        // Mocking the behavior of the flightRepository.findById() method to return the existing flight
         Mockito.when(flightRepository.findById(id)).thenReturn(Optional.of(existingFlight));
 
-        // Mocking the behavior of the flightRepository.save() method to return the updated flight
-        Mockito.when(flightRepository.save(Mockito.any(Flight.class))).thenAnswer(invocation -> invocation.getArgument(0));
+        Mockito.when(flightRepository.save(Mockito.any(Flight.class))).thenReturn(updatedFlight);
 
-        // Call the method to be tested
         Flight returnedFlight = flightService.updateFlight(id, update);
 
-        // Assertions
-        assertEquals(updatedFlight, returnedFlight); // Check if the returned flight matches the expected updated flight
-        Mockito.verify(flightRepository, Mockito.times(1)).save(updatedFlight); // Verify that the save method was called once with the updated flight
+        assertEquals(updatedFlight, returnedFlight);
+        Mockito.verify(flightRepository, Mockito.times(1)).save(updatedFlight);
     }
 
 }