I think adding a pg_partition_root() call to as many pg_partition_tree tests as you modified is overkill ... OTOH I'd have one test that invokes pg_partition_tree(pg_partition_root(some-partition)) to verify that starting from any point in the tree you get the whole tree.
I haven't actually tried to write a query that obtains a tree of constraints using this, mind ... -- Álvaro Herrera https://www.2ndQuadrant.com/ PostgreSQL Development, 24x7 Support, Remote DBA, Training & Services