Add test code for branch nodes

[kanterov]

Read then delete

• Make sure to use a concise title for the pull-request.
• Use #patch, #minor or #major in the pull-request title to bump the corresponding version. Otherwise, the patch version will be bumped. More details

TL;DR

Please replace this text with a description of what this PR accomplishes.

Type

• [ ] Bug Fix
• [ ] Feature
• [ ] Plugin

Are all requirements met?

• [ ] Code completed
• [ ] Smoke tested
• [ ] Unit tests added
• [ ] Code documentation added
• [ ] Any pending items have an associated Issue

Complete description

How did you fix the bug, make the feature etc. Link to any design docs etc

Tracking Issue

https://github.com/lyft/flyte/issues/

Follow-up issue

NA OR https://github.com/lyft/flyte/issues/