#include "kvm_util.h"
 #include "vmx.h"
 
-#define VCPU_ID           1
 #define MAXPHYADDR 36
 
 #define MEM_REGION_GVA 0x0000123456789000
        GUEST_DONE();
 }
 
-static void run_guest(struct kvm_vm *vm)
-{
-       int rc;
-
-       rc = _vcpu_run(vm, VCPU_ID);
-       TEST_ASSERT(rc == 0, "vcpu_run failed: %d\n", rc);
-}
-
 /*
  * Accessors to get R/M, REG, and Mod bits described in the SDM vol 2,
  * figure 2-2 "Table Interpretation of ModR/M Byte (C8H)".
               GET_RM(insn_bytes[1]) != 0x5;
 }
 
-static void process_exit_on_emulation_error(struct kvm_vm *vm)
+static void process_exit_on_emulation_error(struct kvm_vcpu *vcpu)
 {
-       struct kvm_run *run = vcpu_state(vm, VCPU_ID);
+       struct kvm_run *run = vcpu->run;
        struct kvm_regs regs;
        uint8_t *insn_bytes;
        uint8_t insn_size;
                         * contained an flds instruction that is 2-bytes in
                         * length (ie: no prefix, no SIB, no displacement).
                         */
-                       vcpu_regs_get(vm, VCPU_ID, ®s);
+                       vcpu_regs_get(vcpu->vm, vcpu->id, ®s);
                        regs.rip += 2;
-                       vcpu_regs_set(vm, VCPU_ID, ®s);
+                       vcpu_regs_set(vcpu->vm, vcpu->id, ®s);
                }
        }
 }
 
-static void do_guest_assert(struct kvm_vm *vm, struct ucall *uc)
+static void do_guest_assert(struct ucall *uc)
 {
        TEST_FAIL("%s at %s:%ld", (const char *)uc->args[0], __FILE__,
                  uc->args[1]);
 }
 
-static void check_for_guest_assert(struct kvm_vm *vm)
+static void check_for_guest_assert(struct kvm_vcpu *vcpu)
 {
-       struct kvm_run *run = vcpu_state(vm, VCPU_ID);
        struct ucall uc;
 
-       if (run->exit_reason == KVM_EXIT_IO &&
-           get_ucall(vm, VCPU_ID, &uc) == UCALL_ABORT) {
-               do_guest_assert(vm, &uc);
+       if (vcpu->run->exit_reason == KVM_EXIT_IO &&
+           get_ucall(vcpu->vm, vcpu->id, &uc) == UCALL_ABORT) {
+               do_guest_assert(&uc);
        }
 }
 
-static void process_ucall_done(struct kvm_vm *vm)
+static void process_ucall_done(struct kvm_vcpu *vcpu)
 {
-       struct kvm_run *run = vcpu_state(vm, VCPU_ID);
+       struct kvm_run *run = vcpu->run;
        struct ucall uc;
 
-       check_for_guest_assert(vm);
+       check_for_guest_assert(vcpu);
 
        TEST_ASSERT(run->exit_reason == KVM_EXIT_IO,
                    "Unexpected exit reason: %u (%s)",
                    run->exit_reason,
                    exit_reason_str(run->exit_reason));
 
-       TEST_ASSERT(get_ucall(vm, VCPU_ID, &uc) == UCALL_DONE,
+       TEST_ASSERT(get_ucall(vcpu->vm, vcpu->id, &uc) == UCALL_DONE,
                    "Unexpected ucall command: %lu, expected UCALL_DONE (%d)",
                    uc.cmd, UCALL_DONE);
 }
 
-static uint64_t process_ucall(struct kvm_vm *vm)
+static uint64_t process_ucall(struct kvm_vcpu *vcpu)
 {
-       struct kvm_run *run = vcpu_state(vm, VCPU_ID);
+       struct kvm_run *run = vcpu->run;
        struct ucall uc;
 
        TEST_ASSERT(run->exit_reason == KVM_EXIT_IO,
                    run->exit_reason,
                    exit_reason_str(run->exit_reason));
 
-       switch (get_ucall(vm, VCPU_ID, &uc)) {
+       switch (get_ucall(vcpu->vm, vcpu->id, &uc)) {
        case UCALL_SYNC:
                break;
        case UCALL_ABORT:
-               do_guest_assert(vm, &uc);
+               do_guest_assert(&uc);
                break;
        case UCALL_DONE:
-               process_ucall_done(vm);
+               process_ucall_done(vcpu);
                break;
        default:
                TEST_ASSERT(false, "Unexpected ucall");
 {
        struct kvm_cpuid_entry2 *entry;
        struct kvm_cpuid2 *cpuid;
+       struct kvm_vcpu *vcpu;
        struct kvm_vm *vm;
        uint64_t gpa, pte;
        uint64_t *hva;
        /* Tell stdout not to buffer its content */
        setbuf(stdout, NULL);
 
-       vm = vm_create_default(VCPU_ID, 0, guest_code);
-
        if (!kvm_check_cap(KVM_CAP_SMALLER_MAXPHYADDR)) {
                printf("module parameter 'allow_smaller_maxphyaddr' is not set.  Skipping test.\n");
                return 0;
        }
 
+       vm = vm_create_with_one_vcpu(&vcpu, guest_code);
+
        cpuid = kvm_get_supported_cpuid();
 
        entry = kvm_get_supported_cpuid_index(0x80000008, 0);
        entry->eax = (entry->eax & 0xffffff00) | MAXPHYADDR;
        set_cpuid(cpuid, entry);
 
-       vcpu_set_cpuid(vm, VCPU_ID, cpuid);
+       vcpu_set_cpuid(vm, vcpu->id, cpuid);
 
        rc = kvm_check_cap(KVM_CAP_EXIT_ON_EMULATION_FAILURE);
        TEST_ASSERT(rc, "KVM_CAP_EXIT_ON_EMULATION_FAILURE is unavailable");
        virt_map(vm, MEM_REGION_GVA, MEM_REGION_GPA, 1);
        hva = addr_gpa2hva(vm, MEM_REGION_GPA);
        memset(hva, 0, PAGE_SIZE);
-       pte = vm_get_page_table_entry(vm, VCPU_ID, MEM_REGION_GVA);
-       vm_set_page_table_entry(vm, VCPU_ID, MEM_REGION_GVA, pte | (1ull << 36));
+       pte = vm_get_page_table_entry(vm, vcpu->id, MEM_REGION_GVA);
+       vm_set_page_table_entry(vm, vcpu->id, MEM_REGION_GVA, pte | (1ull << 36));
 
-       run_guest(vm);
-       process_exit_on_emulation_error(vm);
-       run_guest(vm);
+       vcpu_run(vm, vcpu->id);
+       process_exit_on_emulation_error(vcpu);
+       vcpu_run(vm, vcpu->id);
 
-       TEST_ASSERT(process_ucall(vm) == UCALL_DONE, "Expected UCALL_DONE");
+       TEST_ASSERT(process_ucall(vcpu) == UCALL_DONE, "Expected UCALL_DONE");
 
        kvm_vm_free(vm);