import org.junit.Before;
import org.junit.Test;
import streams.*;

import java.util.Iterator;
import java.util.stream.*;
import java.util.stream.Stream;

import static org.junit.Assert.assertEquals;

/**
 * Authors:
 * Aggelos Biboudis (@biboudis)
 * Nick Palladinos (@NickPalladinos)
 */
public class TestAlgebrasFlatMapAlgebras {

    public Long[] v_outer, v_inner;
    public Long expected;

    private <C> Long query(StreamTerminalAlg<Id.t, C> algebra){
        Long value = Id.prj(
                algebra.<Long>count(
                        algebra.flatMap(x -> {
                            System.out.print("outer ");
                            return algebra.<Long, Long>map(y -> {
                                System.out.print("inner ");
                                return x * y;
                            }, algebra.source(v_inner));
                        },  algebra.source(v_outer)))).value;

        System.out.println();

        return value;
    }

    @Before
    public void setUp() {
        v_outer = IntStream.range(0, 5).mapToObj(i -> (long) (i % 5)).toArray(Long[]::new);
        v_inner = IntStream.range(0, 2).mapToObj(i -> (long) (i % 5)).toArray(Long[]::new);
        expected = java.util.stream.Stream.of(v_outer)
                .flatMap(x -> java.util.stream.Stream.of(v_inner).map(y -> x * y))
                .count();
    }

    @Test
    public void testPull(){
        PullFactory algebra = new PullFactory();

        Long actual = query(algebra);

        assertEquals(expected, actual);
    }

    @Test
    public void testPush(){
        PushFactory algebra = new PushFactory();

        Long actual = query(algebra);

        assertEquals(expected, actual);
    }

    @Test
    public void testJava8StreamsPush(){
        Long actual = java.util.stream.Stream.of(v_outer)
                .flatMap(x -> {
                    System.out.print("outer ");
                    return java.util.stream.Stream.of(v_inner).map(y -> {
                        System.out.print("inner ");
                        return x * y;
                    });
                })
                .count();

        assertEquals(expected, actual);
    }

    @Test
    public void testJava8StreamsPushWithPull(){
        Long actual = 0L;

        Iterator<Long> iterator = Stream.of(v_outer)
                .flatMap(x -> {
                    System.out.print("outer ");
                    return Stream.of(v_inner).map(y -> {
                        System.out.print("inner ");
                        return x * y;
                    });
                }).iterator();

        while(iterator.hasNext()){
            actual++;
            iterator.next();
        }

        assertEquals(expected, actual);
    }
}

@@ -393,593 +393,8 @@
 d;%0A%0A
-    private %3CC%3E Long query(StreamTerminalAlg%3CId.t, C%3E algebra)%7B%0A        Long value = Id.prj(%0A                algebra.%3CLong%3Ecount(%0A                        algebra.flatMap(x -%3E %7B%0A                            System.out.print(%22outer %22);%0A                            return algebra.%3CLong, Long%3Emap(y -%3E %7B%0A                                System.out.print(%22inner %22);%0A                                return x * y;%0A                            %7D, algebra.source(v_inner));%0A                        %7D,  algebra.source(v_outer)))).value;%0A%0A        System.out.println();%0A%0A        return value;%0A    %7D%0A%0A
     
@@ -560,9 +560,9 @@
 (0, 
-2
+5
 ).ma
@@ -882,719 +882,372 @@
     
-Long actual = query(algebra);%0A%0A        assertEquals(expected, actual
+App%3CPull.t, Long%3E map = algebra.map(y -%3E %7B%0A            System.out.println(%22inner: %22 + y
 );%0A    
-%7D%0A%0A
     
-@Test%0A    public void testPush()%7B%0A        PushFactory algebra = new PushFactory();%0A%0A        Long actual = query(algebra);%0A%0A        assertEquals(expected, actual);%0A    %7D%0A%0A    @Test%0A    public void testJava8StreamsPush()%7B%0A        Long actual = java.util.stream.Stream.of(v_outer)%0A                .flatMap(x -%3E %7B%0A                    System.out.print(%22outer %22);%0A                    return java.util.stream.Stream.of(v_inner).map(y -%3E %7B%0A                        System.out.print(%22inner %22);%0A                        return x * y;%0A                    %7D);%0A                %7D)%0A                .count();%0A%0A        assertEquals(expected, actual
+    return (long) 10 * (long) y;%0A        %7D, algebra.source(v_inner));%0A%0A        App%3CPull.t, Long%3E flatMap = algebra.flatMap(y -%3E map, algebra.source(v_outer));%0A%0A        Pull%3CLong%3E prj = Pull.prj(flatMap);%0A%0A        prj.hasNext();%0A%0A        System.out.println(prj.next()
 );%0A 
@@ -1314,35 +1314,8 @@
 ()%7B%0A
-        Long actual = 0L;%0A%0A
     
@@ -1399,56 +1399,8 @@
 %3E %7B%0A
-                    System.out.print(%22outer %22);%0A
     
@@ -1496,17 +1496,24 @@
 rint
+ln
 (%22inner
+:
  %22
+ + y
 );%0A 
@@ -1617,14 +1617,8 @@
     
-while(
 iter
@@ -1631,18 +1631,18 @@
 asNext()
-)%7B
+;%0A
 %0A       
@@ -1646,99 +1646,42 @@
     
-    actual++;%0A            iterator.next();%0A        %7D%0A%0A        assertEquals(expected, actual
+System.out.println(iterator.next()
 );%0A 

Update Test FlatMap.