From 52759d85376bc32d289fe2bc7bd085ed9892e97d Mon Sep 17 00:00:00 2001
From: Howard Lopez <lopehowa@amazon.com>
Date: Wed, 22 May 2019 15:51:52 -0700
Subject: [PATCH] Add tests for proper header additions

---
 src/StreamRequestPayloadMiddleware.php       |   3 +-
 tests/StreamRequestPayloadMiddlewareTest.php | 132 ++++++++++++++++---
 2 files changed, 119 insertions(+), 16 deletions(-)

diff --git a/src/StreamRequestPayloadMiddleware.php b/src/StreamRequestPayloadMiddleware.php
index 1c6c1616dd..8b46369ec4 100644
--- a/src/StreamRequestPayloadMiddleware.php
+++ b/src/StreamRequestPayloadMiddleware.php
@@ -58,7 +58,8 @@ public function __invoke(CommandInterface $command, RequestInterface $request)
             // to be calculated and not already known
             if (empty($requiresLength)
                 && empty($contentLength)
-                && !empty($operation['v4-unsigned-body'])
+                && isset($operation['authtype'])
+                && $operation['authtype'] == 'v4-unsigned-body'
             ) {
                 $request = $request->withHeader('transfer-encoding', 'chunked');
 
diff --git a/tests/StreamRequestPayloadMiddlewareTest.php b/tests/StreamRequestPayloadMiddlewareTest.php
index 6b217c7986..d0d96ab7b9 100644
--- a/tests/StreamRequestPayloadMiddlewareTest.php
+++ b/tests/StreamRequestPayloadMiddlewareTest.php
@@ -3,12 +3,16 @@
 
 use Aws\Api\Service;
 use Aws\AwsClient;
+use Aws\ClientResolver;
+use Aws\CommandInterface;
 use Aws\HandlerList;
+use Aws\Middleware;
 use Aws\Result;
 use Aws\StreamRequestPayloadMiddleware;
 use GuzzleHttp\Psr7;
 use GuzzleHttp\Psr7\Request;
 use PHPUnit\Framework\TestCase;
+use Psr\Http\Message\RequestInterface;
 
 /**
  * @covers \Aws\StreamRequestPayloadMiddleware
@@ -16,29 +20,107 @@
 class StreamRequestPayloadMiddlewareTest extends TestCase
 {
 
-    public function testAddsProperHeaders()
-    {
+    /**
+     * @dataProvider generateTestCases
+     *
+     * @param CommandInterface $command
+     * @param array $expectedHeaders
+     * @param array $expectedNonHeaders
+     */
+    public function testAddsProperHeaders(
+        CommandInterface $command,
+        array $expectedHeaders,
+        array $expectedNonHeaders
+    ) {
         $service = $this->generateTestService();
-        $client = $this->generateTestClient($service);
-        $command = $client->getCommand(
-            'StreamingOp',
-            [
-                'InputStream' => Psr7\stream_for('test'),
-//                'InputString' => 'some_string'
-            ]
+        $serializer = ClientResolver::_default_serializer([
+            'api' => $service,
+            'endpoint' => ''
+        ]);
+
+        $list = new HandlerList();
+        $list->prependBuild(Middleware::requestBuilder($serializer), 'builder');
+        $list->prependSign(
+            StreamRequestPayloadMiddleware::wrap($service),
+            'StreamRequestPayloadMiddleware'
         );
 
-        $list = $command->getHandlerList();
-//        $list = new HandlerList();
-        $list->setHandler(function ($command, $request) {
-//            var_dump($request);
+        $list->setHandler(function (
+            CommandInterface $command,
+            RequestInterface $request
+        ) use (
+            $expectedHeaders,
+            $expectedNonHeaders
+        ) {
+            $this->assertArraySubset($expectedHeaders, $request->getHeaders());
+            foreach ($expectedNonHeaders as $header) {
+                $this->assertArrayNotHasKey($header, $request->getHeaders());
+            }
             return new Result([]);
         });
-//        $list->appendSign(StreamRequestPayloadMiddleware::wrap($service));
 
         $handler = $list->resolve();
+        $handler($command, new Request('POST', 'https://foo.com'));
+    }
+
+    public function generateTestCases()
+    {
+        $service = $this->generateTestService();
+        $client = $this->generateTestClient($service);
+        $inputStream = Psr7\stream_for('test');
 
-        $result = $handler($command, new Request('POST', 'https://foo.com'));
+        return [
+            [
+                $client->getCommand(
+                    'NonStreamingOp',
+                    [
+                        'InputString' => 'teststring',
+                    ]
+                ),
+                [],
+                [ 'transfer-encoding', 'content-length' ],
+            ],
+            [
+                $client->getCommand(
+                    'StreamingOp',
+                    [
+                        'InputStream' => $inputStream,
+                    ]
+                ),
+                [ 'content-length' => [26] ],
+                [ 'transfer-encoding' ],
+            ],
+            [
+                $client->getCommand(
+                    'StreamingLengthOp',
+                    [
+                        'InputStream' => $inputStream,
+                    ]
+                ),
+                [ 'content-length' => [26] ],
+                [ 'transfer-encoding' ],
+            ],
+            [
+                $client->getCommand(
+                    'StreamingUnsignedOp',
+                    [
+                        'InputStream' => $inputStream,
+                    ]
+                ),
+                [ 'transfer-encoding' => ['chunked'] ],
+                [ 'content-length' ],
+            ],
+            [
+                $client->getCommand(
+                    'StreamingLengthUnsignedOp',
+                    [
+                        'InputStream' => $inputStream,
+                    ]
+                ),
+                [ 'content-length' => [26] ],
+                [ 'transfer-encoding' ],
+            ],
+        ];
     }
 
     private function generateTestClient(Service $service, $args = [])
@@ -76,6 +158,17 @@ private function generateTestService()
                         "type" => "blob",
                         "streaming" => true,
                     ],
+                    "NonStreamingInputShape" => [
+                        "type" => "structure",
+                        "required" => [
+                            "InputString",
+                        ],
+                        "members" => [
+                            "InputString" => [
+                                "shape" => "StringType",
+                            ],
+                        ],
+                    ],
                     "StreamingInputShape" => [
                         "type" => "structure",
                         "required" => [
@@ -103,6 +196,15 @@ private function generateTestService()
                     ],
                 ],
                 'operations' => [
+                    "NonStreamingOp" => [
+                        "name"=> "NonStreamingOp",
+                        "http"=> [
+                            "method"=> "POST",
+                            "requestUri"=> "/",
+                            "responseCode"=> 200
+                        ],
+                        "input"=> ["shape"=> "NonStreamingInputShape"],
+                    ],
                     "StreamingOp" => [
                         "name"=> "StreamingOp",
                         "http"=> [