1 files changed, 23 insertions, 13 deletions
diff --git a/src/cz/crcs/ectester/standalone/output/XMLTestWriter.java b/src/cz/crcs/ectester/standalone/output/XMLTestWriter.java
index 2a35ce3..9332759 100644
--- a/src/cz/crcs/ectester/standalone/output/XMLTestWriter.java
+++ b/src/cz/crcs/ectester/standalone/output/XMLTestWriter.java
@@ -5,10 +5,11 @@ import cz.crcs.ectester.common.test.TestSuite;
 import cz.crcs.ectester.common.test.Testable;
 import cz.crcs.ectester.common.util.ByteUtil;
 import cz.crcs.ectester.standalone.ECTesterStandalone;
-import cz.crcs.ectester.standalone.test.KeyAgreementTestable;
-import cz.crcs.ectester.standalone.test.KeyGeneratorTestable;
-import cz.crcs.ectester.standalone.test.SignatureTestable;
-import cz.crcs.ectester.standalone.test.StandaloneTestSuite;
+import cz.crcs.ectester.standalone.test.base.KeyAgreementTestable;
+import cz.crcs.ectester.standalone.test.base.KeyGeneratorTestable;
+import cz.crcs.ectester.standalone.test.base.SignatureTestable;
+import cz.crcs.ectester.standalone.test.base.StandaloneTestable;
+import cz.crcs.ectester.standalone.test.suites.StandaloneTestSuite;
 import org.w3c.dom.Element;
 
 import javax.xml.parsers.ParserConfigurationException;
@@ -97,18 +98,27 @@ public class XMLTestWriter extends BaseXMLTestWriter {
         return sigElem;
     }
 
+    private Element stageElement(StandaloneTestable t) {
+        Element result = doc.createElement("stage");
+        result.setTextContent(t.getStage().name());
+        return result;
+    }
+
     @Override
     protected Element testableElement(Testable t) {
         Element result = doc.createElement("test");
-        if (t instanceof KeyGeneratorTestable) {
-            result.setAttribute("type", "key-pair-generator");
-            result.appendChild(kgtElement((KeyGeneratorTestable) t));
-        } else if (t instanceof KeyAgreementTestable) {
-            result.setAttribute("type", "key-agreement");
-            result.appendChild(kaElement((KeyAgreementTestable) t));
-        } else if (t instanceof SignatureTestable) {
-            result.setAttribute("type", "signature");
-            result.appendChild(sigElement((SignatureTestable) t));
+        if (t instanceof StandaloneTestable) {
+            if (t instanceof KeyGeneratorTestable) {
+                result.setAttribute("type", "key-pair-generator");
+                result.appendChild(kgtElement((KeyGeneratorTestable) t));
+            } else if (t instanceof KeyAgreementTestable) {
+                result.setAttribute("type", "key-agreement");
+                result.appendChild(kaElement((KeyAgreementTestable) t));
+            } else if (t instanceof SignatureTestable) {
+                result.setAttribute("type", "signature");
+                result.appendChild(sigElement((SignatureTestable) t));
+            }
+            result.appendChild(stageElement((StandaloneTestable) t));
         }
         return result;
     }