Visually Explained: Refactoring to the State Pattern in Java (From If-Else Hell to Clean Code)

Have you ever found yourself trapped in a labyrinth of if-else statements, trying to manage an object that behaves differently depending on its current status? A document can be a DRAFT, then UNDER_REVIEW, then PUBLISHED. Each state allows different actions. As more states and transitions are added, the logic becomes a tangled mess.

This is a classic sign that you need the State Design Pattern.

This guide will walk you through refactoring a chaotic, state-based Document class into a clean, scalable, and maintainable structure using the State Pattern in Java.

The Problem: A Monolith of Conditionals

Let’s start with the code we want to fix. We have a Document class that manages its state with a simple String or enum and a series of conditional checks in every method.

public class Document {
    private String state; // e.g., "DRAFT", "UNDER_REVIEW", "PUBLISHED"

    public Document() {
        this.state = "DRAFT";
    }

    public void edit(String content) {
        if (state.equals("DRAFT") || state.equals("NEEDS_REVISION")) {
            System.out.println("Editing the document...");
            // ... logic to update content
        } else {
            System.out.println("ERROR: Cannot edit in the current state: " + state);
        }
    }

    public void submitForReview() {
        if (state.equals("DRAFT")) {
            this.state = "UNDER_REVIEW";
            System.out.println("Document submitted for review.");
        } else {
            System.out.println("ERROR: Can only submit drafts for review.");
        }
    }

    public void publish() {
        if (state.equals("UNDER_REVIEW")) {
            // Imagine more logic here... user permissions, etc.
            this.state = "PUBLISHED";
            System.out.println("Document has been published!");
        } else if (state.equals("PUBLISHED")) {
            System.out.println("INFO: Document is already published.");
        } 
        else {
            System.out.println("ERROR: Cannot publish directly from state: " + state);
        }
    }

    // ... many more methods (reject, expire, etc.) with similar conditional logic
}

The workflow is hard to follow, and adding a new state requires modifying every single method. This violates the Open/Closed Principle and makes the code brittle.

The Solution: The State Design Pattern

The State Pattern allows an object (the Context) to change its behavior when its internal state changes. The object appears to change its class because its behavior is delegated to a State object.

Here are the key components:

1. Context (Document): The main class that holds a reference to a concrete state object. It delegates actions to this state object.
2. State (IDocumentState): An interface that defines the common methods for all possible actions (edit(), publish(), etc.).
3. Concrete States (DraftState, PublishedState, etc.): Classes that implement the State interface, providing the specific behavior for each state.

Let’s visualize the relationship between these components.

classDiagram
    class Document {
        -IDocumentState currentState
        +changeState(IDocumentState newState)
        +edit(content)
        +submitForReview()
        +publish()
    }
    class IDocumentState {
        <<interface>>
        +edit(Document doc, content)
        +submitForReview(Document doc)
        +publish(Document doc)
    }
    class DraftState {
        +edit(Document doc, content)
        +submitForReview(Document doc)
        +publish(Document doc)
    }
    class UnderReviewState {
        +edit(Document doc, content)
        +submitForReview(Document doc)
        +publish(Document doc)
    }
    class PublishedState {
        +edit(Document doc, content)
        +submitForReview(Document doc)
        +publish(Document doc)
    }

    Document o-- IDocumentState : delegates to
    IDocumentState <|.. DraftState : implements
    IDocumentState <|.. UnderReviewState : implements
    IDocumentState <|.. PublishedState : implements

Step-by-Step Refactoring

Let’s transform our messy Document class into this clean architecture.

Step 1: Define the State Interface

First, we create an interface that declares all the actions a document can undergo. Each method takes the context (Document) as a parameter to allow states to transition the context to a new state.

// src/main/java/com/example/state/IDocumentState.java
package com.example.state;

import com.example.Document;

public interface IDocumentState {
    String getName();
    void edit(Document doc, String content);
    void submitForReview(Document doc);
    void publish(Document doc);
    void reject(Document doc);
    void expire(Document doc);
}

Step 2: Create the Concrete State Classes

Next, we create a separate class for each state in our workflow. This organizes the logic cleanly.

src/main/java/com/example/
├── Document.java
└── state/
    ├── IDocumentState.java
    ├── DraftState.java
    ├── UnderReviewState.java
    ├── PublishedState.java
    ├── NeedsRevisionState.java
    ├── RejectedState.java
    └── ExpiredState.java

Step 3: Implement Logic within Each State

Now, we move the logic from the if-else blocks into the corresponding state classes.

DraftState.java

In the DRAFT state, we can only edit and submit. All other actions are invalid.

// src/main/java/com/example/state/DraftState.java
package com.example.state;

import com.example.Document;

public class DraftState implements IDocumentState {
    @Override
    public String getName() { return "Draft"; }

    @Override
    public void edit(Document doc, String content) {
        System.out.println("Editing the document content...");
        doc.setContent(content); // Assuming Document has a setContent method
    }

    @Override
    public void submitForReview(Document doc) {
        System.out.println("Submitting document for review.");
        doc.changeState(new UnderReviewState()); // Transition to the next state!
    }

    @Override
    public void publish(Document doc) {
        System.out.println("CANNOT publish directly from Draft state.");
    }

    @Override
    public void reject(Document doc) {
        System.out.println("CANNOT reject a Draft.");
    }

    @Override
    public void expire(Document doc) {
        System.out.println("CANNOT expire a Draft.");
    }
}

PublishedState.java

Once a document is published, most actions are forbidden. It can only be expired.

// src/main/java/com/example/state/PublishedState.java
package com.example.state;

import com.example.Document;

public class PublishedState implements IDocumentState {
    @Override
    public String getName() { return "Published"; }

    @Override
    public void edit(Document doc, String content) {
        System.out.println("CANNOT edit a Published document.");
    }

    @Override
    public void submitForReview(Document doc) {
        System.out.println("CANNOT submit a Published document for review.");
    }

    @Override
    public void publish(Document doc) {
        System.out.println("INFO: Document is already published.");
    }
    
    @Override
    public void reject(Document doc) {
        System.out.println("CANNOT reject a Published document.");
    }

    @Override
    public void expire(Document doc) {
        System.out.println("Expiring the document. It will no longer be public.");
        doc.changeState(new ExpiredState());
    }
}

[!TIP] If a state class has no fields of its own, you can make it a Singleton. This prevents creating unnecessary objects and saves memory, as the behavior for a given state never changes.

Step 4: Refactor the Context (Document) Class

Finally, we refactor the Document class. We remove all the conditional logic and delegate the calls to the current state object.

Here is a “diff” showing the transformation.

- public class Document {
-     private String state;
- 
-     public Document() {
-         this.state = "DRAFT";
-     }
- 
-     public void edit(String content) {
-         if (state.equals("DRAFT") || state.equals("NEEDS_REVISION")) {
-             System.out.println("Editing the document...");
-         } else {
-             System.out.println("ERROR: Cannot edit in the current state: " + state);
-         }
-     }
- 
-     public void publish() {
-         if (state.equals("UNDER_REVIEW")) {
-             this.state = "PUBLISHED";
-             System.out.println("Document has been published!");
-         } else {
-             System.out.println("ERROR: Cannot publish directly from state: " + state);
-         }
-     }
-     // ... and so on for other methods
- }
+ import com.example.state.DraftState;
+ import com.example.state.IDocumentState;
+ 
+ public class Document {
+     private IDocumentState currentState;
+     private String content;
+ 
+     public Document() {
+         // The document starts in the Draft state
+         this.currentState = new DraftState();
+         System.out.println("New document created. Current state: " + currentState.getName());
+     }
+ 
+     // This method allows state objects to transition the context
+     public void changeState(IDocumentState newState) {
+         this.currentState = newState;
+         System.out.println("STATE CHANGED TO: " + currentState.getName());
+     }
+ 
+     // All actions are now delegated to the state object
+     public void edit(String newContent) {
+         currentState.edit(this, newContent);
+     }
+ 
+     public void submitForReview() {
+         currentState.submitForReview(this);
+     }
+ 
+     public void publish() {
+         currentState.publish(this);
+     }
+ 
+     public void reject() {
+         currentState.reject(this);
+     }
+ 
+     public void expire() {
+         currentState.expire(this);
+     }
+ 
+     public void setContent(String content) {
+         this.content = content;
+     }
+ }

Our Document class is now incredibly simple. It doesn’t know or care about the rules of any particular state; it only knows how to delegate responsibility to its current state object.

Visualizing the New Workflow

With the State Pattern, the document’s lifecycle becomes a clear, manageable flow.

graph TD
    A[Draft] -- submitForReview() --> B(Under Review);
    B -- reject() --> C{Needs Revision};
    C -- edit() & submitForReview() --> B;
    B -- publish() --> D((Published));
    B -- reject() --> E((Rejected));
    D -- expire() --> F((Expired));

    style A fill:#f9f,stroke:#333,stroke-width:2px
    style D fill:#bbf,stroke:#333,stroke-width:2px
    style E fill:#fbb,stroke:#333,stroke-width:2px
    style F fill:#999,stroke:#333,stroke-width:2px

Conclusion: Why the State Pattern Is a Game-Changer

By refactoring to the State Pattern, we’ve achieved several key benefits:

1. Organized Logic: Each state’s behavior is isolated in its own class, making it easy to understand and manage.
2. Open/Closed Principle: We can add new states (e.g., ArchivedState) without modifying any existing state classes or the Document context. We just create a new class that implements IDocumentState.
3. Simplified Context: The Document class is no longer bloated with conditional logic. Its only job is to maintain the current state and delegate tasks.
4. Clearer State Transitions: State transitions are explicit and self-contained within the state classes themselves (doc.changeState(new ...)).

The State Pattern is a powerful tool for managing objects with complex, state-dependent behavior. It transforms tangled conditional logic into a clean, object-oriented structure that is robust, maintainable, and easy to extend.

Test Your Knowledge