#define SZ_512G (SZ_64G * 8)
 #endif
 
-static struct platform_device *alloc_memdev(int id)
-{
-       struct resource res[] = {
-               [0] = {
-                       .flags = IORESOURCE_MEM,
-               },
-               [1] = {
-                       .flags = IORESOURCE_MEM,
-                       .desc = IORES_DESC_PERSISTENT_MEMORY,
-               },
-       };
-       struct platform_device *pdev;
-       int i, rc;
-
-       for (i = 0; i < ARRAY_SIZE(res); i++) {
-               struct cxl_mock_res *r = alloc_mock_res(SZ_256M);
-
-               if (!r)
-                       return NULL;
-               res[i].start = r->range.start;
-               res[i].end = r->range.end;
-       }
-
-       pdev = platform_device_alloc("cxl_mem", id);
-       if (!pdev)
-               return NULL;
-
-       rc = platform_device_add_resources(pdev, res, ARRAY_SIZE(res));
-       if (rc)
-               goto err;
-
-       return pdev;
-
-err:
-       platform_device_put(pdev);
-       return NULL;
-}
-
 static __init int cxl_test_init(void)
 {
        int rc, i;
                struct platform_device *dport = cxl_switch_dport[i];
                struct platform_device *pdev;
 
-               pdev = alloc_memdev(i);
+               pdev = platform_device_alloc("cxl_mem", i);
                if (!pdev)
                        goto err_mem;
                pdev->dev.parent = &dport->dev;
 
 #include <cxlmem.h>
 
 #define LSA_SIZE SZ_128K
+#define DEV_SIZE SZ_2G
 #define EFFECT(x) (1U << x)
 
 static struct cxl_cel_entry mock_cel[] = {
                .opcode = cpu_to_le16(CXL_MBOX_OP_GET_LSA),
                .effect = cpu_to_le16(0),
        },
+       {
+               .opcode = cpu_to_le16(CXL_MBOX_OP_GET_PARTITION_INFO),
+               .effect = cpu_to_le16(0),
+       },
        {
                .opcode = cpu_to_le16(CXL_MBOX_OP_SET_LSA),
                .effect = cpu_to_le16(EFFECT(1) | EFFECT(2)),
 
 static int mock_id(struct cxl_dev_state *cxlds, struct cxl_mbox_cmd *cmd)
 {
-       struct platform_device *pdev = to_platform_device(cxlds->dev);
        struct cxl_mbox_identify id = {
                .fw_revision = { "mock fw v1 " },
                .lsa_size = cpu_to_le32(LSA_SIZE),
-               /* FIXME: Add partition support */
-               .partition_align = cpu_to_le64(0),
+               .partition_align =
+                       cpu_to_le64(SZ_256M / CXL_CAPACITY_MULTIPLIER),
+               .total_capacity =
+                       cpu_to_le64(DEV_SIZE / CXL_CAPACITY_MULTIPLIER),
        };
-       u64 capacity = 0;
-       int i;
 
        if (cmd->size_out < sizeof(id))
                return -EINVAL;
 
-       for (i = 0; i < 2; i++) {
-               struct resource *res;
-
-               res = platform_get_resource(pdev, IORESOURCE_MEM, i);
-               if (!res)
-                       break;
-
-               capacity += resource_size(res) / CXL_CAPACITY_MULTIPLIER;
+       memcpy(cmd->payload_out, &id, sizeof(id));
 
-               if (le64_to_cpu(id.partition_align))
-                       continue;
+       return 0;
+}
 
-               if (res->desc == IORES_DESC_PERSISTENT_MEMORY)
-                       id.persistent_capacity = cpu_to_le64(
-                               resource_size(res) / CXL_CAPACITY_MULTIPLIER);
-               else
-                       id.volatile_capacity = cpu_to_le64(
-                               resource_size(res) / CXL_CAPACITY_MULTIPLIER);
-       }
+static int mock_partition_info(struct cxl_dev_state *cxlds,
+                              struct cxl_mbox_cmd *cmd)
+{
+       struct cxl_mbox_get_partition_info pi = {
+               .active_volatile_cap =
+                       cpu_to_le64(DEV_SIZE / 2 / CXL_CAPACITY_MULTIPLIER),
+               .active_persistent_cap =
+                       cpu_to_le64(DEV_SIZE / 2 / CXL_CAPACITY_MULTIPLIER),
+       };
 
-       id.total_capacity = cpu_to_le64(capacity);
+       if (cmd->size_out < sizeof(pi))
+               return -EINVAL;
 
-       memcpy(cmd->payload_out, &id, sizeof(id));
+       memcpy(cmd->payload_out, &pi, sizeof(pi));
 
        return 0;
 }
        case CXL_MBOX_OP_GET_LSA:
                rc = mock_get_lsa(cxlds, cmd);
                break;
+       case CXL_MBOX_OP_GET_PARTITION_INFO:
+               rc = mock_partition_info(cxlds, cmd);
+               break;
        case CXL_MBOX_OP_SET_LSA:
                rc = mock_set_lsa(cxlds, cmd);
                break;