<commit_before>package org.hashtrees.test;

import static org.hashtrees.test.utils.HashTreesImplTestUtils.DEFAULT_SEG_DATA_BLOCKS_COUNT;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.DEFAULT_TREE_ID;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.ROOT_NODE;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.SEG_ID_PROVIDER;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.TREE_ID_PROVIDER;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.createHashTree;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.generateInMemoryAndPersistentStores;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.generateInMemoryStore;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.generatePersistentStore;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.randomByteBuffer;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.randomBytes;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

import org.apache.commons.codec.binary.Hex;
import org.hashtrees.HashTrees;
import org.hashtrees.HashTreesConstants;
import org.hashtrees.HashTreesImpl;
import org.hashtrees.SimpleTreeIdProvider;
import org.hashtrees.store.HashTreesManagerStore;
import org.hashtrees.store.HashTreesMemStore;
import org.hashtrees.store.HashTreesPersistentStore;
import org.hashtrees.store.HashTreesStore;
import org.hashtrees.store.SimpleMemStore;
import org.hashtrees.synch.HashTreesManager;
import org.hashtrees.synch.HashTreesThriftClientProvider;
import org.hashtrees.test.utils.HashTreesImplTestUtils;
import org.hashtrees.test.utils.HashTreesImplTestUtils.HTreeComponents;
import org.hashtrees.thrift.generated.SegmentData;
import org.hashtrees.thrift.generated.SegmentHash;
import org.hashtrees.thrift.generated.ServerName;
import org.hashtrees.util.ByteUtils;
import org.hashtrees.util.NonBlockingQueuingTask.QueueReachedMaxCapacityException;
import org.junit.Assert;
import org.junit.Test;

public class HashTreesImplTest {

	private static ByteBuffer generateBytesFrom(int... values) {
		byte[] result = new byte[values.length * ByteUtils.SIZEOF_INT];
		ByteBuffer bb = ByteBuffer.wrap(result);
		for (int value : values)
			bb.putInt(value);
		return bb;
	}

	private static ByteBuffer generateRandomKeyWithPrefix(int prefixValue) {
		byte[] randomBytes = randomBytes();
		byte[] key = new byte[ByteUtils.SIZEOF_INT + randomBytes.length];
		ByteBuffer bb = ByteBuffer.wrap(key);
		bb.putInt(prefixValue);
		bb.put(randomBytes);
		return bb;
	}

	@Test
	public void testPut() throws Exception {

		HashTreesStore[] stores = generateInMemoryAndPersistentStores();

		try {
			for (HashTreesStore store : stores) {
				int segId = 1;

				HTreeComponents components = createHashTree(
						DEFAULT_SEG_DATA_BLOCKS_COUNT, TREE_ID_PROVIDER,
						SEG_ID_PROVIDER, store);
				HashTrees testTree = components.hTree;
				HashTreesStore testTreeStore = components.hTStore;

				ByteBuffer key = generateRandomKeyWithPrefix(segId);
				ByteBuffer value = ByteBuffer.wrap(randomBytes());
				testTree.hPut(key, value);
				ByteBuffer expectedDigest = ByteBuffer.wrap(ByteUtils
						.sha1(value.array()));

				SegmentData actualKeyAndDigest = testTree.getSegmentData(
						SimpleTreeIdProvider.TREE_ID, segId, key);
				Assert.assertNotNull(actualKeyAndDigest);
				Assert.assertTrue(Arrays.equals(key.array(),
						actualKeyAndDigest.getKey()));
				Assert.assertTrue(Arrays.equals(expectedDigest.array(),
						actualKeyAndDigest.getDigest()));

				List<Integer> dirtySegs = testTreeStore
						.getDirtySegments(SimpleTreeIdProvider.TREE_ID);
				Assert.assertEquals(1, dirtySegs.size());
				Assert.assertEquals(1, dirtySegs.get(0).intValue());
			}
		} finally {
			HashTreesImplTestUtils.closeStores(stores);
		}
	}

	@Test
	public void testRemove() throws Exception {

		HashTreesStore[] stores = generateInMemoryAndPersistentStores();
		int segId = 2;

		try {
			for (HashTreesStore store : stores) {
				HTreeComponents components = createHashTree(
						DEFAULT_SEG_DATA_BLOCKS_COUNT, TREE_ID_PROVIDER,
						SEG_ID_PROVIDER, store);
				HashTrees testTree = components.hTree;
				HashTreesStore testTreeStore = components.hTStore;

				ByteBuffer key = generateRandomKeyWithPrefix(segId);
				ByteBuffer value = ByteBuffer.wrap(randomBytes());
				testTree.hPut(key, value);
				testTree.hRemove(key);

				SegmentData segData = testTree.getSegmentData(
						SimpleTreeIdProvider.TREE_ID, segId, key);
				Assert.assertNull(segData);

				List<Integer> dirtySegs = testTreeStore
						.getDirtySegments(SimpleTreeIdProvider.TREE_ID);
				Assert.assertEquals(1, dirtySegs.size());
				Assert.assertEquals(2, dirtySegs.get(0).intValue());
			}
		} finally {
			HashTreesImplTestUtils.closeStores(stores);
		}
	}

	@Test
	public void testUpdateSegmentHashesTest() throws Exception {

		int rootNodeId = 0;
		int nodeId = 2;<commit_msg>Adding test case for full rebuild operation<commit_after>package org.hashtrees.test;

import static org.hashtrees.test.utils.HashTreesImplTestUtils.DEFAULT_SEG_DATA_BLOCKS_COUNT;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.DEFAULT_TREE_ID;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.ROOT_NODE;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.SEG_ID_PROVIDER;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.TREE_ID_PROVIDER;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.createHashTree;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.generateInMemoryAndPersistentStores;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.generateInMemoryStore;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.generatePersistentStore;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.randomByteBuffer;
import static org.hashtrees.test.utils.HashTreesImplTestUtils.randomBytes;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

import org.apache.commons.codec.binary.Hex;
import org.hashtrees.HashTrees;
import org.hashtrees.HashTreesConstants;
import org.hashtrees.HashTreesImpl;
import org.hashtrees.SimpleTreeIdProvider;
import org.hashtrees.store.HashTreesManagerStore;
import org.hashtrees.store.HashTreesMemStore;
import org.hashtrees.store.HashTreesPersistentStore;
import org.hashtrees.store.HashTreesStore;
import org.hashtrees.store.SimpleMemStore;
import org.hashtrees.store.Store;
import org.hashtrees.synch.HashTreesManager;
import org.hashtrees.synch.HashTreesThriftClientProvider;
import org.hashtrees.test.utils.HashTreesImplTestUtils;
import org.hashtrees.test.utils.HashTreesImplTestUtils.HTreeComponents;
import org.hashtrees.thrift.generated.SegmentData;
import org.hashtrees.thrift.generated.SegmentHash;
import org.hashtrees.thrift.generated.ServerName;
import org.hashtrees.util.ByteUtils;
import org.hashtrees.util.NonBlockingQueuingTask.QueueReachedMaxCapacityException;
import org.junit.Assert;
import org.junit.Test;

public class HashTreesImplTest {

	private static ByteBuffer generateBytesFrom(int... values) {
		byte[] result = new byte[values.length * ByteUtils.SIZEOF_INT];
		ByteBuffer bb = ByteBuffer.wrap(result);
		for (int value : values)
			bb.putInt(value);
		return bb;
	}

	private static ByteBuffer generateRandomKeyWithPrefix(int prefixValue) {
		byte[] randomBytes = randomBytes();
		byte[] key = new byte[ByteUtils.SIZEOF_INT + randomBytes.length];
		ByteBuffer bb = ByteBuffer.wrap(key);
		bb.putInt(prefixValue);
		bb.put(randomBytes);
		return bb;
	}

	@Test
	public void testPut() throws Exception {

		HashTreesStore[] stores = generateInMemoryAndPersistentStores();

		try {
			for (HashTreesStore store : stores) {
				int segId = 1;

				HTreeComponents components = createHashTree(
						DEFAULT_SEG_DATA_BLOCKS_COUNT, TREE_ID_PROVIDER,
						SEG_ID_PROVIDER, store);
				HashTrees testTree = components.hTree;
				HashTreesStore testTreeStore = components.hTStore;

				ByteBuffer key = generateRandomKeyWithPrefix(segId);
				ByteBuffer value = ByteBuffer.wrap(randomBytes());
				testTree.hPut(key, value);
				ByteBuffer expectedDigest = ByteBuffer.wrap(ByteUtils
						.sha1(value.array()));

				SegmentData actualKeyAndDigest = testTree.getSegmentData(
						SimpleTreeIdProvider.TREE_ID, segId, key);
				Assert.assertNotNull(actualKeyAndDigest);
				Assert.assertTrue(Arrays.equals(key.array(),
						actualKeyAndDigest.getKey()));
				Assert.assertTrue(Arrays.equals(expectedDigest.array(),
						actualKeyAndDigest.getDigest()));

				List<Integer> dirtySegs = testTreeStore
						.getDirtySegments(SimpleTreeIdProvider.TREE_ID);
				Assert.assertEquals(1, dirtySegs.size());
				Assert.assertEquals(1, dirtySegs.get(0).intValue());
			}
		} finally {
			HashTreesImplTestUtils.closeStores(stores);
		}
	}

	@Test
	public void testRemove() throws Exception {

		HashTreesStore[] stores = generateInMemoryAndPersistentStores();
		int segId = 2;

		try {
			for (HashTreesStore store : stores) {
				HTreeComponents components = createHashTree(
						DEFAULT_SEG_DATA_BLOCKS_COUNT, TREE_ID_PROVIDER,
						SEG_ID_PROVIDER, store);
				HashTrees testTree = components.hTree;
				HashTreesStore testTreeStore = components.hTStore;

				ByteBuffer key = generateRandomKeyWithPrefix(segId);
				ByteBuffer value = ByteBuffer.wrap(randomBytes());
				testTree.hPut(key, value);
				testTree.hRemove(key);

				SegmentData segData = testTree.getSegmentData(
						SimpleTreeIdProvider.TREE_ID, segId, key);
				Assert.assertNull(segData);

				List<Integer> dirtySegs = testTreeStore
						.getDirtySegments(SimpleTreeIdProvider.TREE_ID);
				Assert.assertEquals(1, dirtySegs.size());
				Assert.assertEquals(2, dirtySegs.get(0).intValue());
			}
		} finally {
			HashTreesImplTestUtils.closeStores(stores);
		}
	}

	@Test
	public void testCompleteRebuild() throws Exception {
		int rootNodeId = 0;
		int nodeId = 2;
		int segId = 1;
		int noOfSegments = 2;
		HashTreesStore[] stores = generateInMemoryAndPersistentStores();

		try {
			for (HashTreesStore store : stores) {
				HTreeComponents components = createHashTree(noOfSegments,
						TREE_ID_PROVIDER, SEG_ID_PROVIDER, store);
				HashTrees testTree = components.hTree;
				Store kvStore = components.store;
				kvStore.registerHashTrees(null);

				ByteBuffer expectedKey = generateRandomKeyWithPrefix(segId);
				ByteBuffer expectedValue = ByteBuffer.wrap(randomBytes());
				StringBuffer sb = new StringBuffer();
				ByteBuffer expectedDigest = ByteBuffer.wrap(ByteUtils
						.sha1(expectedValue.array()));
				sb.append(HashTreesImpl.getHexString(expectedKey,
						expectedDigest) + "\n");
				byte[] expectedLeafNodeDigest = ByteUtils.sha1(sb.toString()
						.getBytes());

				testTree.hPut(randomByteBuffer(), randomByteBuffer());
				kvStore.put(expectedKey.array(), expectedValue.array());
				testTree.rebuildHashTree(SimpleTreeIdProvider.TREE_ID, true);
				SegmentHash segHash = testTree.getSegmentHash(
						SimpleTreeIdProvider.TREE_ID, nodeId);
				Assert.assertNotNull(segHash);
				Assert.assertTrue(Arrays.equals(expectedLeafNodeDigest,
						segHash.getHash()));

				sb.setLength(0);
				sb.append(Hex.encodeHexString(expectedLeafNodeDigest) + "\n");
				byte[] expectedRootNodeDigest = ByteUtils.sha1(sb.toString()
						.getBytes());
				SegmentHash actualRootNodeDigest = testTree.getSegmentHash(
						SimpleTreeIdProvider.TREE_ID, rootNodeId);
				Assert.assertNotNull(actualRootNodeDigest);
				Assert.assertTrue(Arrays.equals(expectedRootNodeDigest,
						actualRootNodeDigest.getHash()));
			}
		} finally {
			HashTreesImplTestUtils.closeStores(stores);
		}
	}

	@Test
	public void testUpdateSegmentHashesTest() throws Exception {

		int rootNodeId = 0;
		int nodeId = 2;
{'fraction_non_alphanumeric': 0.09657579787234043, 'fraction_numerical': 0.0024933510638297874, 'mean_word_length': 5.353220696937698, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5003954', 'n_tokens_mistral': 4314, 'n_tokens_neox': 3810, 'n_words': 624}