How to remove switch statements with polymorphism

In this post we will use polymorphism to remove switch statements. Inappropriately using switch statements is a sign of code smell.

Remove switch statements using polymorphism

In this post we will use polymorphism to remove switch statements. Inappropriately using switch statements is a sign of code smell. Robert Martin in his book Clean Code says that he limits himself to one switch statement per object type.

We will use example of a Customer class. Initially we will use switch statements to implement the methods of the class. Then we will use polymorphism to refactor it and improve the code structure.

public class Customer {

  public enum Membership {
  	Bronze, Silver, Gold
  }

  public Membership membership;

  public Customer(Membership membership) {
  	this.membership = membership;
  }

  public void setMembership(Membership membership) {
    this.membership = membership;
  }

  public Membership getMembership() {
    return membership;
  }

  public double getMonthlyPrice() {
    switch(membership) {
    	case Bronze:
      		return 10;
      	case Silver:
      		return 20;
      	case Gold:
      		return 30;
      	default:
      		throw new IllegalArgumentException("Invalid membership!!");
    }
  }
}

Each Customer has a membership type. getMonthlyPrice method calculates the monthly membership cost of Customer based on their membership level. So far its all good. However we have been asked to add functionality to calculate points earned by the customer. Customers earn points based on their membership level and the amount they spent.

public double getPointsEarned(double amountSpent) {
  switch(membership) {
    case Bronze:
    	return amountSpent * 1;
    case Silver:
    	return amountSpent * 2;
    case Gold:
    	return amountSpent * 3;
    default:
    	throw new IllegalArgumentException("Invalid membership!!");
  }
}

Now we can clearly see whats wrong with our approach. We have duplicated switch statements. Our example is really simple, but the real world application will have complex logic for these calculations.

The first thought to solve this would be to use inheritance and introduce BronzeCustomer, SilverCustomer and GoldCustomer which will extend abstract Customer. Each of these classes will implement their version of getMonthlyPrice and getPointsEarned. But inheritance will not work in our case as the membership level might change in runtime after the object creation. What I mean by this is a Customer can be upgraded or downgraded to other membership level anytime.

We can solve this problem by using State Pattern. We will begin refactoring by introducing Member interface and delegate the functionality of price and membership points calculation to it. We will then create BronzeMember, SilverMember and GoldMember which will implement Member interface.

public interface Member {
  double getMonthlyPrice();
  double getPointsEarned(double amountSpent);
}

class BronzeMember implements Member {
  public double getMonthlyPrice() {
    return 10;
  }

  public double getPointsEarned(double amountSpent) {
    return amountSpent * 1;
  }
}

class SilverMember implements Member {
  public double getMonthlyPrice() {
    return 20;
  }

  public double getPointsEarned(double amountSpent) {
    return amountSpent * 2;
  }
}

class GoldMember implements Member {
  public double getMonthlyPrice() {
    return 30;
  }

  public double getPointsEarned(double amountSpent) {
    return amountSpent * 3;
  }
}

Next we will update Customer class to use Member interface. We have introduced a private property member. Customer delegates the work of calculation to its member instance. When a membership is set, the member instance is also updated based on the membership type. The constructor is updated to use the setMembership method. This is how our final Customer class looks like.

public class Customer {

  public enum Membership {
  	Bronze, Silver, Gold
  }

  private Membership membership;
  private Member member;

  public Customer(Membership membership) {
  	this.setMembership(membership);
  }

  public void setMembership(Membership membership) {
    this.membership = membership;

    switch(membership) {
      case Bronze:
      	this.member = new BronzeMember();
      	break;
      case Silver:
      	this.member = new SilverMember();
      	break;
      case Gold:
      	this.member = new GoldMember();
      	break;
      default:
      	throw new IllegalArgumentException("Invalid membership!!");
    }
  }

  public Membership getMembership() {
    return membership;
  }

  public double getMonthlyPrice() {
    return this.member.getMonthlyPrice();
  }

  public double getPointsEarned(double amountSpent) {
  	return this.member.getPointsEarned(amountSpent);
  }
}

We have successfully used polymorphism to refactor the switch statements in the Customer class. All the complex logic of calculating membership cost and points is delegated to respective classes. When a new membership level is added, we can simply update setMembership method to support it. Above all, none of the client classes using Customer will notice that the class has been changed and continue to work as before.