Add write to output stream

datafaker-net · Dec 15, 2024 · 72268ce · 72268ce
1 parent 0be88ad
commit 72268ce
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 8 deletions.
diff --git a/src/main/java/net/datafaker/transformations/CsvTransformer.java b/src/main/java/net/datafaker/transformations/CsvTransformer.java
@@ -44,7 +44,7 @@ public String generate(Iterable<IN> input, Schema<IN, ?> schema) {
         }
 
         StringBuilder sb = new StringBuilder();
-        generateHeader(schema, sb);
+        generateHeader(schema, sb, true);
 
         Iterator<IN> iterator = input.iterator();
         boolean hasNext = iterator.hasNext();
@@ -83,22 +83,24 @@ private void addCharSequence(StringBuilder sb, CharSequence charSequence) {
         sb.append(quote);
     }
 
-    private void generateHeader(Schema<?, ?> schema, StringBuilder sb) {
+    private void generateHeader(Schema<?, ?> schema, StringBuilder sb, boolean insertSeparator) {
         if (withHeader) {
             for (int i = 0; i < schema.getFields().length; i++) {
                 addLine(sb, schema.getFields()[i].getName());
                 if (i < schema.getFields().length - 1) {
                     sb.append(separator);
                 }
             }
-            sb.append(LINE_SEPARATOR);
+            if (insertSeparator) {
+                sb.append(LINE_SEPARATOR);
+            }
         }
     }
 
     @Override
     public String generate(Schema<IN, ?> schema, int limit) {
         StringBuilder sb = new StringBuilder();
-        generateHeader(schema, sb);
+        generateHeader(schema, sb, true);
         for (int i = 0; i < limit; i++) {
             sb.append(apply(null, schema, i));
             if (i < limit - 1) {
@@ -111,13 +113,13 @@ public String generate(Schema<IN, ?> schema, int limit) {
     @Override
     public String getStartStream(Schema<IN, ?> schema) {
         StringBuilder sb = new StringBuilder();
-        generateHeader(null, sb);
+        generateHeader(schema, sb, false);
         return sb.toString();
     }
 
     @Override
     public String getEndStream() {
-        return null;
+        return "";
     }
 
     public static class CsvTransformerBuilder<IN> {

diff --git a/src/main/java/net/datafaker/transformations/JavaObjectTransformer.java b/src/main/java/net/datafaker/transformations/JavaObjectTransformer.java
@@ -2,6 +2,7 @@
 
 import net.datafaker.sequence.FakeSequence;
 
+import java.io.OutputStream;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.RecordComponent;
@@ -138,6 +139,11 @@ public Stream<Object> generateStream(final Schema<Object, ?> schema, long limit)
                 .limit(limit);
     }
 
+    @Override
+    public void writeToOutputStream(OutputStream outputStream, Schema<Object, ?> schema, long limit) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
     @Override
     public Collection<Object> generate(Schema<Object, ?> schema, int limit) {
         return this.generateStream(schema, limit).collect(Collectors.toList());

diff --git a/src/main/java/net/datafaker/transformations/Transformer.java b/src/main/java/net/datafaker/transformations/Transformer.java
@@ -1,5 +1,8 @@
 package net.datafaker.transformations;
 
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
 import java.util.stream.Stream;
 
 public interface Transformer<IN, OUT> {
@@ -29,7 +32,6 @@ default String getElementSeparator() {
         return "";
     }
 
-
     default Stream<OUT> generateStream(final Schema<IN, ?> schema, long limit) {
         Item item = new Item(0);
         return Stream.generate(() -> {
@@ -40,7 +42,11 @@ default Stream<OUT> generateStream(final Schema<IN, ?> schema, long limit) {
             res.append(apply(null, schema, item.current));
 
             if (item.current == limit - 1) {
-                res.append(getLineSeparator()).append(getEndStream());
+                res.append(getLineSeparator());
+                final String endStream = getEndStream();
+                if (endStream != null) {
+                    res.append(endStream);
+                }
             } else {
                 if (!getElementSeparator().isEmpty()) {
                     res.append(getElementSeparator());
@@ -51,6 +57,17 @@ default Stream<OUT> generateStream(final Schema<IN, ?> schema, long limit) {
         }).limit(limit);
     }
 
+    default void writeToOutputStream(OutputStream outputStream, final Schema<IN, ?> schema, long limit) {
+        generateStream(schema, limit).forEach(item -> {
+            byte[] bytes = (item + System.lineSeparator()).getBytes(StandardCharsets.UTF_8);
+            try {
+                outputStream.write(bytes);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
     class Item {
         private long current;
 

diff --git a/src/test/java/net/datafaker/sequence/FakeStreamTest.java b/src/test/java/net/datafaker/sequence/FakeStreamTest.java
@@ -9,11 +9,21 @@
 import net.datafaker.transformations.Field;
 import net.datafaker.transformations.JsonTransformer;
 import net.datafaker.transformations.Schema;
+import net.datafaker.transformations.Transformer;
+import net.datafaker.transformations.sql.SqlTransformer;
+import org.assertj.core.util.Files;
 import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 
+import java.io.BufferedOutputStream;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Random;
 import java.util.function.Supplier;
@@ -335,4 +345,46 @@ void testIteratorInfinite() {
 
         assertThat(count).isEqualTo(amountOfElementsToTake);
     }
+
+    @ParameterizedTest(name = "{index}: {0}")
+    @MethodSource("inputForFilesCreatedTest")
+    void testCsvFilesCreated(final String testName, final Transformer transformer, final List<String> expected) {
+        BaseFaker faker = new BaseFaker(new Random(10L));
+        Name name = faker.name();
+        Schema<Object, ?> schema = Schema.of(
+            field("FirstName", name::firstName),
+            field("LastName", name::lastName)
+        );
+
+        File csvFile = Files.newTemporaryFile();
+        assertThat(csvFile).exists().isEmpty();
+
+        try (BufferedOutputStream fos = new BufferedOutputStream(new FileOutputStream(csvFile))) {
+            transformer.writeToOutputStream(fos, schema, 2);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+
+        assertThat(Files.linesOf(csvFile, StandardCharsets.UTF_8))
+            .containsAll(expected);
+    }
+
+    private static Stream<Arguments> inputForFilesCreatedTest() {
+        return Stream.of(
+            Arguments.of("csv", CsvTransformer.builder().build(), List.of(
+                "\"FirstName\";\"LastName\"",
+                "\"Willis\";\"Huels\"",
+                "\"Carlena\";\"Jenkins\"")),
+            Arguments.of("sql", SqlTransformer.builder().build(), List.of(
+                "INSERT INTO \"MyTable\" (\"FirstName\", \"LastName\") VALUES ('Willis', 'Huels');",
+                "INSERT INTO \"MyTable\" (\"FirstName\", \"LastName\") VALUES ('Carlena', 'Jenkins');"
+            )),
+            Arguments.of("json", JsonTransformer.builder().build(), List.of(
+                "[",
+                "{\"FirstName\": \"Willis\", \"LastName\": \"Huels\"},",
+                "{\"FirstName\": \"Carlena\", \"LastName\": \"Jenkins\"}",
+                "]"
+            ))
+        );
+    }
 }