diff --git a/Assets/Prefabs/Sticks/farmersAssociation.prefab b/Assets/Prefabs/Sticks/farmersAssociation.prefab
index 974a4fb..d8bfae3 100644
--- a/Assets/Prefabs/Sticks/farmersAssociation.prefab
+++ b/Assets/Prefabs/Sticks/farmersAssociation.prefab
@@ -1105,7 +1105,15 @@ MonoBehaviour:
   m_Script: {fileID: 11500000, guid: edcd1782caa73a747b1f060d18ce0f18, type: 3}
   m_Name: 
   m_EditorClassIdentifier: 
-  upgradeUI: {fileID: 4817988606444742093, guid: 1d481c7093a8d2443a04261fd7c72be2, type: 3}
+  _upgradeUI: {fileID: 4817988606444742093, guid: 1d481c7093a8d2443a04261fd7c72be2, type: 3}
+  _upgradeUnitCardList:
+  - {fileID: 11400000, guid: 9ff4b606166b07d4bb3afb15e5a23f41, type: 2}
+  - {fileID: 11400000, guid: ee85fc25d3460c745806649e7f9d166a, type: 2}
+  - {fileID: 11400000, guid: 0334115ab72f6d5449497db1bb50377e, type: 2}
+  _upgradePrefabList:
+  - {fileID: 5531237550998824024, guid: 6cd87b398e7a0e94580f4fcbe2fd310a, type: 3}
+  - {fileID: 5157279992115123224, guid: 85534a2d6c2add54d864073914646192, type: 3}
+  - {fileID: 2230360378425099702, guid: ece16a59d4dca2246817fbb1b23f17ba, type: 3}
 --- !u!114 &8585520847943034727
 MonoBehaviour:
   m_ObjectHideFlags: 0
@@ -1122,9 +1130,6 @@ MonoBehaviour:
   _speed: 0
   _attack_damage: 1
   _attack_interval: 2
-  upgradeList:
-  - {fileID: 9177659942431061517, guid: 6cd87b398e7a0e94580f4fcbe2fd310a, type: 3}
-  - {fileID: 682208058972047871, guid: ece16a59d4dca2246817fbb1b23f17ba, type: 3}
 --- !u!50 &-1491803373025033585
 Rigidbody2D:
   serializedVersion: 4
diff --git a/Assets/Scripts/Ally/AllyUpgrade.cs b/Assets/Scripts/Ally/AllyUpgrade.cs
index 04e9de2..0ff739a 100644
--- a/Assets/Scripts/Ally/AllyUpgrade.cs
+++ b/Assets/Scripts/Ally/AllyUpgrade.cs
@@ -6,11 +6,11 @@ using UnityEngine.EventSystems;
 public class AllyUpgrade : MonoBehaviour, IPointerClickHandler
 {
     [SerializeField]
-    private GameObject upgradeUI;
+    private GameObject _upgradeUI;
     [SerializeField]
-    private List<UnitCard> upgradeUnitCardList = new List<UnitCard>();
+    private List<UnitCard> _upgradeUnitCardList = new List<UnitCard>();
     [SerializeField]
-    private List<GameObject> upgradePrefabList = new List<GameObject>();
+    private List<GameObject> _upgradePrefabList = new List<GameObject>();
     private GameObject _instance = null;
 
     public void OnPointerClick(PointerEventData eventData)
@@ -30,17 +30,13 @@ public class AllyUpgrade : MonoBehaviour, IPointerClickHandler
                     }
                     else
                     {
-                        _instance = Instantiate(upgradeUI, Camera.main.WorldToScreenPoint(transform.position), Quaternion.identity, canvas.transform);
+                        _instance = Instantiate(_upgradeUI, Camera.main.WorldToScreenPoint(transform.position), Quaternion.identity, canvas.transform);
                     }
                 }
                 Debug.Log(_instance.transform.position);
                 Debug.Log(Camera.main.WorldToScreenPoint(transform.position));
-
-                GameObject upgradeParent = _instance.transform.GetChild(0).gameObject;
-                for (int i = 0; i < upgradeParent.transform.childCount; i++)
-                {
-                    GameObject upgradeEnfant = upgradeParent.transform.GetChild(i).gameObject;
-                }
+                AssignUpgrades();
+                
             }
             else
             {
@@ -50,4 +46,17 @@ public class AllyUpgrade : MonoBehaviour, IPointerClickHandler
 
         }
     }
+
+    private void AssignUpgrades()
+    {
+        GameObject upgradeParent = _instance.transform.GetChild(0).gameObject;
+        // Assign upgrade Prefabs
+        for (int i = 0; i < upgradeParent.transform.childCount; i++)
+        {
+            GameObject upgradeEnfant = upgradeParent.transform.GetChild(i).gameObject;
+            UpgradePlacementButton button = upgradeEnfant.GetComponent<UpgradePlacementButton>();
+            button.SetUnit(_upgradeUnitCardList[i], _upgradePrefabList[i]);
+            Debug.Log("Button " + button.name);
+        }
+    }
 }
diff --git a/Assets/Scripts/UnitTree/UnitTreeManager.cs b/Assets/Scripts/UnitTree/UnitTreeManager.cs
index c1a2dce..fa61ce8 100644
--- a/Assets/Scripts/UnitTree/UnitTreeManager.cs
+++ b/Assets/Scripts/UnitTree/UnitTreeManager.cs
@@ -18,7 +18,6 @@ public class UnitTreeManager : MonoBehaviour
         //Instantiate UnitTreeDisplay rect 3 times
         for (int i = 0; i < 3; i++)
         {
-            //oogaDisplay.Add(GetComponent<UnitTreeDisplay>());
             GameObject rect = new GameObject();
             Vector3 rectPosition = m_Position;
             rectPosition.x += 10.0f * i;
@@ -26,7 +25,6 @@ public class UnitTreeManager : MonoBehaviour
             instance.AddComponent<RectTransform>();
             UnitTreeDisplay ooga = instance.AddComponent<UnitTreeDisplay>();
             ooga.SetPosition(rectPosition);
-            //oogaDisplay[i].SetPosition(rectPosition);
         }
 
         //Instantiate line (later curved line? https://forum.unity.com/threads/easy-curved-line-renderer-free-utility.391219/)
diff --git a/Assets/Scripts/UnitTree/UpgradePlacementButton.cs b/Assets/Scripts/UnitTree/UpgradePlacementButton.cs
index f6b4cdf..9473c01 100644
--- a/Assets/Scripts/UnitTree/UpgradePlacementButton.cs
+++ b/Assets/Scripts/UnitTree/UpgradePlacementButton.cs
@@ -58,4 +58,13 @@ public class UpgradePlacementButton : MonoBehaviour, IPointerClickHandler
         Debug.Log("Upgrade clicked...");
 
     }
+
+    public void SetUnit(UnitCard card, GameObject gameObject)
+    {
+        SetTextFor(_foodLabel, card.Food);
+        SetTextFor(_woodLabel, card.Wood);
+        SetTextFor(_rockLabel, card.Rock);
+        _unitCardInformation = card;
+        _prefab = gameObject;
+    }
 }
\ No newline at end of file