// Copyright 2023 Arm Ltd
 
 #include <kunit/device.h>
+#include <kunit/resource.h>
 #include <kunit/test.h>
 #include "internal.h"
 
 #define BLOCK_TEST_SIZE 12
 
+KUNIT_DEFINE_ACTION_WRAPPER(regmap_exit_action, regmap_exit, struct regmap *);
+
 struct regmap_test_priv {
        struct device *dev;
 };
        if (IS_ERR(ret)) {
                kfree(buf);
                kfree(*data);
+       } else {
+               kunit_add_action(test, regmap_exit_action, ret);
        }
 
        return ret;
 
        /* If using a cache the cache satisfied the read */
        KUNIT_EXPECT_EQ(test, config.cache_type == REGCACHE_NONE, data->read[0]);
-
-       regmap_exit(map);
 }
 
 static void bulk_write(struct kunit *test)
        /* If using a cache the cache satisfied the read */
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_EXPECT_EQ(test, config.cache_type == REGCACHE_NONE, data->read[i]);
-
-       regmap_exit(map);
 }
 
 static void bulk_read(struct kunit *test)
        /* If using a cache the cache satisfied the read */
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_EXPECT_EQ(test, config.cache_type == REGCACHE_NONE, data->read[i]);
-
-       regmap_exit(map);
 }
 
 static void write_readonly(struct kunit *test)
        /* Did that match what we see on the device? */
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_EXPECT_EQ(test, i != 5, data->written[i]);
-
-       regmap_exit(map);
 }
 
 static void read_writeonly(struct kunit *test)
 
        /* Did we trigger a hardware access? */
        KUNIT_EXPECT_FALSE(test, data->read[5]);
-
-       regmap_exit(map);
 }
 
 static void reg_defaults(struct kunit *test)
                        break;
                }
        }
-
-       regmap_exit(map);
 }
 
 static void stride(struct kunit *test)
                        KUNIT_EXPECT_TRUE(test, data->written[i]);
                }
        }
-
-       regmap_exit(map);
 }
 
 static struct regmap_range_cfg test_range = {
                KUNIT_EXPECT_FALSE(test, data->read[i]);
                KUNIT_EXPECT_FALSE(test, data->written[i]);
        }
-
-       regmap_exit(map);
 }
 
 /* Try to stress dynamic creation of cache data structures */
                KUNIT_EXPECT_EQ(test, rval, vals[i]);
                KUNIT_EXPECT_EQ(test, config.cache_type == REGCACHE_NONE, data->read[i]);
        }
-
-       regmap_exit(map);
 }
 
 static void cache_bypass(struct kunit *test)
        regcache_cache_bypass(map, false);
        KUNIT_EXPECT_EQ(test, 0, regmap_read(map, param->from_reg, &rval));
        KUNIT_EXPECT_EQ(test, val, rval);
-
-       regmap_exit(map);
 }
 
 static void cache_sync_marked_dirty(struct kunit *test)
        KUNIT_EXPECT_MEMEQ(test, &data->vals[param->from_reg], val, sizeof(val));
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_EXPECT_EQ(test, true, data->written[param->from_reg + i]);
-
-       regmap_exit(map);
 }
 
 static void cache_sync_after_cache_only(struct kunit *test)
        KUNIT_EXPECT_MEMEQ(test, &data->vals[param->from_reg], val, sizeof(val));
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_EXPECT_TRUE(test, data->written[param->from_reg + i]);
-
-       regmap_exit(map);
 }
 
 static void cache_sync_defaults_marked_dirty(struct kunit *test)
        KUNIT_EXPECT_EQ(test, 0, regcache_sync(map));
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_EXPECT_FALSE(test, data->written[param->from_reg + i]);
-
-       regmap_exit(map);
 }
 
 static void cache_sync_default_after_cache_only(struct kunit *test)
        KUNIT_EXPECT_EQ(test, 0, regcache_sync(map));
        KUNIT_EXPECT_TRUE(test, data->written[param->from_reg + 2]);
        KUNIT_EXPECT_EQ(test, data->vals[param->from_reg + 2], orig_val);
-
-       regmap_exit(map);
 }
 
 static void cache_sync_readonly(struct kunit *test)
        /* Did that match what we see on the device? */
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_EXPECT_EQ(test, i != 5, data->written[param->from_reg + i]);
-
-       regmap_exit(map);
 }
 
 static void cache_sync_patch(struct kunit *test)
                        break;
                }
        }
-
-       regmap_exit(map);
 }
 
 static void cache_drop(struct kunit *test)
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_EXPECT_EQ(test, data->read[param->from_reg + i], i >= 3 && i <= 5);
        KUNIT_EXPECT_MEMEQ(test, &data->vals[param->from_reg], rval, sizeof(rval));
-
-       regmap_exit(map);
 }
 
 static void cache_drop_all_and_sync_marked_dirty(struct kunit *test)
        KUNIT_EXPECT_EQ(test, 0, regcache_sync(map));
        for (i = 0; i <= config.max_register; i++)
                KUNIT_EXPECT_FALSE(test, data->written[i]);
-
-       regmap_exit(map);
 }
 
 static void cache_drop_all_and_sync_no_defaults(struct kunit *test)
        KUNIT_EXPECT_EQ(test, 0, regcache_sync(map));
        for (i = 0; i <= config.max_register; i++)
                KUNIT_EXPECT_FALSE(test, data->written[i]);
-
-       regmap_exit(map);
 }
 
 static void cache_drop_all_and_sync_has_defaults(struct kunit *test)
        KUNIT_EXPECT_EQ(test, 0, regcache_sync(map));
        for (i = 0; i <= config.max_register; i++)
                KUNIT_EXPECT_FALSE(test, data->written[i]);
-
-       regmap_exit(map);
 }
 
 static void cache_present(struct kunit *test)
        /* Now everything should be cached */
        for (i = 0; i < BLOCK_TEST_SIZE; i++)
                KUNIT_ASSERT_TRUE(test, regcache_reg_cached(map, param->from_reg + i));
-
-       regmap_exit(map);
 }
 
 /* Check that caching the window register works with sync */
        if (IS_ERR(ret)) {
                kfree(buf);
                kfree(*data);
+       } else {
+               kunit_add_action(test, regmap_exit_action, ret);
        }
 
        return ret;
                KUNIT_EXPECT_EQ(test, 0, regmap_read(map, i, &rval));
                KUNIT_EXPECT_EQ(test, config.reg_defaults[i].def, rval);
        }
-
-       regmap_exit(map);
 }
 
 static void raw_read_defaults(struct kunit *test)
        }
 
        kfree(rval);
-       regmap_exit(map);
 }
 
 static void raw_write_read_single(struct kunit *test)
        KUNIT_EXPECT_EQ(test, 0, regmap_write(map, 0, val));
        KUNIT_EXPECT_EQ(test, 0, regmap_read(map, 0, &rval));
        KUNIT_EXPECT_EQ(test, val, rval);
-
-       regmap_exit(map);
 }
 
 static void raw_write(struct kunit *test)
 
        /* The values should appear in the "hardware" */
        KUNIT_EXPECT_MEMEQ(test, &hw_buf[2], val, sizeof(val));
-
-       regmap_exit(map);
 }
 
 static bool reg_zero(struct device *dev, unsigned int reg)
        /* Make sure we didn't touch the register after the noinc register */
        KUNIT_EXPECT_EQ(test, 0, regmap_read(map, 1, &val));
        KUNIT_ASSERT_EQ(test, val_test, val);
-
-       regmap_exit(map);
 }
 
 static void raw_sync(struct kunit *test)
 
        /* The values should now appear in the "hardware" */
        KUNIT_EXPECT_MEMEQ(test, &hw_buf[2], &val[0], sizeof(val));
-
-       regmap_exit(map);
 }
 
 static void raw_ranges(struct kunit *test)
                KUNIT_EXPECT_FALSE(test, data->read[i]);
                KUNIT_EXPECT_FALSE(test, data->written[i]);
        }
-
-       regmap_exit(map);
 }
 
 static struct kunit_case regmap_test_cases[] = {