diff --git a/drivers/gpu/drm/i915/gem/selftests/i915_gem_client_blt.c b/drivers/gpu/drm/i915/gem/selftests/i915_gem_client_blt.c
index 855481252bdacda7487f44853a60f0b1870edbb4..fa79233093ebf9f63bb93ec02222e575fe232327 100644
--- a/drivers/gpu/drm/i915/gem/selftests/i915_gem_client_blt.c
+++ b/drivers/gpu/drm/i915/gem/selftests/i915_gem_client_blt.c
@@ -11,8 +11,8 @@
 
 static int igt_client_fill(void *arg)
 {
-	struct intel_context *ce = arg;
-	struct drm_i915_private *i915 = ce->gem_context->i915;
+	struct drm_i915_private *i915 = arg;
+	struct intel_context *ce = i915->engine[BCS0]->kernel_context;
 	struct drm_i915_gem_object *obj;
 	struct rnd_state prng;
 	IGT_TIMEOUT(end);
@@ -89,11 +89,6 @@ static int igt_client_fill(void *arg)
 err_put:
 	i915_gem_object_put(obj);
 err_flush:
-	mutex_lock(&i915->drm.struct_mutex);
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
-	mutex_unlock(&i915->drm.struct_mutex);
-
 	if (err == -ENOMEM)
 		err = 0;
 
@@ -112,5 +107,5 @@ int i915_gem_client_blt_live_selftests(struct drm_i915_private *i915)
 	if (!HAS_ENGINE(i915, BCS0))
 		return 0;
 
-	return i915_subtests(tests, i915->engine[BCS0]->kernel_context);
+	return i915_live_subtests(tests, i915);
 }
diff --git a/drivers/gpu/drm/i915/gem/selftests/i915_gem_context.c b/drivers/gpu/drm/i915/gem/selftests/i915_gem_context.c
index 53c81b5dfd696dbbfd8a4b9cbd0ec7949290c0eb..a23c6df9b9f465e99b8a61fa100b896b8d22edea 100644
--- a/drivers/gpu/drm/i915/gem/selftests/i915_gem_context.c
+++ b/drivers/gpu/drm/i915/gem/selftests/i915_gem_context.c
@@ -1731,5 +1731,5 @@ int i915_gem_context_live_selftests(struct drm_i915_private *dev_priv)
 	if (i915_terminally_wedged(dev_priv))
 		return 0;
 
-	return i915_subtests(tests, dev_priv);
+	return i915_live_subtests(tests, dev_priv);
 }
diff --git a/drivers/gpu/drm/i915/gem/selftests/i915_gem_object_blt.c b/drivers/gpu/drm/i915/gem/selftests/i915_gem_object_blt.c
index e23d8c9e929838406ad23221f852101467c37fda..11d37238c62c8db27d4f804c7f5510dba04e0f89 100644
--- a/drivers/gpu/drm/i915/gem/selftests/i915_gem_object_blt.c
+++ b/drivers/gpu/drm/i915/gem/selftests/i915_gem_object_blt.c
@@ -11,8 +11,8 @@
 
 static int igt_fill_blt(void *arg)
 {
-	struct intel_context *ce = arg;
-	struct drm_i915_private *i915 = ce->gem_context->i915;
+	struct drm_i915_private *i915 = arg;
+	struct intel_context *ce = i915->engine[BCS0]->kernel_context;
 	struct drm_i915_gem_object *obj;
 	struct rnd_state prng;
 	IGT_TIMEOUT(end);
@@ -83,11 +83,6 @@ static int igt_fill_blt(void *arg)
 err_put:
 	i915_gem_object_put(obj);
 err_flush:
-	mutex_lock(&i915->drm.struct_mutex);
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
-	mutex_unlock(&i915->drm.struct_mutex);
-
 	if (err == -ENOMEM)
 		err = 0;
 
@@ -106,5 +101,5 @@ int i915_gem_object_blt_live_selftests(struct drm_i915_private *i915)
 	if (!HAS_ENGINE(i915, BCS0))
 		return 0;
 
-	return i915_subtests(tests, i915->engine[BCS0]->kernel_context);
+	return i915_live_subtests(tests, i915);
 }
diff --git a/drivers/gpu/drm/i915/gt/selftest_hangcheck.c b/drivers/gpu/drm/i915/gt/selftest_hangcheck.c
index cf592a049a71d9f0301402a45a4e25f152695c4d..2d9cc3cd1f279f77a0401dc98f3aa3131fd5cff3 100644
--- a/drivers/gpu/drm/i915/gt/selftest_hangcheck.c
+++ b/drivers/gpu/drm/i915/gt/selftest_hangcheck.c
@@ -1744,7 +1744,7 @@ int intel_hangcheck_live_selftests(struct drm_i915_private *i915)
 	saved_hangcheck = fetch_and_zero(&i915_modparams.enable_hangcheck);
 	drain_delayed_work(&i915->gpu_error.hangcheck_work); /* flush param */
 
-	err = i915_subtests(tests, i915);
+	err = i915_live_subtests(tests, i915);
 
 	mutex_lock(&i915->drm.struct_mutex);
 	igt_flush_test(i915, I915_WAIT_LOCKED);
diff --git a/drivers/gpu/drm/i915/gt/selftest_lrc.c b/drivers/gpu/drm/i915/gt/selftest_lrc.c
index 0c97f953e908fb9a7edf03e115335a5c7369956f..11f490502ca6fd8e07660ff9adc6a19db2456a48 100644
--- a/drivers/gpu/drm/i915/gt/selftest_lrc.c
+++ b/drivers/gpu/drm/i915/gt/selftest_lrc.c
@@ -73,7 +73,6 @@ static int live_sanitycheck(void *arg)
 err_spin:
 	igt_spinner_fini(&spin);
 err_unlock:
-	igt_flush_test(i915, I915_WAIT_LOCKED);
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 	return err;
@@ -290,8 +289,6 @@ static int live_timeslice_preempt(void *arg)
 err_obj:
 	i915_gem_object_put(obj);
 err_unlock:
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 
@@ -472,8 +469,6 @@ static int live_busywait_preempt(void *arg)
 err_ctx_hi:
 	kernel_context_close(ctx_hi);
 err_unlock:
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 	return err;
@@ -580,7 +575,6 @@ static int live_preempt(void *arg)
 err_spin_hi:
 	igt_spinner_fini(&spin_hi);
 err_unlock:
-	igt_flush_test(i915, I915_WAIT_LOCKED);
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 	return err;
@@ -687,7 +681,6 @@ static int live_late_preempt(void *arg)
 err_spin_hi:
 	igt_spinner_fini(&spin_hi);
 err_unlock:
-	igt_flush_test(i915, I915_WAIT_LOCKED);
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 	return err;
@@ -826,8 +819,6 @@ static int live_suppress_self_preempt(void *arg)
 err_client_a:
 	preempt_client_fini(&a);
 err_unlock:
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 	return err;
@@ -995,8 +986,6 @@ static int live_suppress_wait_preempt(void *arg)
 err_client_0:
 	preempt_client_fini(&client[0]);
 err_unlock:
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 	return err;
@@ -1143,8 +1132,6 @@ static int live_chain_preempt(void *arg)
 err_client_hi:
 	preempt_client_fini(&hi);
 err_unlock:
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 	return err;
@@ -1270,7 +1257,6 @@ static int live_preempt_hang(void *arg)
 err_spin_hi:
 	igt_spinner_fini(&spin_hi);
 err_unlock:
-	igt_flush_test(i915, I915_WAIT_LOCKED);
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 	return err;
@@ -2054,5 +2040,5 @@ int intel_execlists_live_selftests(struct drm_i915_private *i915)
 	if (i915_terminally_wedged(i915))
 		return 0;
 
-	return i915_subtests(tests, i915);
+	return i915_live_subtests(tests, i915);
 }
diff --git a/drivers/gpu/drm/i915/gt/selftest_timeline.c b/drivers/gpu/drm/i915/gt/selftest_timeline.c
index 193cc564ade2acd49c938be59578f23de19a7209..eae3b1963bf7a9325d0dec84f1e6c88af66006a6 100644
--- a/drivers/gpu/drm/i915/gt/selftest_timeline.c
+++ b/drivers/gpu/drm/i915/gt/selftest_timeline.c
@@ -821,8 +821,6 @@ static int live_hwsp_recycle(void *arg)
 	}
 
 out:
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
 	intel_runtime_pm_put(&i915->runtime_pm, wakeref);
 	mutex_unlock(&i915->drm.struct_mutex);
 
@@ -841,5 +839,5 @@ int intel_timeline_live_selftests(struct drm_i915_private *i915)
 	if (i915_terminally_wedged(i915))
 		return 0;
 
-	return i915_subtests(tests, i915);
+	return i915_live_subtests(tests, i915);
 }
diff --git a/drivers/gpu/drm/i915/i915_selftest.h b/drivers/gpu/drm/i915/i915_selftest.h
index 207e21b478f21afe1165d863f082acd092a7cbf5..d9b17b9e69930c30c905bd18346cd1672e2c827e 100644
--- a/drivers/gpu/drm/i915/i915_selftest.h
+++ b/drivers/gpu/drm/i915/i915_selftest.h
@@ -66,12 +66,28 @@ struct i915_subtest {
 	const char *name;
 };
 
+int __i915_nop_setup(void *data);
+int __i915_nop_teardown(int err, void *data);
+
+int __i915_live_setup(void *data);
+int __i915_live_teardown(int err, void *data);
+
 int __i915_subtests(const char *caller,
+		    int (*setup)(void *data),
+		    int (*teardown)(int err, void *data),
 		    const struct i915_subtest *st,
 		    unsigned int count,
 		    void *data);
 #define i915_subtests(T, data) \
-	__i915_subtests(__func__, T, ARRAY_SIZE(T), data)
+	__i915_subtests(__func__, \
+			__i915_nop_setup, __i915_nop_teardown, \
+			T, ARRAY_SIZE(T), data)
+#define i915_live_subtests(T, data) ({ \
+	typecheck(struct drm_i915_private *, data); \
+	__i915_subtests(__func__, \
+			__i915_live_setup, __i915_live_teardown, \
+			T, ARRAY_SIZE(T), data); \
+})
 
 #define SUBTEST(x) { x, #x }
 
diff --git a/drivers/gpu/drm/i915/selftests/i915_gem.c b/drivers/gpu/drm/i915/selftests/i915_gem.c
index ed0c17bf6613c1c7ea387b9676525efbd9b746c4..b8ffae481730bf25ea78d29db39cea6a624d11d2 100644
--- a/drivers/gpu/drm/i915/selftests/i915_gem.c
+++ b/drivers/gpu/drm/i915/selftests/i915_gem.c
@@ -154,8 +154,6 @@ static int igt_gem_suspend(void *arg)
 
 	mutex_lock(&i915->drm.struct_mutex);
 	err = switch_to_context(i915, ctx);
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
 	mutex_unlock(&i915->drm.struct_mutex);
 out:
 	mock_file_free(i915, file);
@@ -195,8 +193,6 @@ static int igt_gem_hibernate(void *arg)
 
 	mutex_lock(&i915->drm.struct_mutex);
 	err = switch_to_context(i915, ctx);
-	if (igt_flush_test(i915, I915_WAIT_LOCKED))
-		err = -EIO;
 	mutex_unlock(&i915->drm.struct_mutex);
 out:
 	mock_file_free(i915, file);
@@ -213,5 +209,5 @@ int i915_gem_live_selftests(struct drm_i915_private *i915)
 	if (i915_terminally_wedged(i915))
 		return 0;
 
-	return i915_subtests(tests, i915);
+	return i915_live_subtests(tests, i915);
 }
diff --git a/drivers/gpu/drm/i915/selftests/i915_selftest.c b/drivers/gpu/drm/i915/selftests/i915_selftest.c
index b18eaefef7982f465fa7648bf7148775e5aeeb53..f46ccf817ad56adddf043c41b13566d2d4804c92 100644
--- a/drivers/gpu/drm/i915/selftests/i915_selftest.c
+++ b/drivers/gpu/drm/i915/selftests/i915_selftest.c
@@ -26,6 +26,8 @@
 #include "../i915_drv.h"
 #include "../i915_selftest.h"
 
+#include "igt_flush_test.h"
+
 struct i915_selftest i915_selftest __read_mostly = {
 	.timeout_ms = 1000,
 };
@@ -240,7 +242,40 @@ static bool apply_subtest_filter(const char *caller, const char *name)
 	return result;
 }
 
+int __i915_nop_setup(void *data)
+{
+	return 0;
+}
+
+int __i915_nop_teardown(int err, void *data)
+{
+	return err;
+}
+
+int __i915_live_setup(void *data)
+{
+	struct drm_i915_private *i915 = data;
+
+	return i915_terminally_wedged(i915);
+}
+
+int __i915_live_teardown(int err, void *data)
+{
+	struct drm_i915_private *i915 = data;
+
+	mutex_lock(&i915->drm.struct_mutex);
+	if (igt_flush_test(i915, I915_WAIT_LOCKED))
+		err = -EIO;
+	mutex_unlock(&i915->drm.struct_mutex);
+
+	i915_gem_drain_freed_objects(i915);
+
+	return err;
+}
+
 int __i915_subtests(const char *caller,
+		    int (*setup)(void *data),
+		    int (*teardown)(int err, void *data),
 		    const struct i915_subtest *st,
 		    unsigned int count,
 		    void *data)
@@ -255,10 +290,17 @@ int __i915_subtests(const char *caller,
 		if (!apply_subtest_filter(caller, st->name))
 			continue;
 
+		err = setup(data);
+		if (err) {
+			pr_err(DRIVER_NAME "/%s: setup failed for %s\n",
+			       caller, st->name);
+			return err;
+		}
+
 		pr_info(DRIVER_NAME ": Running %s/%s\n", caller, st->name);
 		GEM_TRACE("Running %s/%s\n", caller, st->name);
 
-		err = st->func(data);
+		err = teardown(st->func(data), data);
 		if (err && err != -EINTR) {
 			pr_err(DRIVER_NAME "/%s: %s failed with error %d\n",
 			       caller, st->name, err);