7 files changed, 31 insertions, 10 deletions

diff --git a/tests/libqos/ahci.c b/tests/libqos/ahci.c
index f3be5500e1..716ab7939e 100644
--- a/tests/libqos/ahci.c
+++ b/tests/libqos/ahci.c
@@ -128,7 +128,7 @@ QPCIDevice *get_ahci_device(uint32_t *fingerprint)
     uint32_t ahci_fingerprint;
     QPCIBus *pcibus;
 
-    pcibus = qpci_init_pc();
+    pcibus = qpci_init_pc(NULL);
 
     /* Find the AHCI PCI device and verify it's the right one. */
     ahci = qpci_device_find(pcibus, QPCI_DEVFN(0x1F, 0x02));
diff --git a/tests/libqos/libqos-pc.c b/tests/libqos/libqos-pc.c
index df340928a6..aa17c980d8 100644
--- a/tests/libqos/libqos-pc.c
+++ b/tests/libqos/libqos-pc.c
@@ -1,10 +1,13 @@
 #include "qemu/osdep.h"
 #include "libqos/libqos-pc.h"
 #include "libqos/malloc-pc.h"
+#include "libqos/pci-pc.h"
 
 static QOSOps qos_ops = {
     .init_allocator = pc_alloc_init_flags,
-    .uninit_allocator = pc_alloc_uninit
+    .uninit_allocator = pc_alloc_uninit,
+    .qpci_init = qpci_init_pc,
+    .qpci_free = qpci_free_pc,
 };
 
 QOSState *qtest_pc_vboot(const char *cmdline_fmt, va_list ap)
diff --git a/tests/libqos/libqos-spapr.c b/tests/libqos/libqos-spapr.c
index f19408be00..333e6fbb45 100644
--- a/tests/libqos/libqos-spapr.c
+++ b/tests/libqos/libqos-spapr.c
@@ -1,10 +1,13 @@
 #include "qemu/osdep.h"
 #include "libqos/libqos-spapr.h"
 #include "libqos/malloc-spapr.h"
+#include "libqos/pci-spapr.h"
 
 static QOSOps qos_ops = {
     .init_allocator = spapr_alloc_init_flags,
-    .uninit_allocator = spapr_alloc_uninit
+    .uninit_allocator = spapr_alloc_uninit,
+    .qpci_init = qpci_init_spapr,
+    .qpci_free = qpci_free_spapr,
 };
 
 QOSState *qtest_spapr_vboot(const char *cmdline_fmt, va_list ap)
diff --git a/tests/libqos/libqos.c b/tests/libqos/libqos.c
index a852dc5f8e..d842bf5126 100644
--- a/tests/libqos/libqos.c
+++ b/tests/libqos/libqos.c
@@ -20,8 +20,13 @@ QOSState *qtest_vboot(QOSOps *ops, const char *cmdline_fmt, va_list ap)
     cmdline = g_strdup_vprintf(cmdline_fmt, ap);
     qs->qts = qtest_start(cmdline);
     qs->ops = ops;
-    if (ops && ops->init_allocator) {
-        qs->alloc = ops->init_allocator(ALLOC_NO_FLAGS);
+    if (ops) {
+        if (ops->init_allocator) {
+            qs->alloc = ops->init_allocator(ALLOC_NO_FLAGS);
+        }
+        if (ops->qpci_init && qs->alloc) {
+            qs->pcibus = ops->qpci_init(qs->alloc);
+        }
     }
 
     g_free(cmdline);
@@ -49,9 +54,15 @@ QOSState *qtest_boot(QOSOps *ops, const char *cmdline_fmt, ...)
  */
 void qtest_shutdown(QOSState *qs)
 {
-    if (qs->alloc && qs->ops && qs->ops->uninit_allocator) {
-        qs->ops->uninit_allocator(qs->alloc);
-        qs->alloc = NULL;
+    if (qs->ops) {
+        if (qs->pcibus && qs->ops->qpci_free) {
+            qs->ops->qpci_free(qs->pcibus);
+            qs->pcibus = NULL;
+        }
+        if (qs->alloc && qs->ops->uninit_allocator) {
+            qs->ops->uninit_allocator(qs->alloc);
+            qs->alloc = NULL;
+        }
     }
     qtest_quit(qs->qts);
     g_free(qs);
diff --git a/tests/libqos/libqos.h b/tests/libqos/libqos.h
index 604980d125..a9f699019c 100644
--- a/tests/libqos/libqos.h
+++ b/tests/libqos/libqos.h
@@ -8,11 +8,14 @@
 typedef struct QOSOps {
     QGuestAllocator *(*init_allocator)(QAllocOpts);
     void (*uninit_allocator)(QGuestAllocator *);
+    QPCIBus *(*qpci_init)(QGuestAllocator *alloc);
+    void (*qpci_free)(QPCIBus *bus);
 } QOSOps;
 
 typedef struct QOSState {
     QTestState *qts;
     QGuestAllocator *alloc;
+    QPCIBus *pcibus;
     QOSOps *ops;
 } QOSState;
 
diff --git a/tests/libqos/pci-pc.c b/tests/libqos/pci-pc.c
index 82066b8531..9600ed6e41 100644
--- a/tests/libqos/pci-pc.c
+++ b/tests/libqos/pci-pc.c
@@ -212,7 +212,7 @@ static void qpci_pc_iounmap(QPCIBus *bus, void *data)
     /* FIXME */
 }
 
-QPCIBus *qpci_init_pc(void)
+QPCIBus *qpci_init_pc(QGuestAllocator *alloc)
 {
     QPCIBusPC *ret;
 
diff --git a/tests/libqos/pci-pc.h b/tests/libqos/pci-pc.h
index 26211790cd..9479b51642 100644
--- a/tests/libqos/pci-pc.h
+++ b/tests/libqos/pci-pc.h
@@ -14,8 +14,9 @@
 #define LIBQOS_PCI_PC_H
 
 #include "libqos/pci.h"
+#include "libqos/malloc.h"
 
-QPCIBus *qpci_init_pc(void);
+QPCIBus *qpci_init_pc(QGuestAllocator *alloc);
 void     qpci_free_pc(QPCIBus *bus);
 
 #endif