diff --git a/buildCodonSM.sh b/buildCodonSM.sh
index f308452..85bc1cd 100755
--- a/buildCodonSM.sh
+++ b/buildCodonSM.sh
@@ -3,12 +3,16 @@
 BEASTV="2.7"
 PKG="CodonSubstModels"
 VERSION="2.0.0"
-# mkdir ~/Library/Application\ Support/BEAST/$BEASTV/$PKG
+if [ ! -d "~/Library/Application Support/BEAST/$BEASTV/$PKG" ]; then
+  mkdir ~/Library/Application\ Support/BEAST/$BEASTV/$PKG
+fi
 ant build
 cp dist/$PKG.v$VERSION.zip ~/Library/Application\ Support/BEAST/$BEASTV/$PKG/tmp.zip
 cd ~/Library/Application\ Support/BEAST/$BEASTV/$PKG/
 unzip -o tmp.zip
 cd ~/WorkSpace/$PKG
 
+ls -l ~/Library/Application\ Support/BEAST/$BEASTV/$PKG
+
 echo "----Done----"
 
diff --git a/fxtemplates/CodonSubstModels.xml b/fxtemplates/CodonSubstModels.xml
index 9fd51b8..946b900 100644
--- a/fxtemplates/CodonSubstModels.xml
+++ b/fxtemplates/CodonSubstModels.xml
@@ -104,7 +104,7 @@
             codonmodels.M0Model.rates,
 	        codonmodels.M0Model.eigenSystem,
 	        codonmodels.M0Model.verbose'
-        buttonLabelMap='beast.app.beauti.BeautiInitDlg.>> details=Edit parameters'
+        buttonLabelMap='beastfx.app.beauti.BeautiInitDlg.>> details=Edit parameters'
     >
         <!--disableMenus='Mode'-->
 
@@ -169,7 +169,7 @@
 		<alignmentProvider id="Codon Alignment" spec='codonmodels.app.beauti.CodonAlignmentProvider' template='@CodonPartitionTemplate'/>
 
 
-        <partitiontemplate id='CodonPartitionTemplate' spec='BeautiSubTemplate' class='beast.base.evolution.likelihood.ThreadedTreeLikelihood' mainid='mcmc'>
+        <partitiontemplate id='CodonPartitionTemplate' spec='beastfx.app.inputeditor.BeautiSubTemplate' class='beast.base.evolution.likelihood.ThreadedTreeLikelihood' mainid='mcmc'>
 <![CDATA[
             <!-- site model                                                              -->
             <plugin spec='SiteModel' id="SiteModel.s:$(n)" gammaCategoryCount='0'>
@@ -247,7 +247,7 @@
 
             <connect srcID='treelog.t:$(n)'                 targetID='mcmc' inputName='logger' if='inposterior(Tree.t:$(n)) and Tree.t:$(n)/estimate=true'/>
             <connect srcID='RandomTree.t:$(n)'                targetID='mcmc' inputName='init'   if='isInitializing'/>
-			<connect method="beast.app.beauti.StateNodeInitialiserListInputEditor.customConnector"/>
+			<connect method="beastfx.app.beauti.StateNodeInitialiserListInputEditor.customConnector"/>
 
             <connect srcID='ClockPrior.c:$(n)'                targetID='prior' inputName='distribution' if='inlikelihood(clockRate.c:$(n)) and clockRate.c:$(n)/estimate=true'/>
 <!-- when FixMeanMutationRatesOperator, the prior is uniform due to the operator -->
diff --git a/src/codonmodels/evolution/alignment/CodonAlignment.java b/src/codonmodels/evolution/alignment/CodonAlignment.java
index e70e218..6c3e645 100644
--- a/src/codonmodels/evolution/alignment/CodonAlignment.java
+++ b/src/codonmodels/evolution/alignment/CodonAlignment.java
@@ -92,8 +92,8 @@ public CodonAlignment(Alignment alignment, GeneticCode geneticCode) {
         if (this.m_dataType == null || !this.m_dataType.getTypeDescription().equals(Codon.CODON)) {
             DataType oldType = this.m_dataType;
             this.m_dataType = new Codon(geneticCode);
-            Log.warning.println("Warning: CodonAlignment (" + this.getID() + ") original data type was " +
-                    oldType + ", is corrected to " + this.m_dataType + " - " + geneticCode.getDescription() + " !");
+            Log.warning.println("Warning: find CodonAlignment (" + this.getID() + ") is not a Codon data type (" +
+                    oldType + "), so set it to : " + this.m_dataType + " - " + geneticCode.getDescription() + " !");
         }
         initAndValidate();
     }
diff --git a/src/codonmodels/evolution/datatype/Codon.java b/src/codonmodels/evolution/datatype/Codon.java
index b2a46c6..f7e02e8 100644
--- a/src/codonmodels/evolution/datatype/Codon.java
+++ b/src/codonmodels/evolution/datatype/Codon.java
@@ -28,9 +28,10 @@
 
 import beast.base.core.Input;
 import beast.base.core.Log;
+import beast.base.evolution.datatype.DataType;
+import beast.base.evolution.datatype.Nucleotide;
 
 import java.util.Arrays;
-import beast.base.evolution.datatype.DataType;
 
 /**
  * Implements DataType for codon,
diff --git a/test/codonmodels/evolution/alignment/CodonAlignmentTest.java b/test/codonmodels/evolution/alignment/CodonAlignmentTest.java
index 843accd..a6e33ed 100644
--- a/test/codonmodels/evolution/alignment/CodonAlignmentTest.java
+++ b/test/codonmodels/evolution/alignment/CodonAlignmentTest.java
@@ -2,13 +2,18 @@
 
 import beast.base.evolution.alignment.Alignment;
 import beast.base.evolution.alignment.Sequence;
+import beast.base.evolution.datatype.DataType;
+import beast.pkgmgmt.BEASTClassLoader;
 import codonmodels.evolution.CodonData;
 import codonmodels.evolution.datatype.Codon;
 import org.junit.Test;
 
 import java.util.List;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
-import static junit.framework.Assert.assertEquals;
 
 /**
  * @author Walter Xie
@@ -30,6 +35,13 @@ public void testGetSequence() {
         }
     }
 
+    @Test
+    public void testGettingDataType() {
+        Set<String> dataTypes = BEASTClassLoader.loadService(DataType.class);
+        System.out.println(dataTypes);
+        assertTrue("Cannot load Codon data type !", dataTypes.contains(Codon.CODON));
+    }
+
     @Test
     public void testCodonUsage() {
         Alignment data = CodonData.getAlig6T333();
diff --git a/test/codonmodels/evolution/datatype/CodonTest.java b/test/codonmodels/evolution/datatype/CodonTest.java
index 497f15b..23acb65 100644
--- a/test/codonmodels/evolution/datatype/CodonTest.java
+++ b/test/codonmodels/evolution/datatype/CodonTest.java
@@ -10,7 +10,7 @@
 import java.util.Arrays;
 import java.util.List;
 
-import static junit.framework.Assert.*;
+import static org.junit.Assert.*;
 
 /**
  * @author Walter Xie
diff --git a/test/codonmodels/evolution/datatype/GeneticCodeTest.java b/test/codonmodels/evolution/datatype/GeneticCodeTest.java
index b6af0e1..866a1ae 100644
--- a/test/codonmodels/evolution/datatype/GeneticCodeTest.java
+++ b/test/codonmodels/evolution/datatype/GeneticCodeTest.java
@@ -6,7 +6,6 @@
 
 import java.util.Arrays;
 
-import static junit.framework.Assert.assertEquals;
 import static org.junit.Assert.*;
 
 /**
diff --git a/test/codonmodels/evolution/likelihood/DALikelihoodTest.java b/test/codonmodels/evolution/likelihood/DALikelihoodTest.java
index c81eb98..9d2bc97 100644
--- a/test/codonmodels/evolution/likelihood/DALikelihoodTest.java
+++ b/test/codonmodels/evolution/likelihood/DALikelihoodTest.java
@@ -14,7 +14,7 @@
 
 import java.util.Arrays;
 
-import static junit.framework.Assert.assertEquals;
+import static org.junit.Assert.assertEquals;
 
 
 /**
diff --git a/test/codonmodels/evolution/likelihood/DAStatesLikelihoodCoreIndexTest.java b/test/codonmodels/evolution/likelihood/DAStatesLikelihoodCoreIndexTest.java
index b9166e8..f53ae0f 100644
--- a/test/codonmodels/evolution/likelihood/DAStatesLikelihoodCoreIndexTest.java
+++ b/test/codonmodels/evolution/likelihood/DAStatesLikelihoodCoreIndexTest.java
@@ -6,7 +6,7 @@
 
 import java.util.Arrays;
 
-import static junit.framework.Assert.assertEquals;
+import static org.junit.Assert.assertEquals;
 
 /**
  * Build transition probability matrices under 3-node tree,
diff --git a/test/codonmodels/evolution/likelihood/M0TreeLikelihoodTest.java b/test/codonmodels/evolution/likelihood/M0TreeLikelihoodTest.java
index d1b8f84..27a4c5e 100644
--- a/test/codonmodels/evolution/likelihood/M0TreeLikelihoodTest.java
+++ b/test/codonmodels/evolution/likelihood/M0TreeLikelihoodTest.java
@@ -8,7 +8,7 @@
 import org.junit.Before;
 import org.junit.Test;
 
-import static junit.framework.Assert.assertEquals;
+import static org.junit.Assert.assertEquals;
 
 /**
  * @author Walter Xie
diff --git a/test/codonmodels/evolution/substmodel/M0ModelTest.java b/test/codonmodels/evolution/substmodel/M0ModelTest.java
index 5f00adc..7ba1ae6 100644
--- a/test/codonmodels/evolution/substmodel/M0ModelTest.java
+++ b/test/codonmodels/evolution/substmodel/M0ModelTest.java
@@ -14,7 +14,7 @@
 import java.util.Arrays;
 import java.util.stream.DoubleStream;
 
-import static junit.framework.Assert.assertEquals;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 /**
diff --git a/version.xml b/version.xml
index 9a29608..2dce05f 100644
--- a/version.xml
+++ b/version.xml
@@ -12,6 +12,7 @@
         <provider classname="codonmodels.CodonSubstitutionModel"/>
         <provider classname="codonmodels.evolution.alignment.CodonAlignment"/>
         <provider classname="codonmodels.evolution.datatype.Codon"/>
+<!--        <provider classname="codonmodels.evolution.datatype.GeneticCode"/>-->
         <provider classname="codonmodels.evolution.likelihood.DataAugTreeLikelihood"/>
         <provider classname="codonmodels.evolution.likelihood.GenericDATreeLikelihood"/>
         <provider classname="codonmodels.evolution.operators.GibbsSamplingOperator"/>