2 files changed, 24 insertions, 11 deletions
diff --git a/src/test/java/org/zwobble/couscous/tests/backends/python/PythonCompilerEvalTests.java b/src/test/java/org/zwobble/couscous/tests/backends/python/PythonCompilerEvalTests.java
index 6f19eb3..e4a5ddf 100644
--- a/src/test/java/org/zwobble/couscous/tests/backends/python/PythonCompilerEvalTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/backends/python/PythonCompilerEvalTests.java
@@ -1,5 +1,7 @@
 package org.zwobble.couscous.tests.backends.python;
 
+import java.util.List;
+
 import org.zwobble.couscous.ast.ClassNode;
 import org.zwobble.couscous.ast.ExpressionNode;
 import org.zwobble.couscous.ast.MethodNode;
@@ -8,19 +10,29 @@ import org.zwobble.couscous.tests.BackendEvalTests;
 import org.zwobble.couscous.values.PrimitiveValue;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+
+import static java.util.Arrays.asList;
 
 import lombok.val;
 
 public class PythonCompilerEvalTests extends BackendEvalTests {
     @Override
-    protected PrimitiveValue evalExpression(ExpressionNode expression) {
+    protected PrimitiveValue evalExpression(
+            List<ClassNode> classes,
+            ExpressionNode expression) {
+        
+        val programNode = ClassNode.builder("Program")
+            .method(MethodNode.staticMethod("run")
+                .statement(new ReturnNode(expression))
+                .build())
+            .build();
+        
         val runner = new PythonMethodRunner();
+        
         return runner.runMethod(
-            ClassNode.builder("Program")
-                .method(MethodNode.staticMethod("run")
-                    .statement(new ReturnNode(expression))
-                    .build())
-                .build(),
+            ImmutableList.copyOf(Iterables.concat(classes, asList(programNode))),
+            programNode.getName(),
             "run",
             ImmutableList.of());
     }
diff --git a/src/test/java/org/zwobble/couscous/tests/backends/python/PythonMethodRunner.java b/src/test/java/org/zwobble/couscous/tests/backends/python/PythonMethodRunner.java
index 73c1c52..33de465 100644
--- a/src/test/java/org/zwobble/couscous/tests/backends/python/PythonMethodRunner.java
+++ b/src/test/java/org/zwobble/couscous/tests/backends/python/PythonMethodRunner.java
@@ -7,6 +7,7 @@ import java.nio.file.Files;
 import java.util.List;
 
 import org.zwobble.couscous.ast.ClassNode;
+import org.zwobble.couscous.ast.TypeName;
 import org.zwobble.couscous.backends.python.PythonCodeGenerator;
 import org.zwobble.couscous.backends.python.PythonCompiler;
 import org.zwobble.couscous.backends.python.PythonSerializer;
@@ -19,7 +20,6 @@ import com.google.common.base.Joiner;
 import com.google.common.io.CharStreams;
 
 import static java.lang.Integer.parseInt;
-import static java.util.Arrays.asList;
 import static org.zwobble.couscous.tests.util.ExtraFiles.deleteRecursively;
 import static org.zwobble.couscous.values.PrimitiveValues.value;
 
@@ -30,21 +30,22 @@ public class PythonMethodRunner implements MethodRunner {
     @Override
     @SneakyThrows
     public PrimitiveValue runMethod(
-            ClassNode classNode,
+            List<ClassNode> classNodes,
+            TypeName className,
             String methodName,
             List<PrimitiveValue> arguments) {
         val directoryPath = Files.createTempDirectory(null);
         val compiler = new PythonCompiler(directoryPath);
         try {
-            compiler.compile(asList(classNode));
+            compiler.compile(classNodes);
             
             val argumentsString = Joiner.on(", ").join(arguments.stream()
                 .map(PythonCodeGenerator::generateCode)
                 .map(PythonSerializer::serialize)
                 .iterator());
             
-            val program = "from " + classNode.getName().getQualifiedName() + " import " + classNode.getSimpleName() +
-                ";print(repr(" + classNode.getSimpleName() + "." + methodName + "(" + argumentsString + ")))";
+            val program = "from " + className.getQualifiedName() + " import " + className.getSimpleName() +
+                ";print(repr(" + className.getSimpleName() + "." + methodName + "(" + argumentsString + ")))";
             
             val process = new ProcessBuilder("python3.4", "-c", program)
                 .directory(directoryPath.toFile())