external-table/test/pxfprotocol_test.c - cloudberry-pxf - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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.
  */

 #include <stdarg.h>
 #include <stddef.h>
 #include <setjmp.h>
 #include "cmockery.h"

 #include "postgres.h"
 #include "utils/memutils.h"

 /* Define UNIT_TESTING so that the extension can skip declaring PG_MODULE_MAGIC */
 #define UNIT_TESTING

 /* include unit under test */
 #include "../src/pxfprotocol.c"

 /* include mock files */
 #include "mock/pxfbridge_mock.c"
 #include "mock/pxfuriparser_mock.c"
 #include "mock/pxffragment_mock.c"
 #include "mock/pxffilters_mock.c"

 char *uri_no_profile = "pxf://default/tmp/dummy1?FRAGMENTER=xxx&RESOLVER=yyy&ACCESSOR=zzz";
 char *uri_param = "pxf://localhost:5888/tmp/dummy1";

 static void
 test_pxfprotocol_validate_urls(void **state)
 {
 	/* setup call info with no call context */
 	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
 	fcinfo->context = palloc0(sizeof(ExtProtocolValidatorData));
 	fcinfo->context->type = T_ExtProtocolValidatorData;
 	Value	   *v = makeString(uri_no_profile);
 	List	   *list = list_make1(v);

 	((ExtProtocolValidatorData *) fcinfo->context)->url_list = list;

 	/* set mock behavior for uri parsing */
 	GPHDUri    *gphd_uri = palloc0(sizeof(GPHDUri));

 	expect_string(parseGPHDUri, uri_str, uri_no_profile);
 	will_return(parseGPHDUri, gphd_uri);

 	expect_value(GPHDUri_verify_no_duplicate_options, uri, gphd_uri);
 	will_be_called(GPHDUri_verify_no_duplicate_options);

 	expect_value(GPHDUri_opt_exists, uri, gphd_uri);
 	expect_string(GPHDUri_opt_exists, key, PXF_PROFILE);
 	will_return(GPHDUri_opt_exists, false);

 	expect_value(GPHDUri_verify_core_options_exist, uri, gphd_uri);
 	expect_any(GPHDUri_verify_core_options_exist, coreoptions);
 	will_be_called(GPHDUri_verify_core_options_exist);

 	expect_value(freeGPHDUri, uri, gphd_uri);
 	will_be_called(freeGPHDUri);

 	Datum		d = pxfprotocol_validate_urls(fcinfo);

 	assert_int_equal(DatumGetInt32(d), 0);

 	/* cleanup */
 	list_free_deep(list);
 	pfree(fcinfo->context);
 	pfree(fcinfo);
 }

 static void
 test_pxfprotocol_import_first_call(void **state)
 {
 	/* setup call info with no call context */
 	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
 	fcinfo->context = palloc0(sizeof(ExtProtocolData));
 	fcinfo->context->type = T_ExtProtocolData;
 	EXTPROTOCOL_GET_DATALEN(fcinfo) = 100;
 	EXTPROTOCOL_GET_DATABUF(fcinfo) = palloc0(EXTPROTOCOL_GET_DATALEN(fcinfo));
 	((ExtProtocolData *) fcinfo->context)->prot_last_call = false;
 	((ExtProtocolData *) fcinfo->context)->prot_url = uri_param;

 	ExternalSelectDesc desc = (ExternalSelectDesc) palloc0(sizeof(ExternalSelectDesc));
 	desc->projInfo = NULL;

 	((ExtProtocolData*) fcinfo->context)->desc = desc;

 	Relation	relation = (Relation) palloc0(sizeof(Relation));

 	((ExtProtocolData *) fcinfo->context)->prot_relation = relation;

 	/* set mock behavior for uri parsing */
 	GPHDUri    *gphd_uri = palloc0(sizeof(GPHDUri));

 	expect_string(parseGPHDUri, uri_str, uri_param);
 	will_return(parseGPHDUri, gphd_uri);

 	/* set mock behavior for set fragments */
 	gphd_uri->fragments = palloc0(sizeof(List));
 	expect_value(get_fragments, uri, gphd_uri);
 	expect_value(get_fragments, relation, relation);
 	will_assign_memory(get_fragments, uri, gphd_uri, sizeof(GPHDUri));
 	will_be_called(get_fragments);

 	/* set mock behavior for bridge import start -- nothing here */
 	expect_any(gpbridge_import_start, context);
 	will_be_called(gpbridge_import_start);

 	/* set mock behavior for bridge read */
 	const int	EXPECTED_SIZE = 31;

 	/* expected number of bytes read from the bridge */
 	expect_any(gpbridge_read, context);
 	expect_value(gpbridge_read, databuf, EXTPROTOCOL_GET_DATABUF(fcinfo));
 	expect_value(gpbridge_read, datalen, EXTPROTOCOL_GET_DATALEN(fcinfo));
 	will_return(gpbridge_read, EXPECTED_SIZE);

 	Datum		d = pxfprotocol_import(fcinfo);

 	/* return number of bytes read from the bridge */
 	assert_int_equal(DatumGetInt32(d), EXPECTED_SIZE);
 	gphadoop_context *context = (gphadoop_context *) EXTPROTOCOL_GET_USER_CTX(fcinfo);

 	/* context has been created */
 	assert_true(context != NULL);
 	/* gphduri has been parsed */
 	assert_true(context->gphd_uri != NULL);
 	/* uri has been initialized */
 	assert_true(context->uri.data != NULL);
 	/* no write file name for import case */
 	assert_true(context->relation != NULL);
 	/* relation pointer is copied */
 	assert_int_equal(context->relation, relation);

 	/* cleanup */
 	pfree(desc);
 	pfree(relation);
 	pfree(gphd_uri);
 	pfree(EXTPROTOCOL_GET_USER_CTX(fcinfo));
 	pfree(EXTPROTOCOL_GET_DATABUF(fcinfo));
 	pfree(fcinfo->context);
 	pfree(fcinfo);
 }

 static void
 test_pxfprotocol_import_second_call(void **state)
 {
 	/* setup call info with call context */
 	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
 	fcinfo->context = palloc0(sizeof(ExtProtocolData));
 	fcinfo->context->type = T_ExtProtocolData;
 	EXTPROTOCOL_GET_DATALEN(fcinfo) = 100;
 	EXTPROTOCOL_GET_DATABUF(fcinfo) = palloc0(EXTPROTOCOL_GET_DATALEN(fcinfo));
 	((ExtProtocolData *) fcinfo->context)->prot_last_call = false;
 	gphadoop_context *call_context = palloc0(sizeof(gphadoop_context));

 	EXTPROTOCOL_SET_USER_CTX(fcinfo, call_context);

 	/* set mock behavior for bridge read */
 	const int	EXPECTED_SIZE = 0;

 	/* expected number of bytes read from the bridge */
 	expect_value(gpbridge_read, context, call_context);
 	expect_value(gpbridge_read, databuf, EXTPROTOCOL_GET_DATABUF(fcinfo));
 	expect_value(gpbridge_read, datalen, EXTPROTOCOL_GET_DATALEN(fcinfo));
 	will_return(gpbridge_read, EXPECTED_SIZE);

 	Datum		d = pxfprotocol_import(fcinfo);

 	assert_int_equal(DatumGetInt32(d), EXPECTED_SIZE);
 	/* return number of bytes read from the bridge */
 	assert_true(EXTPROTOCOL_GET_USER_CTX(fcinfo) == call_context);
 	/* context is still the same */

 	/* cleanup */
 	pfree(call_context);
 	pfree(EXTPROTOCOL_GET_DATABUF(fcinfo));
 	pfree(fcinfo->context);
 	pfree(fcinfo);
 }

 static void
 test_pxfprotocol_import_last_call(void **state)
 {
 	/* setup call info with a call context and last call indicator */
 	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
 	fcinfo->context = palloc0(sizeof(ExtProtocolData));
 	fcinfo->context->type = T_ExtProtocolData;
 	gphadoop_context *call_context = palloc0(sizeof(gphadoop_context));

 	EXTPROTOCOL_SET_USER_CTX(fcinfo, call_context);
 	EXTPROTOCOL_SET_LAST_CALL(fcinfo);

 	/* init data in context that will be cleaned up */
 	initStringInfo(&call_context->uri);

 	/* set mock behavior for bridge cleanup */
 	expect_value(gpbridge_cleanup, context, call_context);
 	will_be_called(gpbridge_cleanup);

 	Datum		d = pxfprotocol_import(fcinfo);

 	assert_int_equal(DatumGetInt32(d), 0);
 	/* 0 is returned from function */
 	assert_true(EXTPROTOCOL_GET_USER_CTX(fcinfo) == NULL);
 	/* call context is cleaned up */

 	/* cleanup */
 	pfree(fcinfo->context);
 	pfree(fcinfo);
 }

 static void
 test_pxfprotocol_export_first_call(void **state)
 {
 	/* setup call info with no call context */
 	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
 	fcinfo->context = palloc0(sizeof(ExtProtocolData));
 	fcinfo->context->type = T_ExtProtocolData;
 	EXTPROTOCOL_GET_DATALEN(fcinfo) = 100;
 	EXTPROTOCOL_GET_DATABUF(fcinfo) = palloc0(EXTPROTOCOL_GET_DATALEN(fcinfo));
 	((ExtProtocolData *) fcinfo->context)->prot_last_call = false;
 	((ExtProtocolData *) fcinfo->context)->prot_url = uri_param;

 	ExternalSelectDesc desc = (ExternalSelectDesc) palloc0(sizeof(ExternalSelectDesc));
 	desc->projInfo = NULL;

 	((ExtProtocolData*) fcinfo->context)->desc = desc;

 	Relation	relation = (Relation) palloc0(sizeof(Relation));

 	((ExtProtocolData *) fcinfo->context)->prot_relation = relation;

 	/* set mock behavior for uri parsing */
 	GPHDUri		*gphd_uri = palloc0(sizeof(GPHDUri));

 	expect_string(parseGPHDUri, uri_str, uri_param);
 	will_return(parseGPHDUri, gphd_uri);

 	/* set mock behavior for bridge export start -- nothing here */
 	expect_any(gpbridge_export_start, context);
 	will_be_called(gpbridge_export_start);

 	/* set mock behavior for bridge write */
 	const int	EXPECTED_SIZE = 31;

 	/* expected number of bytes written to the bridge */
 	expect_any(gpbridge_write, context);
 	expect_value(gpbridge_write, databuf, EXTPROTOCOL_GET_DATABUF(fcinfo));
 	expect_value(gpbridge_write, datalen, EXTPROTOCOL_GET_DATALEN(fcinfo));
 	will_return(gpbridge_write, EXPECTED_SIZE);

 	Datum		d = pxfprotocol_export(fcinfo);

 	/* return number of bytes written to the bridge */
 	assert_int_equal(DatumGetInt32(d), EXPECTED_SIZE);
 	gphadoop_context *context = (gphadoop_context *) EXTPROTOCOL_GET_USER_CTX(fcinfo);

 	/* context has been created */
 	assert_true(context != NULL);
 	/* gphduri has been parsed */
 	assert_true(context->gphd_uri != NULL);
 	/* uri has been initialized */
 	assert_true(context->uri.data != NULL);
 	/* write file name initialized, but empty, since it is filled by another component */
 	assert_true(context->relation != NULL);
 	/* relation pointer is copied */
 	assert_int_equal(context->relation, relation);

 	/* cleanup */
 	pfree(desc);
 	pfree(relation);
 	pfree(gphd_uri);
 	pfree(EXTPROTOCOL_GET_USER_CTX(fcinfo));
 	pfree(EXTPROTOCOL_GET_DATABUF(fcinfo));
 	pfree(fcinfo->context);
 	pfree(fcinfo);
 }

 static void
 test_pxfprotocol_export_second_call(void **state)
 {
 	/* setup call info with call context */
 	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
 	fcinfo->context = palloc0(sizeof(ExtProtocolData));
 	fcinfo->context->type = T_ExtProtocolData;
 	EXTPROTOCOL_GET_DATALEN(fcinfo) = 100;
 	EXTPROTOCOL_GET_DATABUF(fcinfo) = palloc0(EXTPROTOCOL_GET_DATALEN(fcinfo));
 	((ExtProtocolData *) fcinfo->context)->prot_last_call = false;
 	gphadoop_context *call_context = palloc0(sizeof(gphadoop_context));

 	EXTPROTOCOL_SET_USER_CTX(fcinfo, call_context);

 	/* set mock behavior for bridge write */
 	const int	EXPECTED_SIZE = 0;

 	/* expected number of bytes written to the bridge */
 	expect_value(gpbridge_write, context, call_context);
 	expect_value(gpbridge_write, databuf, EXTPROTOCOL_GET_DATABUF(fcinfo));
 	expect_value(gpbridge_write, datalen, EXTPROTOCOL_GET_DATALEN(fcinfo));
 	will_return(gpbridge_write, EXPECTED_SIZE);

 	Datum		d = pxfprotocol_export(fcinfo);

 	assert_int_equal(DatumGetInt32(d), EXPECTED_SIZE);
 	/* return number of bytes written to the bridge */
 	assert_true(EXTPROTOCOL_GET_USER_CTX(fcinfo) == call_context);
 	/* context is still the same */

 	/* cleanup */
 	pfree(call_context);
 	pfree(EXTPROTOCOL_GET_DATABUF(fcinfo));
 	pfree(fcinfo->context);
 	pfree(fcinfo);
 }

 static void
 test_pxfprotocol_export_last_call(void **state)
 {
 	/* setup call info with a call context and last call indicator */
 	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
 	fcinfo->context = palloc0(sizeof(ExtProtocolData));
 	fcinfo->context->type = T_ExtProtocolData;
 	gphadoop_context *call_context = palloc0(sizeof(gphadoop_context));

 	EXTPROTOCOL_SET_USER_CTX(fcinfo, call_context);
 	EXTPROTOCOL_SET_LAST_CALL(fcinfo);

 	/* init data in context that will be cleaned up */
 	initStringInfo(&call_context->uri);

 	/* set mock behavior for bridge cleanup */
 	expect_value(gpbridge_cleanup, context, call_context);
 	will_be_called(gpbridge_cleanup);

 	Datum		d = pxfprotocol_export(fcinfo);

 	assert_int_equal(DatumGetInt32(d), 0);
 	/* 0 is returned from function */
 	assert_true(EXTPROTOCOL_GET_USER_CTX(fcinfo) == NULL);
 	/* call context is cleaned up */

 	/* cleanup */
 	pfree(fcinfo->context);
 	pfree(fcinfo);
 }
 /* test setup and teardown methods */
 static void
 before_test(void **state)
 {
 	/* set global variables */
 	GpIdentity.segindex = 0;
 }

 static void
 after_test(void **state)
 {
 	/*
 	 * no-op, but the teardown seems to be required when the test fails,
 	 * otherwise CMockery issues a mismatch error
 	 */
 }

 int
 main(int argc, char *argv[])
 {
 	cmockery_parse_arguments(argc, argv);

 	const		UnitTest tests[] = {
 		unit_test(test_pxfprotocol_validate_urls),
 		unit_test_setup_teardown(test_pxfprotocol_import_first_call, before_test, after_test),
 		unit_test_setup_teardown(test_pxfprotocol_import_second_call, before_test, after_test),
 		unit_test_setup_teardown(test_pxfprotocol_import_last_call, before_test, after_test),
 		unit_test_setup_teardown(test_pxfprotocol_export_first_call, before_test, after_test),
 		unit_test_setup_teardown(test_pxfprotocol_export_second_call, before_test, after_test),
 		unit_test_setup_teardown(test_pxfprotocol_export_last_call, before_test, after_test)

 	};

 	MemoryContextInit();

 	return run_tests(tests);
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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.
	*/

	#include <stdarg.h>
	#include <stddef.h>
	#include <setjmp.h>
	#include "cmockery.h"

	#include "postgres.h"
	#include "utils/memutils.h"

	/* Define UNIT_TESTING so that the extension can skip declaring PG_MODULE_MAGIC */
	#define UNIT_TESTING

	/* include unit under test */
	#include "../src/pxfprotocol.c"

	/* include mock files */
	#include "mock/pxfbridge_mock.c"
	#include "mock/pxfuriparser_mock.c"
	#include "mock/pxffragment_mock.c"
	#include "mock/pxffilters_mock.c"

	char *uri_no_profile = "pxf://default/tmp/dummy1?FRAGMENTER=xxx&RESOLVER=yyy&ACCESSOR=zzz";
	char *uri_param = "pxf://localhost:5888/tmp/dummy1";

	static void
	test_pxfprotocol_validate_urls(void **state)
	{
	/* setup call info with no call context */
	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
	fcinfo->context = palloc0(sizeof(ExtProtocolValidatorData));
	fcinfo->context->type = T_ExtProtocolValidatorData;
	Value *v = makeString(uri_no_profile);
	List *list = list_make1(v);

	((ExtProtocolValidatorData *) fcinfo->context)->url_list = list;

	/* set mock behavior for uri parsing */
	GPHDUri *gphd_uri = palloc0(sizeof(GPHDUri));

	expect_string(parseGPHDUri, uri_str, uri_no_profile);
	will_return(parseGPHDUri, gphd_uri);

	expect_value(GPHDUri_verify_no_duplicate_options, uri, gphd_uri);
	will_be_called(GPHDUri_verify_no_duplicate_options);

	expect_value(GPHDUri_opt_exists, uri, gphd_uri);
	expect_string(GPHDUri_opt_exists, key, PXF_PROFILE);
	will_return(GPHDUri_opt_exists, false);

	expect_value(GPHDUri_verify_core_options_exist, uri, gphd_uri);
	expect_any(GPHDUri_verify_core_options_exist, coreoptions);
	will_be_called(GPHDUri_verify_core_options_exist);

	expect_value(freeGPHDUri, uri, gphd_uri);
	will_be_called(freeGPHDUri);

	Datum d = pxfprotocol_validate_urls(fcinfo);

	assert_int_equal(DatumGetInt32(d), 0);

	/* cleanup */
	list_free_deep(list);
	pfree(fcinfo->context);
	pfree(fcinfo);
	}

	static void
	test_pxfprotocol_import_first_call(void **state)
	{
	/* setup call info with no call context */
	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
	fcinfo->context = palloc0(sizeof(ExtProtocolData));
	fcinfo->context->type = T_ExtProtocolData;
	EXTPROTOCOL_GET_DATALEN(fcinfo) = 100;
	EXTPROTOCOL_GET_DATABUF(fcinfo) = palloc0(EXTPROTOCOL_GET_DATALEN(fcinfo));
	((ExtProtocolData *) fcinfo->context)->prot_last_call = false;
	((ExtProtocolData *) fcinfo->context)->prot_url = uri_param;

	ExternalSelectDesc desc = (ExternalSelectDesc) palloc0(sizeof(ExternalSelectDesc));
	desc->projInfo = NULL;

	((ExtProtocolData*) fcinfo->context)->desc = desc;

	Relation relation = (Relation) palloc0(sizeof(Relation));

	((ExtProtocolData *) fcinfo->context)->prot_relation = relation;

	/* set mock behavior for uri parsing */
	GPHDUri *gphd_uri = palloc0(sizeof(GPHDUri));

	expect_string(parseGPHDUri, uri_str, uri_param);
	will_return(parseGPHDUri, gphd_uri);

	/* set mock behavior for set fragments */
	gphd_uri->fragments = palloc0(sizeof(List));
	expect_value(get_fragments, uri, gphd_uri);
	expect_value(get_fragments, relation, relation);
	will_assign_memory(get_fragments, uri, gphd_uri, sizeof(GPHDUri));
	will_be_called(get_fragments);

	/* set mock behavior for bridge import start -- nothing here */
	expect_any(gpbridge_import_start, context);
	will_be_called(gpbridge_import_start);

	/* set mock behavior for bridge read */
	const int EXPECTED_SIZE = 31;

	/* expected number of bytes read from the bridge */
	expect_any(gpbridge_read, context);
	expect_value(gpbridge_read, databuf, EXTPROTOCOL_GET_DATABUF(fcinfo));
	expect_value(gpbridge_read, datalen, EXTPROTOCOL_GET_DATALEN(fcinfo));
	will_return(gpbridge_read, EXPECTED_SIZE);

	Datum d = pxfprotocol_import(fcinfo);

	/* return number of bytes read from the bridge */
	assert_int_equal(DatumGetInt32(d), EXPECTED_SIZE);
	gphadoop_context context = (gphadoop_context ) EXTPROTOCOL_GET_USER_CTX(fcinfo);

	/* context has been created */
	assert_true(context != NULL);
	/* gphduri has been parsed */
	assert_true(context->gphd_uri != NULL);
	/* uri has been initialized */
	assert_true(context->uri.data != NULL);
	/* no write file name for import case */
	assert_true(context->relation != NULL);
	/* relation pointer is copied */
	assert_int_equal(context->relation, relation);

	/* cleanup */
	pfree(desc);
	pfree(relation);
	pfree(gphd_uri);
	pfree(EXTPROTOCOL_GET_USER_CTX(fcinfo));
	pfree(EXTPROTOCOL_GET_DATABUF(fcinfo));
	pfree(fcinfo->context);
	pfree(fcinfo);
	}

	static void
	test_pxfprotocol_import_second_call(void **state)
	{
	/* setup call info with call context */
	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
	fcinfo->context = palloc0(sizeof(ExtProtocolData));
	fcinfo->context->type = T_ExtProtocolData;
	EXTPROTOCOL_GET_DATALEN(fcinfo) = 100;
	EXTPROTOCOL_GET_DATABUF(fcinfo) = palloc0(EXTPROTOCOL_GET_DATALEN(fcinfo));
	((ExtProtocolData *) fcinfo->context)->prot_last_call = false;
	gphadoop_context *call_context = palloc0(sizeof(gphadoop_context));

	EXTPROTOCOL_SET_USER_CTX(fcinfo, call_context);

	/* set mock behavior for bridge read */
	const int EXPECTED_SIZE = 0;

	/* expected number of bytes read from the bridge */
	expect_value(gpbridge_read, context, call_context);
	expect_value(gpbridge_read, databuf, EXTPROTOCOL_GET_DATABUF(fcinfo));
	expect_value(gpbridge_read, datalen, EXTPROTOCOL_GET_DATALEN(fcinfo));
	will_return(gpbridge_read, EXPECTED_SIZE);

	Datum d = pxfprotocol_import(fcinfo);

	assert_int_equal(DatumGetInt32(d), EXPECTED_SIZE);
	/* return number of bytes read from the bridge */
	assert_true(EXTPROTOCOL_GET_USER_CTX(fcinfo) == call_context);
	/* context is still the same */

	/* cleanup */
	pfree(call_context);
	pfree(EXTPROTOCOL_GET_DATABUF(fcinfo));
	pfree(fcinfo->context);
	pfree(fcinfo);
	}

	static void
	test_pxfprotocol_import_last_call(void **state)
	{
	/* setup call info with a call context and last call indicator */
	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
	fcinfo->context = palloc0(sizeof(ExtProtocolData));
	fcinfo->context->type = T_ExtProtocolData;
	gphadoop_context *call_context = palloc0(sizeof(gphadoop_context));

	EXTPROTOCOL_SET_USER_CTX(fcinfo, call_context);
	EXTPROTOCOL_SET_LAST_CALL(fcinfo);

	/* init data in context that will be cleaned up */
	initStringInfo(&call_context->uri);

	/* set mock behavior for bridge cleanup */
	expect_value(gpbridge_cleanup, context, call_context);
	will_be_called(gpbridge_cleanup);

	Datum d = pxfprotocol_import(fcinfo);

	assert_int_equal(DatumGetInt32(d), 0);
	/* 0 is returned from function */
	assert_true(EXTPROTOCOL_GET_USER_CTX(fcinfo) == NULL);
	/* call context is cleaned up */

	/* cleanup */
	pfree(fcinfo->context);
	pfree(fcinfo);
	}

	static void
	test_pxfprotocol_export_first_call(void **state)
	{
	/* setup call info with no call context */
	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
	fcinfo->context = palloc0(sizeof(ExtProtocolData));
	fcinfo->context->type = T_ExtProtocolData;
	EXTPROTOCOL_GET_DATALEN(fcinfo) = 100;
	EXTPROTOCOL_GET_DATABUF(fcinfo) = palloc0(EXTPROTOCOL_GET_DATALEN(fcinfo));
	((ExtProtocolData *) fcinfo->context)->prot_last_call = false;
	((ExtProtocolData *) fcinfo->context)->prot_url = uri_param;

	ExternalSelectDesc desc = (ExternalSelectDesc) palloc0(sizeof(ExternalSelectDesc));
	desc->projInfo = NULL;

	((ExtProtocolData*) fcinfo->context)->desc = desc;

	Relation relation = (Relation) palloc0(sizeof(Relation));

	((ExtProtocolData *) fcinfo->context)->prot_relation = relation;

	/* set mock behavior for uri parsing */
	GPHDUri *gphd_uri = palloc0(sizeof(GPHDUri));

	expect_string(parseGPHDUri, uri_str, uri_param);
	will_return(parseGPHDUri, gphd_uri);

	/* set mock behavior for bridge export start -- nothing here */
	expect_any(gpbridge_export_start, context);
	will_be_called(gpbridge_export_start);

	/* set mock behavior for bridge write */
	const int EXPECTED_SIZE = 31;

	/* expected number of bytes written to the bridge */
	expect_any(gpbridge_write, context);
	expect_value(gpbridge_write, databuf, EXTPROTOCOL_GET_DATABUF(fcinfo));
	expect_value(gpbridge_write, datalen, EXTPROTOCOL_GET_DATALEN(fcinfo));
	will_return(gpbridge_write, EXPECTED_SIZE);

	Datum d = pxfprotocol_export(fcinfo);

	/* return number of bytes written to the bridge */
	assert_int_equal(DatumGetInt32(d), EXPECTED_SIZE);
	gphadoop_context context = (gphadoop_context ) EXTPROTOCOL_GET_USER_CTX(fcinfo);

	/* context has been created */
	assert_true(context != NULL);
	/* gphduri has been parsed */
	assert_true(context->gphd_uri != NULL);
	/* uri has been initialized */
	assert_true(context->uri.data != NULL);
	/* write file name initialized, but empty, since it is filled by another component */
	assert_true(context->relation != NULL);
	/* relation pointer is copied */
	assert_int_equal(context->relation, relation);

	/* cleanup */
	pfree(desc);
	pfree(relation);
	pfree(gphd_uri);
	pfree(EXTPROTOCOL_GET_USER_CTX(fcinfo));
	pfree(EXTPROTOCOL_GET_DATABUF(fcinfo));
	pfree(fcinfo->context);
	pfree(fcinfo);
	}

	static void
	test_pxfprotocol_export_second_call(void **state)
	{
	/* setup call info with call context */
	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
	fcinfo->context = palloc0(sizeof(ExtProtocolData));
	fcinfo->context->type = T_ExtProtocolData;
	EXTPROTOCOL_GET_DATALEN(fcinfo) = 100;
	EXTPROTOCOL_GET_DATABUF(fcinfo) = palloc0(EXTPROTOCOL_GET_DATALEN(fcinfo));
	((ExtProtocolData *) fcinfo->context)->prot_last_call = false;
	gphadoop_context *call_context = palloc0(sizeof(gphadoop_context));

	EXTPROTOCOL_SET_USER_CTX(fcinfo, call_context);

	/* set mock behavior for bridge write */
	const int EXPECTED_SIZE = 0;

	/* expected number of bytes written to the bridge */
	expect_value(gpbridge_write, context, call_context);
	expect_value(gpbridge_write, databuf, EXTPROTOCOL_GET_DATABUF(fcinfo));
	expect_value(gpbridge_write, datalen, EXTPROTOCOL_GET_DATALEN(fcinfo));
	will_return(gpbridge_write, EXPECTED_SIZE);

	Datum d = pxfprotocol_export(fcinfo);

	assert_int_equal(DatumGetInt32(d), EXPECTED_SIZE);
	/* return number of bytes written to the bridge */
	assert_true(EXTPROTOCOL_GET_USER_CTX(fcinfo) == call_context);
	/* context is still the same */

	/* cleanup */
	pfree(call_context);
	pfree(EXTPROTOCOL_GET_DATABUF(fcinfo));
	pfree(fcinfo->context);
	pfree(fcinfo);
	}

	static void
	test_pxfprotocol_export_last_call(void **state)
	{
	/* setup call info with a call context and last call indicator */
	PG_FUNCTION_ARGS = palloc0(sizeof(FunctionCallInfoData));
	fcinfo->context = palloc0(sizeof(ExtProtocolData));
	fcinfo->context->type = T_ExtProtocolData;
	gphadoop_context *call_context = palloc0(sizeof(gphadoop_context));

	EXTPROTOCOL_SET_USER_CTX(fcinfo, call_context);
	EXTPROTOCOL_SET_LAST_CALL(fcinfo);

	/* init data in context that will be cleaned up */
	initStringInfo(&call_context->uri);

	/* set mock behavior for bridge cleanup */
	expect_value(gpbridge_cleanup, context, call_context);
	will_be_called(gpbridge_cleanup);

	Datum d = pxfprotocol_export(fcinfo);

	assert_int_equal(DatumGetInt32(d), 0);
	/* 0 is returned from function */
	assert_true(EXTPROTOCOL_GET_USER_CTX(fcinfo) == NULL);
	/* call context is cleaned up */

	/* cleanup */
	pfree(fcinfo->context);
	pfree(fcinfo);
	}
	/* test setup and teardown methods */
	static void
	before_test(void **state)
	{
	/* set global variables */
	GpIdentity.segindex = 0;
	}

	static void
	after_test(void **state)
	{
	/*
	* no-op, but the teardown seems to be required when the test fails,
	* otherwise CMockery issues a mismatch error
	*/
	}

	int
	main(int argc, char *argv[])
	{
	cmockery_parse_arguments(argc, argv);

	const UnitTest tests[] = {
	unit_test(test_pxfprotocol_validate_urls),
	unit_test_setup_teardown(test_pxfprotocol_import_first_call, before_test, after_test),
	unit_test_setup_teardown(test_pxfprotocol_import_second_call, before_test, after_test),
	unit_test_setup_teardown(test_pxfprotocol_import_last_call, before_test, after_test),
	unit_test_setup_teardown(test_pxfprotocol_export_first_call, before_test, after_test),
	unit_test_setup_teardown(test_pxfprotocol_export_second_call, before_test, after_test),
	unit_test_setup_teardown(test_pxfprotocol_export_last_call, before_test, after_test)

	};

	MemoryContextInit();

	return run_tests(tests);
	}