diff --git a/Sources/FirebladeECS/FSM.swift b/Sources/FirebladeECS/FSM.swift
index 946f5c2..d11567e 100644
--- a/Sources/FirebladeECS/FSM.swift
+++ b/Sources/FirebladeECS/FSM.swift
@@ -60,11 +60,7 @@ public class ComponentSingletonProvider {
         ObjectIdentifier(instance)
     }
     
-    public init<T: ComponentInitializable>(type: T.Type) {
-        componentType = type
-    }
-    
-    internal init(type: ComponentInitializable.Type) {
+    public init(type: ComponentInitializable.Type) {
         componentType = type
     }
 }
@@ -108,7 +104,7 @@ public class EntityState {
     
     public init() { }
     
-    public func add<C: ComponentInitializable>(_ type: C.Type) -> StateComponentMapping {
+    @discardableResult public func add<C: ComponentInitializable>(_ type: C.Type) -> StateComponentMapping {
         StateComponentMapping(creatingState: self, type: type)
     }
     
@@ -133,32 +129,27 @@ public class StateComponentMapping {
         provider = ComponentTypeProvider(type: type)
     }
     
-    public func withInstance(_ component: Component) -> StateComponentMapping {
+    @discardableResult public func withInstance(_ component: Component) -> StateComponentMapping {
         setProvider(ComponentInstanceProvider(instance: component))
         return self
     }
     
-    public func withType<T: ComponentInitializable>(_ type: T.Type) -> Self {
+    @discardableResult public func withType<T: ComponentInitializable>(_ type: T.Type) -> Self {
         setProvider(ComponentTypeProvider(type: type))
         return self
     }
     
-    public func withSingleton<T: ComponentInitializable>(_ type: T.Type?) -> Self {
-        if let type = type {
-            setProvider(ComponentSingletonProvider(type: type))
-        } else {
-            setProvider(ComponentSingletonProvider(type: componentType))
-        }
-        
+    @discardableResult public func withSingleton<T: ComponentInitializable>(_ type: T.Type?) -> Self {
+        setProvider(ComponentSingletonProvider(type: type ?? componentType))
         return self
     }
     
-    public func withMethod(_ closure: DynamicComponentProvider.Closure) -> Self {
+    @discardableResult public func withMethod(_ closure: DynamicComponentProvider.Closure) -> Self {
         setProvider(DynamicComponentProvider(closure: closure))
         return self
     }
     
-    public func withProvider(_ provider: ComponentProvider) -> Self {
+    @discardableResult public func withProvider(_ provider: ComponentProvider) -> Self {
         setProvider(provider)
         return self
     }
@@ -220,7 +211,6 @@ public class EntityStateMachine {
                 } else {
                     entity.remove(t.key)
                 }
-                
             }
         } else {
             toAdd = newState.providers
diff --git a/Tests/FirebladeECSTests/FSMTests.swift b/Tests/FirebladeECSTests/FSMTests.swift
index bfd3833..4402448 100644
--- a/Tests/FirebladeECSTests/FSMTests.swift
+++ b/Tests/FirebladeECSTests/FSMTests.swift
@@ -154,6 +154,74 @@ class DynamicComponentProviderTests: XCTestCase {
     }
 }
 
+@testable import class FirebladeECS.EntityState
+
+class EntityStateTests: XCTestCase {
+    private var state = EntityState()
+    
+    override func setUp() {
+        state = EntityState()
+    }
+    
+    override func tearDown() {
+        state = EntityState()
+    }
+    
+    func testAddWithNoQualifierCreatesTypeProvider() {
+        state.add(MockComponent.self)
+        let provider = state.providers[MockComponent.identifier]
+        XCTAssertTrue(provider is ComponentTypeProvider?)
+        XCTAssertTrue(provider?.getComponent() is MockComponent?)
+    }
+    
+    func testAddWithTypeQualifierCreatesTypeProvider() {
+        state.add(MockComponent.self).withType(MockComponent2.self)
+        let provider = state.providers[MockComponent.identifier]
+        XCTAssertTrue(provider is ComponentTypeProvider?)
+        XCTAssertTrue(provider?.getComponent() is MockComponent2?)
+    }
+    
+    func testAddWithInstanceQualifierCreatesInstanceProvider() {
+        let component = MockComponent()
+        state.add(MockComponent.self).withInstance(component)
+        let provider = state.providers[MockComponent.identifier]
+        XCTAssertTrue(provider is ComponentInstanceProvider?)
+        XCTAssertTrue(provider?.getComponent() === component)
+    }
+    
+    func testAddWithSingletonQualifierCreatesSingletonProvider() {
+        state.add(MockComponent.self).withSingleton(MockComponent.self)
+        let  provider = state.providers[MockComponent.identifier]
+        XCTAssertTrue(provider is ComponentSingletonProvider?)
+        XCTAssertTrue(provider?.getComponent() is MockComponent?)
+    }
+    
+    func testAddWithMethodQualifierCreatesDynamicProvider() {
+        let dynamickProvider = DynamicComponentProvider.Closure {
+            MockComponent()
+        }
+        
+        state.add(MockComponent.self).withMethod(dynamickProvider)
+        let provider = state.providers[MockComponent.identifier]
+        XCTAssertTrue(provider is DynamicComponentProvider?)
+        XCTAssertTrue(provider?.getComponent() is MockComponent)
+    }
+    
+    class MockComponent: ComponentInitializable {
+        let value: Int
+        
+        init(value: Int) {
+            self.value = value
+        }
+        
+        required init() {
+            self.value = 0
+        }
+    }
+    
+    class MockComponent2: MockComponent {}
+}
+
 class EntityStateMachineTests: XCTestCase {
     // TODO:
 }