diff --git a/src/Squidex.Infrastructure/States/StateFactory.cs b/src/Squidex.Infrastructure/States/StateFactory.cs
index 12b257bc9..1e7f282c5 100644
--- a/src/Squidex.Infrastructure/States/StateFactory.cs
+++ b/src/Squidex.Infrastructure/States/StateFactory.cs
@@ -75,22 +75,18 @@ namespace Squidex.Infrastructure.States
 
                         await state.ActivateAsync(stateHolder);
 
-                        var stateEntry = statesCache.CreateEntry(key);
-
-                        stateEntry.Value = state;
-                        stateEntry.AbsoluteExpirationRelativeToNow = CacheDuration;
-
-                        stateEntry.PostEvictionCallbacks.Add(new PostEvictionCallbackRegistration
-                        {
-                            EvictionCallback = (k, v, r, s) =>
+                        statesCache.CreateEntry(key)
+                            .SetValue(state)
+                            .SetAbsoluteExpiration(CacheDuration)
+                            .RegisterPostEvictionCallback((k, v, r, s) =>
                             {
                                 dispatcher.DispatchAsync(() =>
                                 {
                                     state.Dispose();
                                     states.Remove(state);
                                 }).Forget();
-                            }
-                        });
+                            })
+                            .Dispose();
 
                         states.Add(state);
 
diff --git a/tests/Squidex.Infrastructure.Tests/States/StatesTests.cs b/tests/Squidex.Infrastructure.Tests/States/StatesTests.cs
new file mode 100644
index 000000000..47656e0dd
--- /dev/null
+++ b/tests/Squidex.Infrastructure.Tests/States/StatesTests.cs
@@ -0,0 +1,162 @@
+// ==========================================================================
+//  StatesTests.cs
+//  Squidex Headless CMS
+// ==========================================================================
+//  Copyright (c) Squidex Group
+//  All rights reserved.
+// ==========================================================================
+
+using System;
+using System.Threading.Tasks;
+using FakeItEasy;
+using Microsoft.Extensions.Caching.Memory;
+using Microsoft.Extensions.Options;
+using Xunit;
+
+namespace Squidex.Infrastructure.States
+{
+    public class StatesTests
+    {
+        private class MyStatefulObject : StatefulObject<int>
+        {
+            public void SetState(int value)
+            {
+                State = value;
+            }
+        }
+
+        private readonly string key = Guid.NewGuid().ToString();
+        private readonly MyStatefulObject state = new MyStatefulObject();
+        private readonly IMemoryCache cache = new MemoryCache(Options.Create(new MemoryCacheOptions()));
+        private readonly IPubSub pubSub = new InMemoryPubSub();
+        private readonly IServiceProvider services = A.Fake<IServiceProvider>();
+        private readonly IStateStore store = A.Fake<IStateStore>();
+        private readonly StateFactory sut;
+
+        public StatesTests()
+        {
+            A.CallTo(() => services.GetService(typeof(MyStatefulObject)))
+                .Returns(state);
+
+            sut = new StateFactory(pubSub, services, store, cache);
+            sut.Connect();
+        }
+
+        [Fact]
+        public async Task Should_provide_object_from_cache()
+        {
+            cache.Set(key, state);
+
+            var actual = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            Assert.Same(state, actual);
+        }
+
+        [Fact]
+        public async Task Should_read_from_store()
+        {
+            A.CallTo(() => store.ReadAsync<int>(key))
+                .Returns((123, Guid.NewGuid().ToString()));
+
+            var actual = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            Assert.Same(state, actual);
+            Assert.Same(state, cache.Get<MyStatefulObject>(key));
+
+            Assert.Equal(123, state.State);
+        }
+
+        [Fact]
+        public async Task Should_provide_state_from_services_and_add_to_cache()
+        {
+            var actual = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            Assert.Same(state, actual);
+            Assert.Same(state, cache.Get<MyStatefulObject>(key));
+        }
+
+        [Fact]
+        public async Task Should_serve_next_request_from_cache()
+        {
+            var actual1 = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            Assert.Same(state, actual1);
+            Assert.Same(state, cache.Get<MyStatefulObject>(key));
+
+            var actual2 = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            Assert.Same(state, actual2);
+            Assert.Same(state, cache.Get<MyStatefulObject>(key));
+
+            A.CallTo(() => services.GetService(typeof(MyStatefulObject)))
+                .MustHaveHappened(Repeated.Exactly.Once);
+        }
+
+        [Fact]
+        public async Task Should_write_to_store_with_previous_etag()
+        {
+            var etag = Guid.NewGuid().ToString();
+
+            A.CallTo(() => store.ReadAsync<int>(key))
+                .Returns((123, etag));
+
+            var actual = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            Assert.Same(state, actual);
+            Assert.Same(state, cache.Get<MyStatefulObject>(key));
+
+            Assert.Equal(123, state.State);
+
+            state.SetState(456);
+
+            await state.WriteStateAsync();
+
+            A.CallTo(() => store.WriteAsync(key, 456, etag, A<string>.That.Matches(x => x != null)))
+                .MustHaveHappened();
+        }
+
+        [Fact]
+        public async Task Should_remove_from_cache_when_message_sent()
+        {
+            var actual = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            await InvalidateCacheAsync();
+
+            Assert.True(actual.IsDisposed);
+        }
+
+        [Fact]
+        public async Task Should_dispose_states_if_exired()
+        {
+            var actual = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            await RemoveFromCacheAsync();
+
+            Assert.True(actual.IsDisposed);
+        }
+
+        [Fact]
+        public async Task Should_dispose_states_if_disposed()
+        {
+            var actual = await sut.GetAsync<MyStatefulObject, int>(key);
+
+            sut.Dispose();
+
+            Assert.True(actual.IsDisposed);
+        }
+
+        private async Task RemoveFromCacheAsync()
+        {
+            cache.Remove(key);
+
+            await Task.Delay(400);
+        }
+
+        private async Task InvalidateCacheAsync()
+        {
+            pubSub.Publish(new InvalidateMessage { Key = key }, true);
+
+            await Task.Delay(400);
+        }
+    }
+}