diff --git a/store-implementation/federated-store/src/test/java/uk/gov/gchq/gaffer/federatedstore/operation/handler/impl/FederatedAddGraphHandlerTest.java b/store-implementation/federated-store/src/test/java/uk/gov/gchq/gaffer/federatedstore/operation/handler/impl/FederatedAddGraphHandlerTest.java
index 879e24768b3..60d9d898773 100644
--- a/store-implementation/federated-store/src/test/java/uk/gov/gchq/gaffer/federatedstore/operation/handler/impl/FederatedAddGraphHandlerTest.java
+++ b/store-implementation/federated-store/src/test/java/uk/gov/gchq/gaffer/federatedstore/operation/handler/impl/FederatedAddGraphHandlerTest.java
@@ -57,7 +57,6 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
-import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static uk.gov.gchq.gaffer.federatedstore.FederatedGraphStorage.USER_IS_ATTEMPTING_TO_OVERWRITE;
 import static uk.gov.gchq.gaffer.federatedstore.FederatedStoreTestUtil.SCHEMA_EDGE_BASIC_JSON;
 import static uk.gov.gchq.gaffer.federatedstore.FederatedStoreTestUtil.loadSchemaFromJson;
@@ -460,7 +459,8 @@ public void shouldAddGraphWithHandler() throws Exception {
                 new Context(testUser),
                 store);
 
-        assertThatThrownBy(() -> store.execute(new GetAllElements(), new Context(testUser)))
-                .hasStackTraceContaining(TestErrorHandler.TEST_ERROR_HANDLER_FLAG);
+        assertThatExceptionOfType(Exception.class)
+                .isThrownBy(() -> store.execute(new GetAllElements(), new Context(testUser)))
+                .withStackTraceContaining(TestErrorHandler.TEST_ERROR_HANDLER_FLAG);
     }
 }