#include "../i915_selftest.h"
 #include "igt_flush_test.h"
+#include "igt_live_test.h"
 #include "igt_spinner.h"
 #include "i915_random.h"
 
                I915_USER_PRIORITY(I915_CONTEXT_MIN_USER_PRIORITY);
 
        for_each_engine(engine, i915, id) {
+               struct igt_live_test t;
                struct i915_request *rq;
 
                if (!intel_engine_has_preemption(engine))
                        continue;
 
+               if (igt_live_test_begin(&t, i915, __func__, engine->name)) {
+                       err = -EIO;
+                       goto err_ctx_lo;
+               }
+
                rq = igt_spinner_create_request(&spin_lo, ctx_lo, engine,
                                                MI_ARB_CHECK);
                if (IS_ERR(rq)) {
 
                igt_spinner_end(&spin_hi);
                igt_spinner_end(&spin_lo);
-               if (igt_flush_test(i915, I915_WAIT_LOCKED)) {
+
+               if (igt_live_test_end(&t)) {
                        err = -EIO;
                        goto err_ctx_lo;
                }
                goto err_ctx_hi;
 
        for_each_engine(engine, i915, id) {
+               struct igt_live_test t;
                struct i915_request *rq;
 
                if (!intel_engine_has_preemption(engine))
                        continue;
 
+               if (igt_live_test_begin(&t, i915, __func__, engine->name)) {
+                       err = -EIO;
+                       goto err_ctx_lo;
+               }
+
                rq = igt_spinner_create_request(&spin_lo, ctx_lo, engine,
                                                MI_ARB_CHECK);
                if (IS_ERR(rq)) {
 
                igt_spinner_end(&spin_hi);
                igt_spinner_end(&spin_lo);
-               if (igt_flush_test(i915, I915_WAIT_LOCKED)) {
+
+               if (igt_live_test_end(&t)) {
                        err = -EIO;
                        goto err_ctx_lo;
                }
                struct i915_sched_attr attr = {
                        .priority = I915_USER_PRIORITY(I915_PRIORITY_MAX),
                };
+               struct igt_live_test t;
                struct i915_request *rq;
                int ring_size, count, i;
 
                        goto err_wedged;
                }
 
+               if (igt_live_test_begin(&t, i915, __func__, engine->name)) {
+                       err = -EIO;
+                       goto err_wedged;
+               }
+
                for_each_prime_number_from(count, 1, ring_size) {
                        rq = igt_spinner_create_request(&hi.spin,
                                                        hi.ctx, engine,
                                goto err_wedged;
                        }
                }
+
+               if (igt_live_test_end(&t)) {
+                       err = -EIO;
+                       goto err_wedged;
+               }
        }
 
        err = 0;
        };
        const unsigned int phase[] = { 0, BATCH };
        intel_wakeref_t wakeref;
+       struct igt_live_test t;
        int err = -ENOMEM;
        u32 *cs;
        int n;
        i915_gem_object_flush_map(smoke.batch);
        i915_gem_object_unpin_map(smoke.batch);
 
+       if (igt_live_test_begin(&t, smoke.i915, __func__, "all")) {
+               err = -EIO;
+               goto err_batch;
+       }
+
        for (n = 0; n < smoke.ncontext; n++) {
                smoke.contexts[n] = kernel_context(smoke.i915);
                if (!smoke.contexts[n])
        }
 
 err_ctx:
-       if (igt_flush_test(smoke.i915, I915_WAIT_LOCKED))
+       if (igt_live_test_end(&t))
                err = -EIO;
 
        for (n = 0; n < smoke.ncontext; n++) {