SK_LightningTreeApp.app

<aura:application extends="force:slds">
    <c:SK_LightningTreeCmp ltngcurrentRecId="0019000000ld4kO"
          ltngSobjectname="Account"
          ltngParentFieldAPIName="ParentId"
          ltngLabelFieldAPIName="Name"
          ltngHierarchyHeader="Account Hierarchy"/>
    
    
    <c:SK_LightningTreeCmp ltngcurrentRecId="5009000000GJkJG"
          ltngSobjectname="Case"
          ltngParentFieldAPIName="ParentId"
          ltngLabelFieldAPIName="CaseNumber"
          ltngHierarchyHeader="Case Hierarchy"/>
</aura:application>

SK_LightningTreeCmp.cmp

<aura:component controller="SK_LightningTreeCmpController">
    <aura:attribute name="ltngcurrentRecId" type="String" required="true"/>
    <aura:attribute name="ltngSobjectname" type="String" required="true"/>
    <aura:attribute name="ltngParentFieldAPIName" type="String" required="true"/>
    <aura:attribute name="ltngLabelFieldAPIName" type="String" required="true"/>
    <aura:attribute name="ltngHierarchyHeader" type="String" required="true"/>
    <aura:attribute name="items" type="Object"/>
    
    <aura:handler event="aura:waiting" action="{!c.showSpinner}"/>
    <aura:handler event="aura:doneWaiting" action="{!c.hideSpinner}"/>
    <div class="slds-align--absolute-center">
        <lightning:spinner aura:id="spinner" variant="brand" size="large" class="slds=hide"/>
    </div>
    <aura:handler name="init" value="{!this}" action="{!c.doInit}" />
    <lightning:tree items="{!v.items }" header="{!v.ltngHierarchyHeader}" onselect="{!c.handleSelect}"/>
</aura:component>

SK_LightningTreeCmpController.cls

public class SK_LightningTreeCmpController {
	@AuraEnabled
    public static List<HierarchyData> findHierarchyData(string recId, string sObjectName, string parentFieldAPIname, string labelFieldAPIName){
        List<HierarchyData> returnValue = new List<HierarchyData>();
        string queryString = 'select id, '+labelFieldAPIName+ ' ,' +parentFieldAPIname+ ' from '+sObjectName;
        //Section to get all child account details from ultimate parent starts-------------------------
        List<String> currentParent      = new List<String>{};
        Integer level               = 0;
        Boolean endOfStructure      = false;
        //method to find ultimate parent of account
        string topMostparent = GetUltimateParentId(recId, sObjectName,parentFieldAPIname );
        currentParent.add(topMostparent);
        system.debug('**********topMostparent:'+ currentParent);
        //Loop though all children
        string finalQueryString = '';
        List<sObject> queryOutput = new List<sObject>();
        while ( !endOfStructure ){  
            if( level == 0 ){
                finalQueryString = queryString + ' where id IN : CurrentParent ORDER BY '+parentFieldAPIname +'  Limit 1000';
            } 
            else {
                finalQueryString = queryString + ' where ParentID IN : CurrentParent ORDER BY '+parentFieldAPIname+' Limit 1000';
            }
            system.debug('********finalQueryString:'+finalQueryString);
            if(finalQueryString != null && finalQueryString !=''){
                try{
                    if(Limits.getLimitQueries()-Limits.getQueries()>0){
                        queryOutput = database.query(finalQueryString);
                    }else{
                        endOfStructure = true;
                    }
                }catch(exception ex){ 
                    endOfStructure = true;
                }
            }
            if( queryOutput.size() == 0 ){
                endOfStructure = true;
            }
            else{
                currentParent.clear();
                //iterating through query output
                for ( Integer i = 0 ; i < queryOutput.size(); i++ ){
                    sobject sb= queryOutput[i];
                    currentParent.add(string.valueof(sb.get('id')) );
                    HierarchyData ss = new HierarchyData();
                    if(sb.get('id') == recId || level == 0){
                        ss.expanded = true;
                    }else{
                        ss.expanded = false;
                    }
                    ss.rec = sb;
                    returnValue.add(ss);
                 }
            }
            level++;
        }
        system.debug('**********returnValue:'+returnValue);
        return returnValue;
    }
    
	// Find the tom most element in Heirarchy  
    // @return objId
    public static String GetUltimateParentId( string recId, string sObjectName, string parentFieldAPIname ){
        Boolean top = false;
        while ( !top ) {
            string queryString = 'select id , ' +parentFieldAPIname+ ' from '+sObjectName + ' where Id =:recId LIMIT 1';
            system.debug('**********queryString GetUltimateParentId:'+queryString);
            sobject sb = database.query(queryString);
            
            if ( sb.get(parentFieldAPIname) != null ) {
                recId = string.valueof(sb.get(parentFieldAPIname));
            }else {
                top = true;
            }
        }
        return recId ;
    }

    public class HierarchyData{
        @AuraEnabled
        public sObject rec{get;set;}
        @AuraEnabled
        public boolean expanded{get;set;}
        public HierarchyData(){
            expanded = false;
        }
    }
}

SK_LightningTreeCmpController.js

({
    doInit: function (component, event, helper) { 
        console.log('doInit of component called');
        var trecid = component.get('v.ltngcurrentRecId');
        var tsObjectName= component.get('v.ltngSobjectname');
        var tparentFieldAPIname= component.get('v.ltngParentFieldAPIName');
        var tlabelFieldAPIName= component.get('v.ltngLabelFieldAPIName');
        helper.callToServer(
            component,
            "c.findHierarchyData",
            function(response) {
                var apexResponse = response;
                var roles = {};
                console.log('*******apexResponse:'+JSON.stringify(apexResponse));
                var results = apexResponse;
                
                roles[undefined] = { Name: "Root", items: [] };
                
                apexResponse.forEach(function(v) {
                    //var sk = v.rec;
                    //var labelValue= sk[tlabelFieldAPIName];
                    roles[v.rec.Id] = { 
                        label: v.rec[tlabelFieldAPIName] ,
                        name: v.rec.Id, 
                        expanded:v.expanded,
                        items: [] };
                });
                console.log('*******roles:'+JSON.stringify(roles));
                apexResponse.forEach(function(v) {
                    roles[v.rec[tparentFieldAPIname]].items.push(roles[v.rec.Id]);   
                });                
                component.set("v.items", roles[undefined].items);
                console.log('*******roles[undefined].items:'+JSON.stringify(roles[undefined].items));
            }, 
            {
                recId: component.get('v.ltngcurrentRecId'),
                sObjectName: component.get('v.ltngSobjectname'),
                parentFieldAPIname: component.get('v.ltngParentFieldAPIName'),
                labelFieldAPIName: component.get('v.ltngLabelFieldAPIName')
            }
        );
        
    },
    showSpinner: function(component, event, helper) {
        var spinner = component.find("spinner");
        $A.util.removeClass(spinner, "slds-hide");
    },
    hideSpinner: function(component, event, helper) {
        var spinner = component.find("spinner");
        $A.util.addClass(spinner, "slds-hide");
    },
    handleSelect: function (cmp, event, helper) {
        //return name of selected tree item
        var myName = event.getParam('name');
        alert("You selected: " + myName);
    }
})

SK_LightningTreeCmpHelper.js

({
	callToServer : function(component, method, callback, params) {
        console.log('Calling helper callToServer function');
		var action = component.get(method);
        if(params){
            action.setParams(params);
        }
        console.log(JSON.stringify(params));
        action.setCallback(this, function(response) {
            var state = response.getState();
            if (state === "SUCCESS") {
                //alert('Processed successfully at server');
                callback.call(this,response.getReturnValue());
            }else if(state === "ERROR"){
                alert('Problem with connection. Please try again.');
            }
        });
		$A.enqueueAction(action);
    }
})

Is there a sample TestClass for the Apex controller?