[RFC,v1,12/12] kunit: mock: implement nice, strict and naggy mock distinctions

diff --git a/include/kunit/mock.h b/include/kunit/mock.h
index 4a03f0627afa..955c4267d726 100644
--- a/include/kunit/mock.h
+++ b/include/kunit/mock.h
@@ -95,10 +95,17 @@  struct mock_method {
 	struct list_head expectations;
 };
 
+enum mock_type {
+	MOCK_TYPE_NICE,
+	MOCK_TYPE_NAGGY,
+	MOCK_TYPE_STRICT
+};
+
 struct mock {
 	struct kunit_post_condition parent;
 	struct kunit *test;
 	struct list_head methods;
+	enum mock_type type;
 	/* TODO(brendanhiggins@google.com): add locking to do_expect. */
 	const void *(*do_expect)(struct mock *mock,
 				 const char *method_name,
@@ -108,6 +115,8 @@  struct mock {
 				 int len);
 };
 
+#define DEFAULT_MOCK_TYPE MOCK_TYPE_NAGGY
+
 void mock_init_ctrl(struct kunit *test, struct mock *mock);
 
 void mock_validate_expectations(struct mock *mock);
@@ -125,6 +134,60 @@  struct mock_expectation *mock_add_matcher(struct mock *mock,
 
 #define MOCK(name) name##_mock
 
+/**
+ * STRICT_MOCK() - sets the mock to be strict and returns the mock
+ * @mock: the mock
+ *
+ * For an example, see ``The Nice, the Strict, and the Naggy`` under
+ * ``Using KUnit``.
+ */
+#define STRICT_MOCK(mock) \
+({ \
+	mock_get_ctrl(mock)->type = MOCK_TYPE_STRICT; \
+	mock; \
+})
+
+static inline bool is_strict_mock(struct mock *mock)
+{
+	return mock->type == MOCK_TYPE_STRICT;
+}
+
+/**
+ * NICE_MOCK() - sets the mock to be nice and returns the mock
+ * @mock: the mock
+ *
+ * For an example, see ``The Nice, the Strict, and the Naggy`` under
+ * ``Using KUnit``.
+ */
+#define NICE_MOCK(mock) \
+({ \
+	mock_get_ctrl(mock)->type = MOCK_TYPE_NICE; \
+	mock; \
+})
+
+static inline bool is_nice_mock(struct mock *mock)
+{
+	return mock->type == MOCK_TYPE_NICE;
+}
+
+/**
+ * NAGGY_MOCK() - sets the mock to be naggy and returns the mock
+ * @mock: the mock
+ *
+ * For an example, see ``The Nice, the Strict, and the Naggy`` under
+ * ``Using KUnit``.
+ */
+#define NAGGY_MOCK(mock) \
+({ \
+	mock_get_ctrl(mock)->type = MOCK_TYPE_NAGGY; \
+	mock; \
+})
+
+static inline bool is_naggy_mock(struct mock *mock)
+{
+	return mock->type == MOCK_TYPE_NAGGY;
+}
+
 /**
  * KUNIT_EXPECT_CALL() - Declares a *call expectation* on a mock function.
  * @expectation_call: a mocked method or function with parameters replaced with
diff --git a/lib/kunit/mock-test.c b/lib/kunit/mock-test.c
index df0969b43ade..1c2aa2aa9c1b 100644
--- a/lib/kunit/mock-test.c
+++ b/lib/kunit/mock-test.c
@@ -243,6 +243,50 @@  static void mock_test_do_expect_default_return(struct kunit *test)
 	KUNIT_EXPECT_EQ(test, 0, expectation->times_called);
 }
 
+/**
+ * DOC: Testing the failure condition of different mock types.
+ *
+ * The following tests will test the behaviour of expectations under different
+ * conditions. For example, what happens when an expectation:
+ * - is not satisfied at the end of the test
+ * - is fulfilled but the expected function is called again
+ * - a function is called without expectations set on it
+ *
+ * For each of these conditions, there may be variations between the different
+ * types of mocks: nice mocks, naggy mocks (the default) and strict mocks.
+ *
+ * More information about these mocks can be found in the kernel documentation
+ * under Documentation/test/api/class-and-function-mocking
+ */
+
+/* Method called on strict mock with no expectations will fail */
+static void mock_test_strict_no_expectations_will_fail(struct kunit *test)
+{
+	struct mock_test_context *ctx = test->priv;
+	struct kunit *failing_test = ctx->failing_test;
+	struct mock *mock = ctx->mock;
+	int param0 = 5, param1 = -5;
+	static const char * const two_param_types[] = {"int", "int"};
+	const void *two_params[] = {&param0, &param1};
+
+	mock->type = MOCK_TYPE_STRICT;
+
+	mock_set_default_action(mock,
+				"add",
+				mock_add,
+				kunit_int_return(failing_test, -4));
+
+	mock->do_expect(mock,
+			"add",
+			mock_add,
+			two_param_types,
+			two_params,
+			ARRAY_SIZE(two_params));
+	mock_validate_expectations(mock);
+
+	KUNIT_EXPECT_FALSE(test, failing_test->success);
+}
+
 /*
  * Method called on naggy mock with no expectations will not fail, but will show
  * a warning message
@@ -257,6 +301,8 @@  static void mock_test_naggy_no_expectations_no_fail(struct kunit *test)
 	static const char * const two_param_types[] = {"int", "int"};
 	const void *two_params[] = {&param0, &param1};
 
+	mock->type = MOCK_TYPE_NAGGY;
+
 	mock_set_default_action(mock,
 				"add",
 				real_add,
@@ -282,6 +328,77 @@  static void mock_test_naggy_no_expectations_no_fail(struct kunit *test)
 	KUNIT_EXPECT_FALSE(test, failing_test->success);
 }
 
+/* Method called on nice mock with no expectations will do nothing. */
+static void mock_test_nice_no_expectations_do_nothing(struct kunit *test)
+{
+	struct mock_test_context *ctx = test->priv;
+	struct kunit *failing_test = ctx->failing_test;
+	struct mock *mock = ctx->mock;
+	int param0 = 5, param1 = -5;
+	static const char * const two_param_types[] = {"int", "int"};
+	const void *two_params[] = {&param0, &param1};
+
+	mock->type = MOCK_TYPE_NICE;
+
+	mock->do_expect(mock,
+			"add",
+			mock_add,
+			two_param_types,
+			two_params,
+			ARRAY_SIZE(two_params));
+	mock_validate_expectations(mock);
+
+	KUNIT_EXPECT_TRUE(test, failing_test->success);
+}
+
+/* Test that method called on a mock (of any type) with no matching expectations
+ * will fail test and print all the tried expectations.
+ */
+static void
+run_method_called_but_no_matching_expectation_test(struct kunit *test,
+						   enum mock_type mock_type)
+{
+	struct mock_test_context *ctx = test->priv;
+	struct kunit *failing_test = ctx->failing_test;
+	struct mock *mock = ctx->mock;
+	int param0 = 5, param1 = -5;
+	static const char * const two_param_types[] = {"int", "int"};
+	const void *two_params[] = {&param0, &param1};
+	struct mock_param_matcher *two_matchers[] = {
+		kunit_int_eq(failing_test, 100),
+		kunit_int_eq(failing_test, 100)
+	};
+
+	mock_add_matcher(mock, "add", mock_add, two_matchers,
+			 ARRAY_SIZE(two_matchers));
+
+	mock->type = mock_type;
+
+	mock->do_expect(mock, "add", mock_add, two_param_types, two_params,
+			ARRAY_SIZE(two_params));
+
+	/* Even nice mocks should fail if there's an unmet expectation. */
+	KUNIT_EXPECT_FALSE(test, failing_test->success);
+}
+
+static void mock_test_naggy_no_matching_expectations_fail(struct kunit *test)
+{
+	run_method_called_but_no_matching_expectation_test(test,
+							   MOCK_TYPE_NAGGY);
+}
+
+static void mock_test_strict_no_matching_expectations_fail(struct kunit *test)
+{
+	run_method_called_but_no_matching_expectation_test(test,
+							   MOCK_TYPE_STRICT);
+}
+
+static void mock_test_nice_no_matching_expectations_fail(struct kunit *test)
+{
+	run_method_called_but_no_matching_expectation_test(test,
+							   MOCK_TYPE_NICE);
+}
+
 static void mock_test_mock_validate_expectations(struct kunit *test)
 {
 	struct mock_test_context *ctx = test->priv;
@@ -309,6 +426,54 @@  static void mock_test_mock_validate_expectations(struct kunit *test)
 	KUNIT_EXPECT_FALSE(test, failing_test->success);
 }
 
+static void mock_test_validate_clears_expectations(struct kunit *test)
+{
+	struct mock_test_context *ctx = test->priv;
+	struct kunit *failing_test = ctx->failing_test;
+	struct mock *mock = ctx->mock;
+	struct mock_param_matcher *matchers[] = {
+		kunit_int_eq(failing_test, 5),
+		kunit_int_eq(failing_test, -4)
+	};
+	int param0 = 5, param1 = -4;
+	static const char * const two_param_types[] = {"int", "int"};
+	const void *two_params[] = {&param0, &param1};
+
+	struct mock_expectation *expectation;
+
+	mock->type = MOCK_TYPE_STRICT;
+
+	/* Add an arbitrary matcher for 0 calls */
+	expectation = mock_add_matcher(mock, "add", mock_add,
+				       matchers, ARRAY_SIZE(matchers));
+	expectation->times_called = 0;
+	expectation->min_calls_expected = 0;
+	expectation->max_calls_expected = 0;
+
+	/* Should have 0 calls and should clear the previous expectation */
+	mock_validate_expectations(mock);
+
+	/* Add a new matcher for 1 call */
+	expectation = mock_add_matcher(mock, "add", mock_add,
+				       matchers, ARRAY_SIZE(matchers));
+	expectation->times_called = 0;
+	expectation->min_calls_expected = 1;
+	expectation->max_calls_expected = 1;
+
+	/* Satisfy previous matcher */
+	mock->do_expect(mock, "add", mock_add, two_param_types, two_params,
+			ARRAY_SIZE(two_params));
+
+	/*
+	 * Validate previous satisfy; if we didn't clear the previous
+	 * expectation, it would fail the mock_test.
+	 */
+	mock_validate_expectations(mock);
+
+	/* If all goes well, shouldn't fail the test. */
+	KUNIT_EXPECT_TRUE(test, failing_test->success);
+}
+
 static int mock_test_init(struct kunit *test)
 {
 	struct mock_test_context *ctx;
@@ -346,7 +511,13 @@  static struct kunit_case mock_test_cases[] = {
 	KUNIT_CASE(mock_test_failed_expect_call_fails_test),
 	KUNIT_CASE(mock_test_do_expect_default_return),
 	KUNIT_CASE(mock_test_mock_validate_expectations),
+	KUNIT_CASE(mock_test_strict_no_expectations_will_fail),
 	KUNIT_CASE(mock_test_naggy_no_expectations_no_fail),
+	KUNIT_CASE(mock_test_nice_no_expectations_do_nothing),
+	KUNIT_CASE(mock_test_strict_no_matching_expectations_fail),
+	KUNIT_CASE(mock_test_naggy_no_matching_expectations_fail),
+	KUNIT_CASE(mock_test_nice_no_matching_expectations_fail),
+	KUNIT_CASE(mock_test_validate_clears_expectations),
 	{}
 };
 
diff --git a/lib/kunit/mock.c b/lib/kunit/mock.c
index 12fb88899451..f1fa7a5b9dd4 100644
--- a/lib/kunit/mock.c
+++ b/lib/kunit/mock.c
@@ -85,6 +85,7 @@  void mock_init_ctrl(struct kunit *test, struct mock *mock)
 	mock->test = test;
 	INIT_LIST_HEAD(&mock->methods);
 	mock->do_expect = mock_do_expect;
+	mock->type = DEFAULT_MOCK_TYPE;
 	mock->parent.validate = mock_validate_wrapper;
 	list_add_tail(&mock->parent.node, &test->post_conditions);
 }
@@ -283,7 +284,12 @@  static struct mock_expectation *mock_apply_expectations(
 		mock_add_method_expectation_error(test, stream,
 			"Method was called with no expectations declared: ",
 			mock, method, type_names, params, len);
-		kunit_stream_commit(stream);
+		if (is_strict_mock(mock))
+			fail_and_flush(test, stream);
+		else if (is_naggy_mock(mock))
+			kunit_stream_commit(stream);
+		else
+			kunit_stream_clear(stream);
 		return NULL;
 	}
 
@@ -313,7 +319,7 @@  static struct mock_expectation *mock_apply_expectations(
 		}
 	}
 
-	if (expectations_all_saturated) {
+	if (expectations_all_saturated && !is_nice_mock(mock)) {
 		mock_add_method_expectation_error(test, stream,
 			"Method was called with fully saturated expectations: ",
 			mock, method, type_names, params, len);