com.linkedin.restli.client.testutils.MockBatchKVResponseFactory Maven / Gradle / Ivy

Go to download
/*
   Copyright (c) 2014 LinkedIn Corp.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package com.linkedin.restli.client.testutils;


import com.linkedin.data.DataMap;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.restli.client.response.BatchKVResponse;
import com.linkedin.restli.common.BatchResponse;
import com.linkedin.restli.common.ComplexResourceKey;
import com.linkedin.restli.common.CompoundKey;
import com.linkedin.restli.common.ErrorResponse;
import java.util.Collections;
import java.util.Map;


/**
 * Factory for creating a {@link BatchKVResponse} that can be used for tests.
 *
 * @author jflorencio
 * @author kparikh
 */
public class MockBatchKVResponseFactory
{
  private MockBatchKVResponseFactory() { }

  /**
   * Create a {@link BatchKVResponse} where the key is a class that extends {@link CompoundKey}
   *
   * @param keyClass the class of the key
   * @param keyParts the pieces that make up the {@link CompoundKey}
   * @param valueClass the value class
   * @param recordTemplates the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getResults()}
   * @param errorResponses the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getErrors()}
   * @param 
   * @param 
   * @return
   */
  public static  BatchKVResponse createWithCompoundKey
      (Class keyClass,
       Map keyParts,
       Class valueClass,
       Map recordTemplates,
       Map errorResponses)
  {
    return create(keyClass, valueClass, keyParts, recordTemplates, errorResponses);
  }

  /**
   * Create a {@link BatchKVResponse} where the key is a {@link ComplexResourceKey}
   *
   * @param valueClass the value class
   * @param keyKeyClass the class of the key part of the {@link ComplexResourceKey}
   * @param keyParamsClass the class of the params part of the {@link ComplexResourceKey}
   * @param recordTemplates the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getResults()}
   *                        NOTE: the params part of the {@link ComplexResourceKey} is removed in this map. A new
   *                        instance of the params class is created with no data in it.
   * @param errorResponses the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getErrors()}
   *                       NOTE: the params part of the {@link ComplexResourceKey} is removed in this map. A new
   *                       instance of the params class is created with no data in it.
   * @param 
   * @return
   */
  @SuppressWarnings("rawtypes")
  public static  BatchKVResponse createWithComplexKey
      (Class valueClass,
       Class keyKeyClass,
       Class keyParamsClass,
       Map, V> recordTemplates,
       Map, ErrorResponse> errorResponses)
  {
    DataMap batchResponseDataMap = buildDataMap(recordTemplates, errorResponses);

    return new BatchKVResponse(batchResponseDataMap,
                                                      ComplexResourceKey.class,
                                                      valueClass,
                                                      null,
                                                      keyKeyClass,
                                                      keyParamsClass);
  }

  /**
   * Creates a {@link BatchKVResponse} where the key is a primitive, or a typeref to a primitive.
   *
   * @param keyClass class for the key
   * @param valueClass class for the value
   * @param recordTemplates the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getResults()}
   * @param errorResponses the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getErrors()}
   * @param 
   * @param 
   * @return
   */
  public static  BatchKVResponse createWithPrimitiveKey(Class keyClass,
                                                                                           Class valueClass,
                                                                                           Map recordTemplates,
                                                                                           Map errorResponses)
  {
    return create(keyClass, valueClass, null, recordTemplates, errorResponses);
  }

  /**
   * Creates a {@link BatchKVResponse} where the key is a typeref to a custom Java class.
   *
   * @param keyClass the custom Java class
   * @param typerefClass the typeref class (the generated class that extends {@link RecordTemplate})
   * @param valueClass class for the value
   * @param recordTemplates the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getResults()}
   * @param errorResponses the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getErrors()}
   * @param 
   * @param 
   * @param 
   * @return
   */
  @SuppressWarnings({"unchecked", "rawtypes"})
  public static  BatchKVResponse createWithCustomTyperefKey(Class keyClass,
                                                                                                   Class typerefClass,
                                                                                                   Class valueClass,
                                                                                                   Map recordTemplates,
                                                                                                   Map errorResponses)
  {
    DataMap batchResponseDataMap = buildDataMap(recordTemplates, errorResponses);

    return new BatchKVResponse(batchResponseDataMap,
                               typerefClass,
                               valueClass,
                               Collections.emptyMap());
  }

  private static  DataMap buildDataMap(Map recordTemplates,
                                                                    Map errorResponses)
  {
    DataMap batchResponseDataMap = new DataMap();
    DataMap rawBatchData = new DataMap();
    for (Map.Entry entry : recordTemplates.entrySet())
    {
      rawBatchData.put(String.valueOf(entry.getKey()), entry.getValue().data());
    }
    batchResponseDataMap.put(BatchResponse.RESULTS, rawBatchData);

    DataMap rawErrorData = new DataMap();
    for (Map.Entry errorResponse : errorResponses.entrySet())
    {
      rawErrorData.put(String.valueOf(errorResponse.getKey()), errorResponse.getValue().data());
    }
    batchResponseDataMap.put(BatchResponse.ERRORS, rawErrorData);
    return batchResponseDataMap;
  }

  private static  BatchKVResponse create
      (Class keyClass,
       Class valueClass,
       Map keyParts,
       Map recordTemplates,
       Map errorResponses)
  {
    DataMap batchResponseDataMap = buildDataMap(recordTemplates, errorResponses);

    return new BatchKVResponse(batchResponseDataMap,
                                     keyClass,
                                     valueClass,
                                     keyParts,
                                     null,
                                     null);
  }
}