Fix multipart transformer in package:mime

dart/pkg/mime/test/mime_multipart_transformer_test.dart

 // Copyright (c) 2013, the Dart project authors.  Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
 import 'dart:async';
 import 'dart:math';
 
 import "package:unittest/unittest.dart";
 import "package:mime/mime.dart";
 
+void _writeInChunks(List<int> data,
+                    int chunkSize,
+                    StreamController<List<int>> controller) {
+  if (chunkSize == -1) chunkSize = data.length;
+
+  int written = 0;
+  for (int pos = 0; pos < data.length; pos += chunkSize) {
+    int remaining = data.length - pos;
+    int writeLength = min(chunkSize, remaining);
+    controller.add(data.sublist(pos, pos + writeLength));
+    written += writeLength;
+  }
+  controller.close();
+}
+
+
 void _testParse(String message,
                String boundary,
                [List<Map> expectedHeaders,
                 List expectedParts,
                 bool expectError = false]) {
   Future testWrite(List<int> data, [int chunkSize = -1]) {
     StreamController controller = new StreamController(sync: true);
 
     var stream = controller.stream.transform(
         new MimeMultipartTransformer(boundary));
@@ skipping 13 matching lines @@
           }));
     }, onError: (error) {
       if (!expectError) throw error;
     }, onDone: () {
       if (expectedParts != null) {
         expect(i, equals(expectedParts.length));
       }
       Future.wait(futures).then(completer.complete);
     });
 
-    if (chunkSize == -1) chunkSize = data.length;
-
-    int written = 0;
-    for (int pos = 0; pos < data.length; pos += chunkSize) {
-      int remaining = data.length - pos;
-      int writeLength = min(chunkSize, remaining);
-      controller.add(data.sublist(pos, pos + writeLength));
-      written += writeLength;
-    }
-    controller.close();
+    _writeInChunks(data, chunkSize, controller);
 
     return completer.future;
   }
+
+  Future testFirstPartOnly(List<int> data, [int chunkSize = -1]) {
+    var completer = new Completer();
+    var controller = new StreamController(sync: true);
+
+    var stream = controller.stream.transform(
+        new MimeMultipartTransformer(boundary));
+
+    var subscription;
+    subscription = stream.first.then((multipart) {
+      if (expectedHeaders != null) {
+        expect(multipart.headers, equals(expectedHeaders[0]));
+      }
+      return (multipart.fold([], (b, d) => b..addAll(d)).then((data) {
+        if (expectedParts != null && expectedParts[0] != null) {
+          expect(data, equals(expectedParts[0].codeUnits));
+        }
+      }));
+    }).then((_) {
+      completer.complete();
+    });
+
+    _writeInChunks(data, chunkSize, controller);
+
+    return completer.future;
+  }
+
+  Future testCompletePartAfterCancel(List<int> data,
+                                     int parts,
+                                     [int chunkSize = -1]) {
+    var completer = new Completer();
+    var controller = new StreamController(sync: true);
+    var stream = controller.stream.transform(
+        new MimeMultipartTransformer(boundary));
+    var subscription;
+    int i = 0;
+    var futures = [];
+    subscription = stream.listen((multipart) {
+      int partIndex = i;
+
+      if (partIndex >= parts) {
+        throw 'Expected no more parts, but got one.';
+      }
+
+      if (expectedHeaders != null) {
+        expect(multipart.headers, equals(expectedHeaders[partIndex]));
+      }
+      futures.add((multipart.fold([], (b, d) => b..addAll(d)).then((data) {
+        if (expectedParts != null && expectedParts[partIndex] != null) {
+          expect(data, equals(expectedParts[partIndex].codeUnits));
+        }
+      })));
+
+      if (partIndex == (parts - 1)) {
+        subscription.cancel();
+        Future.wait(futures).then(completer.complete);
+      }
+      i++;
+    });
+
+    _writeInChunks(data, chunkSize, controller);
+
+    return completer.future;
+  }
 
   // Test parsing the data three times delivering the data in
   // different chunks.
   List<int> data = message.codeUnits;
   test('test', () {
     expect(Future.wait([
         testWrite(data),
         testWrite(data, 10),
         testWrite(data, 2),
-        testWrite(data, 1)]),
-        completes);
+        testWrite(data, 1),
+    ]), completes);
   });
+
+  if (expectedParts.length > 0) {
+    test('test-first-part-only', () {
+      expect(Future.wait([
+          testFirstPartOnly(data),
+          testFirstPartOnly(data, 10),
+          testFirstPartOnly(data, 2),
+          testFirstPartOnly(data, 1),
+      ]), completes);
+    });
+
+    test('test-n-parts-only', () {
+      int numPartsExpected = expectedParts.length - 1;
+      if (numPartsExpected == 0) numPartsExpected = 1;
+
+      expect(Future.wait([
+        testCompletePartAfterCancel(data, numPartsExpected),
+        testCompletePartAfterCancel(data, numPartsExpected, 10),
+        testCompletePartAfterCancel(data, numPartsExpected, 2),
+        testCompletePartAfterCancel(data, numPartsExpected, 1),
+      ]), completes);
+    });
+  }
 }
 
 void _testParseValid() {
   // Empty message from Chrome form post.
   var message = '------WebKitFormBoundaryU3FBruSkJKG0Yor1--\r\n';
   _testParse(message, "----WebKitFormBoundaryU3FBruSkJKG0Yor1", [], []);
 
   // Sample from Wikipedia.
   message = """
 This is a message with multiple parts in MIME format.\r
@@ skipping 226 matching lines @@
 \r
 Body2\r
 --xxx\r\n""";
   _testParse(message, "xxx", null, [null, null], true);
 }
 
 void main() {
   _testParseValid();
   _testParseInvalid();
 }