diff --git a/drivers/base/platform-mock.c b/drivers/base/platform-mock.c
index 3df9f1b..e30c1af 100644
--- a/drivers/base/platform-mock.c
+++ b/drivers/base/platform-mock.c
@@ -6,8 +6,10 @@
  * Author: Brendan Higgins <brendanhiggins@google.com>
  */
 
+#include <linux/klist.h>
 #include <linux/platform_device_mock.h>
 #include <linux/of_platform.h>
+#include "base.h"
 
 struct device_node *of_fake_node(struct test *test, const char *name)
 {
@@ -24,7 +26,7 @@ struct device_node *of_fake_node(struct test *test, const char *name)
 
 struct platform_device *
 of_fake_probe_platform(struct test *test,
-		       const struct platform_driver *driver,
+		       struct platform_driver *driver,
 		       const char *node_name)
 {
 	struct platform_device *pdev;
@@ -45,6 +47,12 @@ of_fake_probe_platform(struct test *test,
 	if (ret)
 		return ERR_PTR(ret);
 
+	pdev->dev.driver = &driver->driver;
+	klist_add_tail(&pdev->dev.p->knode_driver,
+		       &pdev->dev.driver->p->klist_devices);
+	if (ret)
+		return ERR_PTR(ret);
+
 	return pdev;
 }
 
diff --git a/include/linux/platform_device_mock.h b/include/linux/platform_device_mock.h
index 6cac125..9885ec7 100644
--- a/include/linux/platform_device_mock.h
+++ b/include/linux/platform_device_mock.h
@@ -45,7 +45,7 @@ struct device_node *of_fake_node(struct test *test, const char *name);
  */
 struct platform_device *
 of_fake_probe_platform(struct test *test,
-		       const struct platform_driver *driver,
+		       struct platform_driver *driver,
 		       const char *node_name);
 
 /**