Merge pull request #5 from cmbrose/more-unit-tests

More unit tests

Author: cmbrose

CHANGELOG.md

@@ -1,6 +1,9 @@
 # Release Notes
 All notable changes and release history of the "cosmos-db" module will be documented in this file.
 
+## 1.7
+* Minor interface improvements
+
 ## 1.6
 * Adds `Use-CosmosDbInternalFlag` command to access some experimental features and debugging helpers
 * Fixes support for extra query features in `Search-CosmosDbRecords` for DBs with more than one partition key range

cosmos-db/cosmos-db.psd1

@@ -3,7 +3,7 @@
 #
 # Generated by: cmbrose
 #
-# Generated on: 12/18/2020
+# Generated on: 12/27/2020
 #
 
 @{
@@ -12,7 +12,7 @@
 # RootModule = ''
 
 # Version number of this module.
-ModuleVersion = '1.6'
+ModuleVersion = '1.7'
 
 # Supported PSEditions
 # CompatiblePSEditions = @()

cosmos-db/cosmos-db.psm1

@@ -127,7 +127,7 @@ Function Get-AuthorizationHeader([string]$ResourceGroup, [string]$SubscriptionId
 Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$contentType="application/json", [bool]$isQuery=$false, [string]$PartitionKey=$null)
 {
     $headers = @{ 
-        "x-ms-date"=$now;
+        "x-ms-date" = $now;
         "x-ms-version" = $API_VERSION;
         "Authorization" = $encodedAuthString;
         "Cache-Control" = "No-Cache";
@@ -136,12 +136,12 @@ Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$co
 
     if ($isQuery)
     {
-        $headers[ "x-ms-documentdb-isquery"] = "true"
+        $headers["x-ms-documentdb-isquery"] = "true"
     }
 
     if ($PartitionKey)
     {
-        $headers["x-ms-documentdb-partitionkey"] = "[`"$requestPartitionKey`"]"
+        $headers["x-ms-documentdb-partitionkey"] = "[`"$PartitionKey`"]"
     }
 
     $headers
@@ -204,7 +204,7 @@ Function Invoke-CosmosDbApiRequestWithContinuation([string]$verb, [string]$url,
     }
 }
 
-Function Get-PartitionKeyRanges
+Function Get-PartitionKeyRangesOrError
 (
     [parameter(Mandatory=$true)][string]$ResourceGroup,
     [parameter(Mandatory=$true)][string]$Database, 
@@ -213,33 +213,40 @@ Function Get-PartitionKeyRanges
     [parameter(Mandatory=$true)][string]$SubscriptionId
 )
 {
-    $baseUrl = Get-BaseDatabaseUrl $Database
-    $collectionsUrl=Get-CollectionsUrl $Container $Collection
-    $pkRangeUrl="$collectionsUrl/$PARTITIONKEYRANGE_TYPE"
+    try
+    {
+        $baseUrl = Get-BaseDatabaseUrl $Database
+        $collectionsUrl=Get-CollectionsUrl $Container $Collection
+        $pkRangeUrl="$collectionsUrl/$PARTITIONKEYRANGE_TYPE"
 
-    $url = "$baseUrl/$pkRangeUrl"
+        $url = "$baseUrl/$pkRangeUrl"
 
-    $cacheKey = $url
-    $cacheResult = Get-CacheValue -Key $cacheKey -Cache $PARTITION_KEY_RANGE_CACHE
-    if ($cacheResult)
-    {
-        return $cacheResult
-    }
+        $cacheKey = $url
+        $cacheResult = Get-CacheValue -Key $cacheKey -Cache $PARTITION_KEY_RANGE_CACHE
+        if ($cacheResult)
+        {
+            return $cacheResult
+        }
 
-    $now = Get-Time
+        $now = Get-Time
 
-    $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $PARTITIONKEYRANGE_TYPE -resourceUrl $collectionsUrl -now $now
+        $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $PARTITIONKEYRANGE_TYPE -resourceUrl $collectionsUrl -now $now
 
-    $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey
-    $headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
+        $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey
+        $headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
 
-    $response = Invoke-CosmosDbApiRequest -Verb $GET_VERB -Url $url -Headers $headers | Get-CosmosDbRecordContent
+        $response = Invoke-CosmosDbApiRequest -Verb $GET_VERB -Url $url -Headers $headers | Get-CosmosDbRecordContent
 
-    $ranges = $response.partitionKeyRanges
+        $ranges = $response.partitionKeyRanges
 
-    Set-CacheValue -Key $cacheKey -Value $ranges -Cache $PARTITION_KEY_RANGE_CACHE -ExpirationHours 6
+        Set-CacheValue -Key $cacheKey -Value $ranges -Cache $PARTITION_KEY_RANGE_CACHE -ExpirationHours 6
 
-    $ranges
+        $ranges
+    } 
+    catch [System.Net.WebException]
+    {
+        $_.Exception
+    } 
 }
 
 Function Get-FilteredPartitionKeyRangesForQuery($allRanges, $queryRanges)
@@ -537,10 +544,21 @@ Function Search-CosmosDbRecords(
     }
 }
 
-Function Search-CosmosDbRecordsWithExtraFeatures([string]$ResourceGroup, [string]$Database, [string]$Container, [string]$Collection, [string]$Query, $Parameters, [string]$SubscriptionId)
+Function Search-CosmosDbRecordsWithExtraFeatures
+(
+    [string]$ResourceGroup,
+    [string]$Database, 
+    [string]$Container, 
+    [string]$Collection, 
+    [string]$Query, 
+    $Parameters, 
+    [string]$SubscriptionId
+)
 {
     begin
     {
+        $Parameters = @(Get-QueryParametersAsNameValuePairs $Parameters)
+
         $baseUrl=Get-BaseDatabaseUrl $Database
         $collectionsUrl=Get-CollectionsUrl $Container $Collection
         $docsUrl="$collectionsUrl/$DOCS_TYPE"
@@ -551,12 +569,17 @@ Function Search-CosmosDbRecordsWithExtraFeatures([string]$ResourceGroup, [string
 
         $encodedAuthString=Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
 
-        $allPartitionKeyRanges = Get-PartitionKeyRanges -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -SubscriptionId $SubscriptionId
+        $allPartitionKeyRangesOrError = Get-PartitionKeyRangesOrError -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -SubscriptionId $SubscriptionId
     }
     process
     {
         try 
         {
+            if ($allPartitionKeyRangesOrError -is [System.Net.WebException])
+            {
+                throw $allPartitionKeyRangesOrError
+            }
+
             $body = @{
                 query = $Query;
                 parameters = $Parameters;
@@ -571,32 +594,32 @@ Function Search-CosmosDbRecordsWithExtraFeatures([string]$ResourceGroup, [string
                 "x-ms-cosmos-is-query-plan-request" = "True";
             }
 
-            $response = Invoke-CosmosDbApiRequest -verb $POST_VERB -url $url -Body $body -Headers $headers | Get-CosmosDbRecordContent
+            $queryPlan = Invoke-CosmosDbApiRequest -verb $POST_VERB -url $url -Body $body -Headers $headers | Get-CosmosDbRecordContent
 
-            $rewrittenQuery = $response.QueryInfo.RewrittenQuery
+            $rewrittenQuery = $queryPlan.QueryInfo.RewrittenQuery
             $searchQuery = if ($rewrittenQuery) { $rewrittenQuery } else { $Query };
 
+            $body = @{
+                query = $searchQuery;
+                parameters = $Parameters;
+            }
+
             $headers.Remove("x-ms-cosmos-is-query-plan-request")
 
             $partitionKeyRanges = 
                 if ($env:COSMOS_DB_FLAG_ENABLE_PARTITION_KEY_RANGE_SEARCHES -eq 1)
                 {
-                    Get-FilteredPartitionKeyRangesForQuery -AllRanges $partitionKeyRanges -QueryRanges $response.QueryRanges
+                    Get-FilteredPartitionKeyRangesForQuery -AllRanges $allPartitionKeyRangesOrError -QueryRanges $queryPlan.QueryRanges
                 }
                 else 
                 {
-                    $allPartitionKeyRanges
+                    $allPartitionKeyRangesOrError
                 }
 
             foreach ($partitionKeyRange in $partitionKeyRanges)
             {
                 $headers["x-ms-documentdb-partitionkeyrangeid"] = $partitionKeyRange.id
 
-                $body = @{
-                    query = $searchQuery;
-                    parameters = $Parameters;
-                }
-
                 Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
             }
         }
@@ -888,11 +911,6 @@ Function Get-CosmosDbRecordContent([parameter(ValueFromPipeline)]$RecordResponse
     }
 }
 
-Function Use-CosmosDbFiddlerDebugging()
-{
-    $env:AZURE_CLI_DISABLE_CONNECTION_VERIFICATION=1
-}
-
 Function Use-CosmosDbInternalFlag
 (
     $enableFiddlerDebugging=$null,

tests/Get-AllCosmosDbRecords.Tests.ps1

@@ -0,0 +1,132 @@
+Get-Module cosmos-db | Remove-Module -Force
+Import-Module $PSScriptRoot\..\cosmos-db\cosmos-db.psm1 -Force
+
+InModuleScope cosmos-db {
+    Describe "Get-AllCosmosDbRecords" {                    
+        BeforeAll {
+            Use-CosmosDbInternalFlag -EnableCaching $false
+
+            . $PSScriptRoot\Utils.ps1    
+
+            $global:capturedNow = $null
+
+            $MOCK_SUB = "MOCK_SUB"
+            $MOCK_RG = "MOCK_RG"
+            $MOCK_DB = "MOCK_DB"
+            $MOCK_CONTAINER = "MOCK_CONTAINER"
+            $MOCK_COLLECTION = "MOCK_COLLECTION"
+
+            $MOCK_AUTH_HEADER = "MockAuthHeader"
+
+            Function VerifyGetAuthHeader($ResourceGroup, $SubscriptionId, $Database, $verb, $resourceType, $resourceUrl, $now)
+            {
+                $ResourceGroup | Should -Be $MOCK_RG
+                $SubscriptionId | Should -Be $MOCK_SUB
+
+                $verb | Should -Be "get"
+                $resourceType | Should -Be "docs"
+                $resourceUrl | Should -Be "dbs/$MOCK_CONTAINER/colls/$MOCK_COLLECTION"
+            }
+
+            Function VerifyInvokeCosmosDbApiRequest($verb, $url, $body, $headers, $partitionKey=$MOCK_RECORD_ID)
+            {
+                $verb | Should -Be "get"
+                $url | Should -Be "https://$MOCK_DB.documents.azure.com/dbs/$MOCK_CONTAINER/colls/$MOCK_COLLECTION/docs"        
+                $body | Should -Be $null
+                    
+                $global:capturedNow | Should -Not -Be $null
+
+                $expectedHeaders = Get-CommonHeaders -now $global:capturedNow -encodedAuthString $MOCK_AUTH_HEADER -isQuery $true
+            
+                AssertHashtablesEqual $expectedHeaders $headers
+            }
+
+            Mock Get-AuthorizationHeader {
+                param($ResourceGroup, $SubscriptionId, $Database, $verb, $resourceType, $resourceUrl, $now)
+        
+                VerifyGetAuthHeader $ResourceGroup $SubscriptionId $Database $verb $resourceType $resourceUrl $now | Out-Null
+        
+                $global:capturedNow = $now
+        
+                $MOCK_AUTH_HEADER
+            }
+        }
+
+        It "Sends correct request" {    
+            $response = @{
+                StatusCode = 200;
+                Content = "{}";
+                Headers = @{};
+            }
+
+            Mock Invoke-CosmosDbApiRequestWithContinuation {
+                param($verb, $url, $body, $headers) 
+                
+                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
+        
+                $response
+            }
+
+            $result = Get-AllCosmosDbRecords -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Container $MOCK_CONTAINER -Collection $MOCK_COLLECTION
+
+            $result | Should -BeExactly $response
+
+            Assert-MockCalled Invoke-CosmosDbApiRequestWithContinuation -Times 1
+        }
+
+        It "Returns multiple responses" {
+            $response1 = @{
+                StatusCode = 200;
+                Content = "1";
+                Headers = @{};
+            }
+
+            $response2 = @{
+                StatusCode = 200;
+                Content = "1";
+                Headers = @{};
+            }
+
+            $response3 = @{
+                StatusCode = 200;
+                Content = "1";
+                Headers = @{};
+            }
+
+            Mock Invoke-CosmosDbApiRequestWithContinuation {
+                param($verb, $url, $body, $headers) 
+                
+                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
+        
+                $response1
+                $response2
+                $response3
+            }
+
+            $result = Get-AllCosmosDbRecords -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Container $MOCK_CONTAINER -Collection $MOCK_COLLECTION
+
+            $result.Count | Should -Be 3
+            $result[0] | Should -BeExactly $response1
+            $result[1] | Should -BeExactly $response2
+            $result[2] | Should -BeExactly $response3
+
+            Assert-MockCalled Invoke-CosmosDbApiRequestWithContinuation -Times 1
+        }
+
+        It "Should handle exceptions gracefully" {    
+            $response = [System.Net.HttpWebResponse]@{}
+
+            Mock Invoke-CosmosDbApiRequest {
+                param($verb, $url, $body, $headers) 
+                
+                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
+        
+                throw [System.Net.WebException]::new("", $null, [System.Net.WebExceptionStatus]::UnknownError, $response)
+            }
+
+            $result = Get-AllCosmosDbRecords -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Container $MOCK_CONTAINER -Collection $MOCK_COLLECTION
+
+            $result | Should -BeExactly $response
+        }
+    }
+}

tests/Get-AuthorizationHeader.Tests.ps1

@@ -0,0 +1,47 @@
+Get-Module cosmos-db | Remove-Module -Force
+Import-Module $PSScriptRoot\..\cosmos-db\cosmos-db.psm1 -Force
+
+InModuleScope cosmos-db {
+    Describe "Get-AuthorizationHeader" {                    
+        BeforeAll {
+            Use-CosmosDbInternalFlag -EnableCaching $false
+
+            $MOCK_SUB = "MOCK_SUB"
+            $MOCK_RG = "MOCK_RG"
+            $MOCK_DB = "MOCK_DB"
+            $MOCK_RESOURCE_URL = "MOCK_RESOURCE_URL"
+            $MOCK_RESOURCE_TYPE = "MOCK_RESOURCE_TYPE"
+            $MOCK_VERB = "MOCK_VERB"
+            $MOCK_NOW = "MOCK_NOW"
+
+            $MOCK_MASTER_KEY_BYTES = [System.Text.Encoding]::UTF8.GetBytes('gVkYp3s6v9y$B&E)H@MbQeThWmZq4t7w')
+
+            Mock Get-Base64Masterkey {
+                param($ResourceGroup, $Database, $SubscriptionId)
+        
+                $ResourceGroup | Should -Be $MOCK_RG | Out-Null
+                $Database | Should -Be $MOCK_DB | Out-Null
+                $SubscriptionId | Should -Be $MOCK_SUB | Out-Null
+                
+                [System.Convert]::ToBase64String($MOCK_MASTER_KEY_BYTES)
+            }
+        }
+
+        It "Returns the correct signature hashed with the master key" {    
+            $result = Get-AuthorizationHeader  -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Verb $MOCK_VERB -ResourceType $MOCK_RESOURCE_TYPE -ResourceUrl $MOCK_RESOURCE_URL -Now $MOCK_NOW
+
+            $expectedSignature = "$MOCK_VERB`n$MOCK_RESOURCE_TYPE`n$MOCK_RESOURCE_URL`n$MOCK_NOW`n`n".ToLower()
+
+            $hasher = New-Object System.Security.Cryptography.HMACSHA256 -Property @{ Key = $MOCK_MASTER_KEY_BYTES }
+            $sigBinary=[System.Text.Encoding]::UTF8.GetBytes($expectedSignature)
+            $hashBytes=$hasher.ComputeHash($sigBinary)
+            $expectedBase64Hash=[System.Convert]::ToBase64String($hashBytes)
+            
+            $expectedHeader = [uri]::EscapeDataString("type=master&ver=1.0&sig=$expectedBase64Hash")
+
+            $result | Should -Be $expectedHeader
+
+            Assert-MockCalled Get-Base64Masterkey -Times 1
+        }
+    }
+}