3 files changed, 59 insertions, 11 deletions

diff --git a/test/dm/core.c b/test/dm/core.c
index edd55b05d6..f74c430843 100644
--- a/test/dm/core.c
+++ b/test/dm/core.c
@@ -749,8 +749,7 @@ static int dm_test_uclass_devices_find(struct unit_test_state *uts)
 		ut_assert(dev);
 	}
 
-	ret = uclass_find_first_device(UCLASS_TEST_DUMMY, &dev);
-	ut_assert(ret == -ENODEV);
+	ut_assertok(uclass_find_first_device(UCLASS_TEST_DUMMY, &dev));
 	ut_assert(!dev);
 
 	return 0;
diff --git a/test/dm/pci.c b/test/dm/pci.c
index c325f6600e..fb93e4c78a 100644
--- a/test/dm/pci.c
+++ b/test/dm/pci.c
@@ -38,7 +38,7 @@ static int dm_test_pci_busdev(struct unit_test_state *uts)
 	ut_assertok(dm_pci_bus_find_bdf(PCI_BDF(0, 0x1f, 0), &swap));
 	device = 0;
 	ut_assertok(dm_pci_read_config16(swap, PCI_DEVICE_ID, &device));
-	ut_asserteq(SANDBOX_PCI_DEVICE_ID, device);
+	ut_asserteq(SANDBOX_PCI_SWAP_CASE_EMUL_ID, device);
 
 	/* Test bus#1 and its devices */
 	ut_assertok(uclass_get_device_by_seq(UCLASS_PCI, 1, &bus));
@@ -50,7 +50,7 @@ static int dm_test_pci_busdev(struct unit_test_state *uts)
 	ut_assertok(dm_pci_bus_find_bdf(PCI_BDF(1, 0x0c, 0), &swap));
 	device = 0;
 	ut_assertok(dm_pci_read_config16(swap, PCI_DEVICE_ID, &device));
-	ut_asserteq(SANDBOX_PCI_DEVICE_ID, device);
+	ut_asserteq(SANDBOX_PCI_SWAP_CASE_EMUL_ID, device);
 
 	return 0;
 }
@@ -170,7 +170,7 @@ static int dm_test_pci_mixed(struct unit_test_state *uts)
 	ut_assertok(dm_pci_bus_find_bdf(PCI_BDF(2, 0x1f, 0), &swap));
 	device = 0;
 	ut_assertok(dm_pci_read_config16(swap, PCI_DEVICE_ID, &device));
-	ut_asserteq(SANDBOX_PCI_DEVICE_ID, device);
+	ut_asserteq(SANDBOX_PCI_SWAP_CASE_EMUL_ID, device);
 
 	/* First test I/O */
 	io_addr = dm_pci_read_bar32(swap, 0);
@@ -294,3 +294,48 @@ static int dm_test_pci_ea(struct unit_test_state *uts)
 	return 0;
 }
 DM_TEST(dm_test_pci_ea, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT);
+
+/* Test the dev_read_addr_pci() function */
+static int dm_test_pci_addr_flat(struct unit_test_state *uts)
+{
+	struct udevice *swap1f, *swap1;
+	ulong io_addr, mem_addr;
+
+	ut_assertok(dm_pci_bus_find_bdf(PCI_BDF(0, 0x1f, 0), &swap1f));
+	io_addr = dm_pci_read_bar32(swap1f, 0);
+	ut_asserteq(io_addr, dev_read_addr_pci(swap1f));
+
+	/*
+	 * This device has both I/O and MEM spaces but the MEM space appears
+	 * first
+	 */
+	ut_assertok(dm_pci_bus_find_bdf(PCI_BDF(0, 0x1, 0), &swap1));
+	mem_addr = dm_pci_read_bar32(swap1, 1);
+	ut_asserteq(mem_addr, dev_read_addr_pci(swap1));
+
+	return 0;
+}
+DM_TEST(dm_test_pci_addr_flat, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT |
+		DM_TESTF_FLAT_TREE);
+
+/*
+ * Test the dev_read_addr_pci() function with livetree. That function is
+ * not currently fully implemented, in that it fails to return the BAR address.
+ * Once that is implemented this test can be removed and dm_test_pci_addr_flat()
+ * can be used for both flattree and livetree by removing the DM_TESTF_FLAT_TREE
+ * flag above.
+ */
+static int dm_test_pci_addr_live(struct unit_test_state *uts)
+{
+	struct udevice *swap1f, *swap1;
+
+	ut_assertok(dm_pci_bus_find_bdf(PCI_BDF(0, 0x1f, 0), &swap1f));
+	ut_asserteq(FDT_ADDR_T_NONE, dev_read_addr_pci(swap1f));
+
+	ut_assertok(dm_pci_bus_find_bdf(PCI_BDF(0, 0x1, 0), &swap1));
+	ut_asserteq(FDT_ADDR_T_NONE, dev_read_addr_pci(swap1));
+
+	return 0;
+}
+DM_TEST(dm_test_pci_addr_live, DM_TESTF_SCAN_PDATA | DM_TESTF_SCAN_FDT |
+		DM_TESTF_LIVE_TREE);
diff --git a/test/dm/test-main.c b/test/dm/test-main.c
index 5d79ce641d..72648162a9 100644
--- a/test/dm/test-main.c
+++ b/test/dm/test-main.c
@@ -64,7 +64,7 @@ static int dm_test_destroy(struct unit_test_state *uts)
 
 		/*
 		 * If the uclass doesn't exist we don't want to create it. So
-		 * check that here before we call uclass_find_device()/
+		 * check that here before we call uclass_find_device().
 		 */
 		uc = uclass_find(id);
 		if (!uc)
@@ -130,7 +130,7 @@ static int dm_test_main(const char *test_name)
 	const int n_ents = ll_entry_count(struct unit_test, dm_test);
 	struct unit_test_state *uts = &global_dm_test_state;
 	struct unit_test *test;
-	int run_count;
+	int found;
 
 	uts->priv = &_global_priv_dm_test_state;
 	uts->fail_count = 0;
@@ -148,7 +148,7 @@ static int dm_test_main(const char *test_name)
 	if (!test_name)
 		printf("Running %d driver model tests\n", n_ents);
 
-	run_count = 0;
+	found = 0;
 #ifdef CONFIG_OF_LIVE
 	uts->of_root = gd->of_root;
 #endif
@@ -180,16 +180,20 @@ static int dm_test_main(const char *test_name)
 			ut_assertok(dm_do_test(uts, test, false));
 			runs++;
 		}
-		run_count += runs;
+		found++;
 	}
 
-	if (test_name && !run_count)
+	if (test_name && !found)
 		printf("Test '%s' not found\n", test_name);
 	else
 		printf("Failures: %d\n", uts->fail_count);
 
+	/* Put everything back to normal so that sandbox works as expected */
+#ifdef CONFIG_OF_LIVE
+	gd->of_root = uts->of_root;
+#endif
 	gd->dm_root = NULL;
-	ut_assertok(dm_init(false));
+	ut_assertok(dm_init(IS_ENABLED(CONFIG_OF_LIVE)));
 	dm_scan_platdata(false);
 	dm_scan_fdt(gd->fdt_blob, false);