fix(@angular-devkit/schematics): support delegating trees in partition (angular#34539)

`partition()`, which backs the `partitionApplyMerge` rule, threw "Tree type is not supported."
for a `UnitTestTree` and other delegating trees. Branch the tree they wrap instead, as
`filter()` does.

PR Close angular#34539

Author: Thanatos9404

packages/angular_devkit/schematics/src/rules/base_spec.ts

@@ -226,6 +226,32 @@ describe('partitionApplyMerge', () => {
     expect(result.exists('/test1')).toBe(false);
     expect(result.exists('/test2')).toBe(false);
   });
+
+  it('works with a UnitTestTree', async () => {
+    const tree = new HostTree();
+    tree.create('/test1', '');
+    tree.create('/test2', '');
+    const predicate = (path: Path) => path.indexOf('1') != -1;
+
+    const ruleYes: Rule = (tree: Tree) => {
+      expect(tree.exists('/test1')).toBe(true);
+      expect(tree.exists('/test2')).toBe(false);
+    };
+    const ruleNo: Rule = (tree: Tree) => {
+      expect(tree.exists('/test1')).toBe(false);
+      expect(tree.exists('/test2')).toBe(true);
+    };
+
+    const result = await lastValueFrom(
+      callRule(
+        partitionApplyMerge(predicate, ruleYes, ruleNo),
+        observableOf(new UnitTestTree(new UnitTestTree(tree))),
+        context,
+      ),
+    );
+    expect(result.exists('/test1')).toBe(true);
+    expect(result.exists('/test2')).toBe(true);
+  });
 });
 
 describe('applyToSubtree', () => {

packages/angular_devkit/schematics/src/tree/static.ts

@@ -7,6 +7,7 @@
  */
 
 import { SchematicsException } from '../exception/exception';
+import { DelegateTree } from './delegate';
 import { FilterHostTree, HostTree } from './host-tree';
 import { FilePredicate, MergeStrategy, Tree } from './interface';
 
@@ -29,6 +30,10 @@ export function merge(
 }
 
 export function partition(tree: Tree, predicate: FilePredicate<boolean>): [Tree, Tree] {
+  if (tree instanceof DelegateTree) {
+    tree = tree.branch();
+  }
+
   if (tree instanceof HostTree) {
     return [
       new FilterHostTree(tree, predicate),