diff --git a/tools/testing/selftests/kvm/x86/vmx_invalid_nested_guest_state.c b/tools/testing/selftests/kvm/x86/vmx_invalid_nested_guest_state.c
index 87f242704030..0dc616ef41e5 100644
--- a/tools/testing/selftests/kvm/x86/vmx_invalid_nested_guest_state.c
+++ b/tools/testing/selftests/kvm/x86/vmx_invalid_nested_guest_state.c
@@ -11,8 +11,6 @@
 
 #define ARBITRARY_IO_PORT 0x80
 
-static struct kvm_vm *vm;
-
 static void l2_guest_code(void)
 {
 	/*
@@ -21,6 +19,7 @@ static void l2_guest_code(void)
 	 */
 	asm volatile("inb $" __stringify(ARBITRARY_IO_PORT) ", %%al"
 		     ::: "rax");
+	GUEST_FAIL("L2 resumed after stuffing invalid guest state");
 }
 
 static void l1_guest_code(struct vmx_pages *vmx_pages)
@@ -46,35 +45,50 @@ static void l1_guest_code(struct vmx_pages *vmx_pages)
 	GUEST_DONE();
 }
 
-int main(int argc, char *argv[])
+static void vcpu_run_to_io(struct kvm_vcpu *vcpu, bool want_l2)
 {
-	gva_t vmx_pages_gva;
-	struct kvm_sregs sregs;
-	struct kvm_vcpu *vcpu;
-	struct kvm_run *run;
-	struct ucall uc;
-
-	TEST_REQUIRE(kvm_cpu_has(X86_FEATURE_VMX));
-
-	vm = vm_create_with_one_vcpu(&vcpu, l1_guest_code);
-
-	/* Allocate VMX pages and shared descriptors (vmx_pages). */
-	vcpu_alloc_vmx(vm, &vmx_pages_gva);
-	vcpu_args_set(vcpu, 1, vmx_pages_gva);
+	struct kvm_run *run = vcpu->run;
 
 	vcpu_run(vcpu);
 
-	run = vcpu->run;
+	TEST_ASSERT_KVM_EXIT_REASON(vcpu, KVM_EXIT_IO);
+
+	TEST_ASSERT(run->io.port == ARBITRARY_IO_PORT &&
+		    (!!(run->flags & KVM_RUN_X86_GUEST_MODE) == want_l2  ||
+		     !kvm_has_cap(KVM_CAP_X86_GUEST_MODE)),
+		    "Expected IN from port 0x%x from L%u, got port 0x%x from L%u",
+		    ARBITRARY_IO_PORT, 1 + want_l2, run->io.port,
+		    1 + !!(run->flags & KVM_RUN_X86_GUEST_MODE));
+}
+
+static struct kvm_vm *vm_create_and_run_l2(struct kvm_vcpu **vcpu)
+{
+	gva_t vmx_pages_gva;
+	struct kvm_vm *vm;
+
+	vm = vm_create_with_one_vcpu(vcpu, l1_guest_code);
+
+	/* Allocate VMX pages and shared descriptors (vmx_pages). */
+	vcpu_alloc_vmx(vm, &vmx_pages_gva);
+	vcpu_args_set(*vcpu, 1, vmx_pages_gva);
 
 	/*
 	 * The first exit to L0 userspace should be an I/O access from L2.
 	 * Running L1 should launch L2 without triggering an exit to userspace.
 	 */
-	TEST_ASSERT_KVM_EXIT_REASON(vcpu, KVM_EXIT_IO);
+	vcpu_run_to_io(*vcpu, true);
 
-	TEST_ASSERT(run->io.port == ARBITRARY_IO_PORT,
-		    "Expected IN from port %d from L2, got port %d",
-		    ARBITRARY_IO_PORT, run->io.port);
+	return vm;
+}
+
+static void test_invalid_l2_guest_state(void)
+{
+	struct kvm_sregs sregs;
+	struct kvm_vcpu *vcpu;
+	struct kvm_vm *vm;
+	struct ucall uc;
+
+	vm = vm_create_and_run_l2(&vcpu);
 
 	/*
 	 * Stuff invalid guest state for L2 by making TR unusuable.  The next
@@ -96,4 +110,13 @@ int main(int argc, char *argv[])
 	default:
 		TEST_FAIL("Unexpected ucall: %lu", uc.cmd);
 	}
+
+	kvm_vm_free(vm);
+}
+
+int main(int argc, char *argv[])
+{
+	TEST_REQUIRE(kvm_cpu_has(X86_FEATURE_VMX));
+
+	test_invalid_l2_guest_state();
 }