Class diagram & fix relation between units

extension-objects/src/main/java/units/CommanderUnit.java

@@ -9,8 +9,6 @@ import concreteextensions.Commander;
  */
 public class CommanderUnit extends Unit {
 
-  private CommanderExtension commanderExtension;
-
   public CommanderUnit(String name) {
     super(name);
   }
@@ -19,10 +17,10 @@ public class CommanderUnit extends Unit {
   public UnitExtension getUnitExtension(String extensionName) {
 
     if (extensionName.equals("CommanderExtension")) {
-      if (commanderExtension == null) {
-        commanderExtension = new Commander(this);
+      if (unitExtension == null) {
+        unitExtension = new Commander(this);
       }
-      return commanderExtension;
+      return unitExtension;
     }
 
     return super.getUnitExtension(extensionName);

extension-objects/src/main/java/units/SergeantUnit.java

@@ -9,8 +9,6 @@ import concreteextensions.Sergeant;
  */
 public class SergeantUnit extends Unit {
 
-  private SergeantExtension sergeantExtension;
-
   public SergeantUnit(String name) {
     super(name);
   }
@@ -19,10 +17,10 @@ public class SergeantUnit extends Unit {
   public UnitExtension getUnitExtension(String extensionName) {
 
     if (extensionName.equals("SergeantExtension")) {
-      if (sergeantExtension == null) {
-        sergeantExtension = new Sergeant(this);
+      if (unitExtension == null) {
+        unitExtension = new Sergeant(this);
       }
-      return sergeantExtension;
+      return unitExtension;
     }
 
     return super.getUnitExtension(extensionName);

extension-objects/src/main/java/units/SoldierUnit.java

@@ -9,8 +9,6 @@ import concreteextensions.Soldier;
  */
 public class SoldierUnit extends Unit {
 
-  private SoldierExtension soldierExtension;
-
   public SoldierUnit(String name) {
     super(name);
   }
@@ -19,11 +17,11 @@ public class SoldierUnit extends Unit {
   public UnitExtension getUnitExtension(String extensionName) {
 
     if (extensionName.equals("SoldierExtension")) {
-      if (soldierExtension == null) {
-        soldierExtension = new Soldier(this);
+      if (unitExtension == null) {
+        unitExtension = new Soldier(this);
       }
 
-      return soldierExtension;
+      return unitExtension;
     }
     return super.getUnitExtension(extensionName);
   }

extension-objects/src/main/java/units/Unit.java

@@ -8,6 +8,7 @@ import abstractextensions.UnitExtension;
 public class Unit {
 
   private String name;
+  protected UnitExtension unitExtension = null;
 
   public Unit(String name) {
     this.name = name;

extension-objects/src/test/java/units/UnitTest.java

@@ -11,8 +11,6 @@ import static org.junit.Assert.*;
  */
 public class UnitTest {
 
-  private Unit unit;
-
   @Test
   public void testConstGetSet() throws Exception {
     final String name = "testName";