--- a/refactoring_custom/SmallSense__CustomTestCaseSetUpCodeGenerator.st	Fri Jul 07 11:29:55 2017 +0200
+++ b/refactoring_custom/SmallSense__CustomTestCaseSetUpCodeGenerator.st	Mon Aug 07 16:06:08 2017 +0100
@@ -124,48 +124,37 @@
 
 !CustomTestCaseSetUpCodeGenerator methodsFor:'executing'!
 
-buildForClass: class
-    | source category superHasSetup current package |
+buildForClass:class 
+    | source  category  superHasSetup  current  package |
 
     current := class.
     superHasSetup := false.
-    [ superHasSetup not and:[ current isNil not ] ] whileTrue:[
-        superHasSetup := current includesSelector: #setUp.
+    [
+        superHasSetup not and:[ current isNil not ]
+    ] whileTrue:[
+        superHasSetup := current includesSelector:#setUp.
         superHasSetup ifFalse:[
             current := current superclass.
         ].
     ].
-    superHasSetup ifTrue:[ 
-        source := 'setUp
-    super setUp.
-
-    "Add your own code here..."
-'.
-        category := (current compiledMethodAt: #setUp) category.
-    ] ifFalse:[ 
-        source := 'setUp
-    "/ super setUp.
-
-    "Add your own code here..."
-'.
-        category := ((Smalltalk at:#TestCase) compiledMethodAt: #setUp) category.
+    source := self methodSource:superHasSetup.
+    superHasSetup ifTrue:[
+        category := (current compiledMethodAt:#setUp) category.
+    ] ifFalse:[
+        category := ((Smalltalk at:#TestCase) compiledMethodAt:#setUp) category.
     ].
-
     package := PackageId noProjectID.
-    samePackageAsTestedClass ? self defaultSamePackageAsTestedClass ifTrue: [ 
+    samePackageAsTestedClass ? self defaultSamePackageAsTestedClass ifTrue:[
         package := class package
     ].
-
-    model createMethod
-        class: class;
-        source: source;
-        category: category;
-        package: package;
+    (model createMethod)
+        class:class;
+        source:source;
+        category:category;
+        package:package;
         compile.
 
-    "Created: / 05-08-2014 / 14:17:17 / Jan Vrany <jan.vrany@fit.cvut.cz>"
-    "Modified: / 31-01-2015 / 18:32:53 / Jakub Nesveda <nesvejak@fit.cvut.cz>"
-    "Modified: / 12-06-2015 / 20:46:10 / Jan Vrany <jan.vrany@fit.cvut.cz>"
+    "Created: / 07-08-2017 / 14:21:27 / Jan Vrany <jan.vrany@fit.cvut.cz>"
 !
 
 buildInContext:aCustomContext
@@ -178,3 +167,24 @@
     "Modified: / 12-06-2015 / 20:46:16 / Jan Vrany <jan.vrany@fit.cvut.cz>"
 ! !
 
+!CustomTestCaseSetUpCodeGenerator methodsFor:'private'!
+
+methodSource:callSuper 
+    | source |
+
+    callSuper ifTrue:[
+        source := 'setUp
+    super setUp.
+
+    "Add your own code here..."
+'.
+    ] ifFalse:[
+        source := 'setUp
+    "/ super setUp.
+
+    "Add your own code here..."
+'.
+    ].
+    ^ source
+! !
+

--- a/refactoring_custom/SmallSense__CustomUITestCaseSetUpCodeGenerator.st	Fri Jul 07 11:29:55 2017 +0200
+++ b/refactoring_custom/SmallSense__CustomUITestCaseSetUpCodeGenerator.st	Mon Aug 07 16:06:08 2017 +0100
@@ -66,3 +66,37 @@
     "Created: / 16-09-2014 / 11:23:50 / Jan Vrany <jan.vrany@fit.cvut.cz>"
 ! !
 
+!CustomUITestCaseSetUpCodeGenerator methodsFor:'private'!
+
+methodSource:callSuper 
+    | source |
+
+    callSuper ifTrue:[
+        source := 'setUp
+    self skipIf: Screen current isNil description: ''No display connection''.
+    Smalltalk loadPackage: ''stx:goodies/sunit/ext/ui''.
+    super setUp.
+
+    "Add your own code here..."
+'.
+    ] ifFalse:[
+        source := 'setUp
+    self skipIf: Screen current isNil description: ''No display connection''.
+    Smalltalk loadPackage: ''stx:goodies/sunit/ext/ui''.
+    "/ super setUp.
+
+    "Add your own code here..."
+'.
+    ].
+    ^ source
+
+    "Created: / 07-08-2017 / 15:09:49 / Jan Vrany <jan.vrany@fit.cvut.cz>"
+! !
+
+!CustomUITestCaseSetUpCodeGenerator class methodsFor:'documentation'!
+
+version_HG
+
+    ^ '$Changeset: <not expanded> $'
+! !
+