Test Driven Development with JUnit 5. Part 4
4. Refactoring the flight-management application
We would like to do some refactoring and replace the conditional statements with polymorphism.
The key to refactoring is to move the design to using polymorphism instead of procedural-style conditional code. With polymorphism (the ability of one object to pass more than one IS-A test), the method you are calling is determined not at compile-time, but at runtime, depending on the effective object type.
The principle in action here is called the open/closed principle. Practically, it means the design shown on the left will require changes to the existing class each time we add a new flight type. These changes may reflect in each conditional decision made based on the flight type. In addition, we are forced to rely on the flightType field and introduce unexecuted default cases.
With the design on the right—which is refactored by replacing conditional with polymorphism—we do not need a flightType evaluation or a default value in the switch instructions. We can even add a new type—let’s anticipate a little and call it PremiumFlight—by simply extending the base class and defining its behavior. According to the open/closed principle, the hierarchy will be open for extensions (we can easily add new classes) but closed for modifications (existing classes, starting with the Flight base class, will not be modified).
How can we be sure we are doing the right thing and not affecting already-working functionality? The answer is that passing the tests provides assurance that existing functionality is untouched. The benefits of the TDD approach really show themselves!
The refactoring will be achieved by keeping the base Flight class and, for each conditional type, adding a separate class to extend Flight. we will change addPassenger and removePassenger to abstract methods and delegate their implementation to subclasses. The flightType field is no longer significant and will be removed.
public abstract class Flight { #1
private String id;
List<Passenger> passengers = new ArrayList<Passenger>(); #2
public Flight(String id) {
this.id = id;
}
public String getId() {
return id;
}
public List<Passenger> getPassengersList() {
return Collections.unmodifiableList(passengers);
}
public abstract boolean addPassenger(Passenger passenger); #3
public abstract boolean removePassenger(Passenger passenger); #3
}
In this listing:
- We declare the class as abstract, making it the basis of the flight hierarchy #1.
- We make the passengers list package-private, allowing it to be directly inherited by the subclasses in the same package #2.
- We declare addPassenger and removePassenger as abstract methods, delegating their implementation to the subclasses #3.
We introduce an EconomyFlight class that extends Flight and implements the inherited addPassenger and removePassenger abstract methods.
public class EconomyFlight extends Flight { #1
public EconomyFlight(String id) { #2
super(id); #2
} #2
@Override
public boolean addPassenger(Passenger passenger) { #3
return passengers.add(passenger); #3
} #3
@Override
public boolean removePassenger(Passenger passenger) { #4
if (!passenger.isVip()) { #4
return passengers.remove(passenger); #4
} #4
return false; #4
} #4
}
In this listing:
- We declare the EconomyFlight class extending the Flight abstract class #1 and create a constructor calling the constructor of the superclass #2.
- We implement the addPassenger method according to the business logic: we simply add a passenger to an economy flight with no restrictions #3.
- We implement the removePassenger method according to the business logic: a passenger can be removed from a flight only if the passenger is not a VIP #4.
We also introduce a BusinessFlight class that extends Flight and implements the inherited addPassenger and removePassenger abstract methods.
public class BusinessFlight extends Flight { #1
public BusinessFlight(String id) { #2
super(id); #2
} #2
@Override
public boolean addPassenger(Passenger passenger) { #3
if (passenger.isVip()) { #3
return passengers.add(passenger); #3
} #3
return false; #3
} #3
@Override
public boolean removePassenger(Passenger passenger) { #4
return false; #4
} #4
}
In this listing:
- We declare the BusinessFlight class extending the Flight abstract class #1 and create a constructor calling the constructor of the superclass #2.
- We implement the addPassenger method according to the business logic: only a VIP passenger can be added to a business flight #3.
- We implement the removePassenger method according to the business logic: a passenger cannot be removed from a business flight #4.
Refactoring by replacing the conditional with polymorphism, we immediately see that the methods now look much shorter and clearer, not cluttered with decision-making. Also, we are not forced to treat the previous default case that was never expected and that threw an exception. Of course, the refactoring and the API changes propagate into the tests, as shown next.
public class AirportTest {
@DisplayName("Given there is an economy flight")
@Nested
class EconomyFlightTest {
private Flight economyFlight;
@BeforeEach
void setUp() {
economyFlight = new EconomyFlight("1"); #1
}
[...]
}
@DisplayName("Given there is a business flight")
@Nested
class BusinessFlightTest {
private Flight businessFlight;
@BeforeEach
void setUp() {
businessFlight = new BusinessFlight("2"); #2
}
[...]
}
}
In this listing, we replace the previous Flight instantiations with instantiations of EconomyFlight #1 and BusinessFlight #2. We also remove the Airport class that served as a client for the Passenger and Flight classes—it is no longer needed, now that we introduced the tests. It previously served to declare the main method that created different types of flights and passengers and made them act together.