From fd1237c34f8022179423f2b8db86e95ea92b6a80 Mon Sep 17 00:00:00 2001
From: A Googler <no-reply@google.com>
Date: Mon, 19 Feb 2024 01:06:10 -0800
Subject: [PATCH] Fix expect.that_struct. It is currently broken, since it
 won't take in attrs, and attrs in a mandatory attribute for the new function
 on a struct.

Also add expect.that_value to allow you to work with arbitrary types

PiperOrigin-RevId: 608252913
---
 lib/private/expect.bzl | 30 +++++++++++++++++++++++++++---
 tests/truth_tests.bzl  | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/lib/private/expect.bzl b/lib/private/expect.bzl
index ab90fd9..ac6cb9e 100644
--- a/lib/private/expect.bzl
+++ b/lib/private/expect.bzl
@@ -81,6 +81,7 @@ def _expect_new(env, meta):
         that_str = lambda *a, **k: _expect_that_str(self, *a, **k),
         that_struct = lambda *a, **k: _expect_that_struct(self, *a, **k),
         that_target = lambda *a, **k: _expect_that_target(self, *a, **k),
+        that_value = lambda *a, **k: _expect_that_value(self, *a, **k),
         where = lambda *a, **k: _expect_where(self, *a, **k),
         # keep sorted end
         # Attributes used by Subject classes and internal helpers
@@ -209,17 +210,24 @@ def _expect_that_str(self, value):
     """
     return StrSubject.new(value, self.meta.derive("string"))
 
-def _expect_that_struct(self, value):
+def _expect_that_struct(self, value, *, attrs, expr = "struct"):
     """Creates a subject for asserting a `struct`.
 
     Args:
         self: implicitly added.
         value: ([`struct`]) the value to check against.
+        expr: ([`str`]) The starting "value of" expression to report in errors.
+        attrs: ([`dict`] of [`str`] to [`callable`]) the functions to convert
+            attributes to subjects. The keys are attribute names that must
+            exist on `actual`. The values are functions with the signature
+            `def factory(value, *, meta)`, where `value` is the actual attribute
+            value of the struct, and `meta` is an [`ExpectMeta`] object.
+
 
     Returns:
         [`StructSubject`] object.
     """
-    return StructSubject.new(value, self.meta.derive("string"))
+    return StructSubject.new(value, meta = self.meta.derive(expr), attrs = attrs)
 
 def _expect_that_target(self, target):
     """Creates a subject for asserting a `Target`.
@@ -244,6 +252,21 @@ def _expect_that_target(self, target):
         },
     ))
 
+def _expect_that_value(self, value, *, factory, expr = "value"):
+    """Creates a subject for asserting an arbitrary value of a custom type.
+
+    Args:
+        self: implicitly added.
+        value: ([`struct`]) the value to check against.
+        factory: A subject factory (a function that takes value and meta).
+            Eg. subjects.collection
+        expr: ([`str`]) The starting "value of" expression to report in errors.
+
+    Returns:
+        A subject corresponding to the type returned by the factory.
+    """
+    return factory(value, self.meta.derive(expr))
+
 def _expect_where(self, **details):
     """Add additional information about the assertion.
 
@@ -272,8 +295,8 @@ def _expect_where(self, **details):
 # buildifier: disable=name-conventions
 Expect = struct(
     # keep sorted start
-    new_from_env = _expect_new_from_env,
     new = _expect_new,
+    new_from_env = _expect_new_from_env,
     that_action = _expect_that_action,
     that_bool = _expect_that_bool,
     that_collection = _expect_that_collection,
@@ -284,6 +307,7 @@ Expect = struct(
     that_str = _expect_that_str,
     that_struct = _expect_that_struct,
     that_target = _expect_that_target,
+    that_value = _expect_that_value,
     where = _expect_where,
     # keep sorted end
 )
diff --git a/tests/truth_tests.bzl b/tests/truth_tests.bzl
index a544034..acc4db4 100644
--- a/tests/truth_tests.bzl
+++ b/tests/truth_tests.bzl
@@ -1356,6 +1356,40 @@ def _str_subject_test(env, _target):
 
 _suite.append(str_subject_test)
 
+def struct_subject_test(name):
+    analysis_test(name = name, impl = _struct_subject_test, target = "truth_tests_helper")
+
+def _struct_subject_test(env, _target):
+    fake_env = _fake_env(env)
+    subject = truth.expect(fake_env).that_struct(
+        struct(a = "a"),
+        attrs = {"a": subjects.str},
+    )
+
+    subject.a().equals("a")
+    _assert_no_failures(fake_env, env = env)
+
+    _end(env, fake_env)
+
+_suite.append(struct_subject_test)
+
+def custom_subject_test(name):
+    analysis_test(name = name, impl = _custom_subject_test, target = "truth_tests_helper")
+
+def _custom_subject_test(env, _target):
+    fake_env = _fake_env(env)
+    subject = truth.expect(fake_env).that_value(
+        ["a"],
+        factory = subjects.collection,
+    )
+
+    subject.contains_exactly(["a"])
+    _assert_no_failures(fake_env, env = env)
+
+    _end(env, fake_env)
+
+_suite.append(custom_subject_test)
+
 def target_subject_test(name):
     analysis_test(name = name, impl = _target_subject_test, target = "truth_tests_helper")  #TODO also file target