Add the ability to mock document model operations in DynamoDB

sdk/src/Services/DynamoDBv2/Custom/DataModel/BatchWrite.cs

@@ -128,7 +128,10 @@ internal BatchWrite(DynamoDBContext context, Type valuesType, DynamoDBFlatConfig
             }
 
             Table table = _context.GetTargetTable(_storageConfig, _config);
-            DocumentBatch = table.CreateBatchWrite();
+
+            // Table.CreateBatchWrite() returns the IDocumentBatchWrite interface.
+            // But since we rely on the internal behavior of DocumentBatchWrite, we instantiate it via the constructor.
+            DocumentBatch = new DocumentBatchWrite(table);
         }
 
         /// <inheritdoc/>

sdk/src/Services/DynamoDBv2/Custom/DataModel/ContextInternal.cs

@@ -189,7 +189,7 @@ internal Table GetUnconfiguredTable(string tableName, bool disableFetchingTableM
 
                 var emptyConfig = new TableConfig(tableName, conversion: null, consumer: Table.DynamoDBConsumer.DataModel,
                     storeAsEpoch: null, isEmptyStringValueEnabled: false, metadataCachingMode: Config.MetadataCachingMode);
-                table = Table.LoadTable(Client, emptyConfig);
+                table = Table.LoadTable(Client, emptyConfig) as Table;
                 tablesMap[tableName] = table;
 
                 return table;
@@ -1178,7 +1178,9 @@ private ContextSearch ConvertScan<T>(IEnumerable<ScanCondition> conditions, Dyna
                 IndexName = flatConfig.IndexName,
                 ConsistentRead = flatConfig.ConsistentRead.GetValueOrDefault(false)
             };
-            Search scan = table.Scan(scanConfig);
+
+            // table.Scan() returns the ISearch interface but we explicitly cast it to a Search object since we rely on its internal behavior
+            Search scan = table.Scan(scanConfig) as Search;
             return new ContextSearch(scan, flatConfig);
         }
 
@@ -1187,7 +1189,9 @@ private ContextSearch ConvertFromScan<T>(ScanOperationConfig scanConfig, DynamoD
             DynamoDBFlatConfig flatConfig = new DynamoDBFlatConfig(operationConfig, Config);
             ItemStorageConfig storageConfig = StorageConfigCache.GetConfig<T>(flatConfig);
             Table table = GetTargetTable(storageConfig, flatConfig);
-            Search search = table.Scan(scanConfig);
+
+            // table.Scan() returns the ISearch interface but we explicitly cast it to a Search object since we rely on its internal behavior
+            Search search = table.Scan(scanConfig) as Search;
             return new ContextSearch(search, flatConfig);
         }
 
@@ -1196,7 +1200,9 @@ private ContextSearch ConvertFromQuery<T>(QueryOperationConfig queryConfig, Dyna
             DynamoDBFlatConfig flatConfig = new DynamoDBFlatConfig(operationConfig, Config);
             ItemStorageConfig storageConfig = StorageConfigCache.GetConfig<T>(flatConfig);
             Table table = GetTargetTable(storageConfig, flatConfig);
-            Search search = table.Query(queryConfig);
+
+            // table.Query() returns the ISearch interface but we explicitly cast it to a Search object since we rely on its internal behavior
+            Search search = table.Query(queryConfig) as Search;
             return new ContextSearch(search, flatConfig);
         }
 
@@ -1241,7 +1247,9 @@ private ContextSearch ConvertQueryHelper<T>(DynamoDBFlatConfig currentConfig, It
             {
                 queryConfig.Select = SelectValues.AllProjectedAttributes;
             }
-            Search query = table.Query(queryConfig);
+
+            // table.Query() returns the ISearch interface but we explicitly cast it to a Search object since we rely on its internal behavior
+            Search query = table.Query(queryConfig) as Search;
 
             return new ContextSearch(query, currentConfig);
         }

sdk/src/Services/DynamoDBv2/Custom/DataModel/TransactGet.cs

@@ -157,7 +157,10 @@ internal TransactGet(DynamoDBContext context, DynamoDBFlatConfig config)
             _config = config;
             _storageConfig = context.StorageConfigCache.GetConfig<T>(config);
             var table = context.GetTargetTable(_storageConfig, config);
-            DocumentTransaction = table.CreateTransactGet();
+
+            // Table.CreateTransactGet() returns the IDocumentTransactGet interface.
+            // But since we rely on the internal behavior of DocumentTransactGet, we instantiate it via the constructor.
+            DocumentTransaction = new DocumentTransactGet(table);
         }
 
         private void ExecuteHelper()

sdk/src/Services/DynamoDBv2/Custom/DataModel/TransactWrite.cs

@@ -188,7 +188,10 @@ internal TransactWrite(DynamoDBContext context, DynamoDBFlatConfig config)
             _config = config;
             _storageConfig = context.StorageConfigCache.GetConfig<T>(config);
             Table table = _context.GetTargetTable(_storageConfig, _config);
-            DocumentTransaction = table.CreateTransactWrite();
+
+            // table.CreateTransactWrite() return the IDocumentTransactWrite interface.
+            // But since we rely on the internal behavior of DocumentTransactWrite, we instatiate it via the constructor.
+            DocumentTransaction = new DocumentTransactWrite(table);
         }
 
         /// <inheritdoc/>

sdk/src/Services/DynamoDBv2/Custom/DataModel/_bcl/Context.Sync.cs

@@ -348,20 +348,20 @@ public IEnumerable<T> FromQuery<T>(QueryOperationConfig queryConfig, FromQueryCo
         #region Table methods
 
         /// <inheritdoc/>
-        public Table GetTargetTable<T>()
+        public ITable GetTargetTable<T>()
         {
             return GetTargetTable<T>((GetTargetTableConfig)null);
         }
 
         /// <inheritdoc/>
         [Obsolete("Use the GetTargetTable overload that takes GetTargetTableConfig instead, since DynamoDBOperationConfig contains properties that are not applicable to GetTargetTable.")]
-        public Table GetTargetTable<T>(DynamoDBOperationConfig operationConfig = null)
+        public ITable GetTargetTable<T>(DynamoDBOperationConfig operationConfig = null)
         {
             return GetTargetTableInternal<T>(new DynamoDBFlatConfig(operationConfig, Config));
         }
 
         /// <inheritdoc/>
-        public Table GetTargetTable<T>(GetTargetTableConfig getTargetTableConfig)
+        public ITable GetTargetTable<T>(GetTargetTableConfig getTargetTableConfig)
         {
             return GetTargetTableInternal<T>(new DynamoDBFlatConfig(getTargetTableConfig?.ToDynamoDBOperationConfig(), Config));
         }

sdk/src/Services/DynamoDBv2/Custom/DataModel/_bcl/IDynamoDBContext.Sync.cs

@@ -593,7 +593,7 @@ partial interface IDynamoDBContext
         /// </summary>
         /// <typeparam name="T">Type to retrieve table for</typeparam>
         /// <returns>Table object</returns>
-        Table GetTargetTable<T>();
+        ITable GetTargetTable<T>();
 
         /// <summary>
         /// Retrieves the target table for the specified type
@@ -602,15 +602,15 @@ partial interface IDynamoDBContext
         /// <param name="operationConfig">Config object which can be used to override that table used.</param>
         /// <returns>Table object</returns>
         [Obsolete("Use the GetTargetTable overload that takes GetTargetTableConfig instead, since DynamoDBOperationConfig contains properties that are not applicable to GetTargetTable.")]
-        Table GetTargetTable<T>(DynamoDBOperationConfig operationConfig = null);
+        ITable GetTargetTable<T>(DynamoDBOperationConfig operationConfig = null);
 
         /// <summary>
         /// Retrieves the target table for the specified type
         /// </summary>
         /// <typeparam name="T">Type to retrieve table for</typeparam>
         /// <param name="getTargetTableConfig">Config object that can be used to override properties on the table's context for this request.</param>
         /// <returns>Table object</returns>
-        Table GetTargetTable<T>(GetTargetTableConfig getTargetTableConfig);
+        ITable GetTargetTable<T>(GetTargetTableConfig getTargetTableConfig);
 
         #endregion
     }

sdk/src/Services/DynamoDBv2/Custom/DataModel/_netstandard/Context.Sync.cs

@@ -14,11 +14,8 @@
  */
 
 using System;
-using System.Collections.Generic;
-using System.Linq;
 
 using Amazon.DynamoDBv2.DocumentModel;
-using Amazon.DynamoDBv2.Model;
 
 namespace Amazon.DynamoDBv2.DataModel
 {
@@ -27,20 +24,20 @@ public partial class DynamoDBContext : IDynamoDBContext
         #region Table methods
 
         /// <inheritdoc/>
-        public Table GetTargetTable<T>()
+        public ITable GetTargetTable<T>()
         {
             return GetTargetTableInternal<T>(new DynamoDBFlatConfig(null, Config));
         }
 
         /// <inheritdoc/>
         [Obsolete("Use the GetTargetTable overload that takes GetTargetTableConfig instead, since DynamoDBOperationConfig contains properties that are not applicable to GetTargetTable.")]
-        public Table GetTargetTable<T>(DynamoDBOperationConfig operationConfig = null)
+        public ITable GetTargetTable<T>(DynamoDBOperationConfig operationConfig = null)
         {
             return GetTargetTableInternal<T>(new DynamoDBFlatConfig(operationConfig, Config));
         }
 
         /// <inheritdoc/>
-        public Table GetTargetTable<T>(GetTargetTableConfig getTargetTableConfig)
+        public ITable GetTargetTable<T>(GetTargetTableConfig getTargetTableConfig)
         {
             return GetTargetTableInternal<T>(new DynamoDBFlatConfig(getTargetTableConfig?.ToDynamoDBOperationConfig(), Config));
         }

sdk/src/Services/DynamoDBv2/Custom/DataModel/_netstandard/IDynamoDBContext.Sync.cs

@@ -14,10 +14,8 @@
  */
 
 using System;
-using System.Collections.Generic;
 
 using Amazon.DynamoDBv2.DocumentModel;
-using Amazon.DynamoDBv2.Model;
 
 namespace Amazon.DynamoDBv2.DataModel
 {
@@ -34,22 +32,22 @@ partial interface IDynamoDBContext
         /// </summary>
         /// <typeparam name="T">Type to retrieve table for</typeparam>
         /// <returns>Table object</returns>
-        Table GetTargetTable<T>();
+        ITable GetTargetTable<T>();
 
         /// <summary>
         /// Retrieves the target table for the specified type
         /// </summary>
         /// <typeparam name="T">Type to retrieve table for</typeparam>
         /// <returns>Table object</returns>
         [Obsolete("Use the GetTargetTable overload that takes GetTargetTableConfig instead, since DynamoDBOperationConfig contains properties that are not applicable to GetTargetTable.")]
-        Table GetTargetTable<T>(DynamoDBOperationConfig operationConfig = null);
+        ITable GetTargetTable<T>(DynamoDBOperationConfig operationConfig = null);
 
         /// <summary>
         /// Retrieves the target table for the specified type
         /// </summary>
         /// <typeparam name="T">Type to retrieve table for</typeparam>
         /// <returns>Table object</returns>
-        Table GetTargetTable<T>(GetTargetTableConfig getTargetTableConfig);
+        ITable GetTargetTable<T>(GetTargetTableConfig getTargetTableConfig);
 
         #endregion
     }