@Test
public void testRetrieveDocumentSet() throws Throwable {

	// Arrange
	when(xdsbRegistry.getPatientUniqueId(anyString(), anyString()))
	.thenReturn(patientUniqueId);
	when(contextHandler.enforcePolicy(isA(XacmlRequest.class))).thenReturn(
			xacmlResponse).thenReturn(xacmlResponse);
	when(
			xdsbRepository.retrieveDocumentSet(
					any(RetrieveDocumentSetRequest.class), anyString(),
					anyString())).thenReturn(retrieveDocumentSetResponse);
	when(marshaller.marshal(isA(Object.class))).thenReturn(
			enforcementPolicies).thenReturn(enforcementPolicies);
	when(documentSegmentation.segmentDocument(isA(DSSRequest.class)))
	.thenReturn(segmentDocumentResponse1).thenReturn(
			segmentDocumentResponse2);
	when(context.getUserPrincipal()).thenReturn(samlTokenPrincipal);
	when(
			samlTokenParser.parse(samlTokenPrincipal.getToken(),
					"urn:oasis:names:tc:xacml:1.0:resource:resource-id"))
					.thenReturn(
							"d3bb3930-7241-11e3-b4f7-00155d3a2124^^^&2.16.840.1.113883.4.357&ISO");
	when(
			samlTokenParser.parse(samlTokenPrincipal.getToken(),
					"urn:oasis:names:tc:xacml:1.0:subject-category:intermediary-subject"))
					.thenReturn("1114252178");
	when(
			samlTokenParser.parse(samlTokenPrincipal.getToken(),
					"urn:oasis:names:tc:xspa:1.0:subject:purposeofuse"))
					.thenReturn("TREATMENT");
	when(
			samlTokenParser.parse(samlTokenPrincipal.getToken(),
					"urn:oasis:names:tc:xacml:1.0:subject-category:recipient-subject"))
					.thenReturn("1760717789");
	final Map<PredicateKey, String> mockMap = mock(Map.class);
	when(auditService.createPredicateMap()).thenReturn(mockMap);
	when(aspectUtils.getInXMLFormat(any())).thenReturn(anyString());
	// Act
	final RetrieveDocumentSetResponse response = pep
			.retrieveDocumentSet(retrieveDocumentSetRequest);
	// Assert
	assertEquals("1111048131312.1312011.43613.91155.1111312013196213",
			response.getDocumentResponse().get(0).getDocumentUniqueId());
	assertEquals("44212111915.14792.41394.914113.0152268155381113",
			response.getDocumentResponse().get(1).getDocumentUniqueId());
	verify(documentSegmentation, times(2)).segmentDocument(
			argThat(matching((DSSRequest req) -> enforcementPolicies
					.equals(req.getEnforcementPoliciesXml())
					&& req.isAudited().equals(true)
					&& req.isAuditFailureByPass().equals(false))));
	// assertTrue(Arrays.equals(segmentDocumentResponse2String.getBytes(),
	// response.getDocumentResponse().get(1).getDocument()));
}