#include "ThreadPinning.hpp"

#include "UnitTest++/UnitTest++.h"

#include <algorithm>
#include <thread>
#include <vector>

using namespace tp;

namespace {
	std::vector<int> sorted(std::vector<int> values) {
		std::sort(values.begin(), values.end());
		return values;
	}
}

SUITE(ThreadPinning) {
	TEST(SiblingGroupsCoverEveryAllowedCpuExactlyOnce) {
		const std::vector<int> allowed = ThreadPinning::allowedCpus();
		if (allowed.empty()) return;
		const auto groups = ThreadPinning::siblingGroups(allowed);
		std::vector<int> seen;
		for (const auto& group : groups)
			seen.insert(seen.end(), group.begin(), group.end());
		CHECK(sorted(seen) == sorted(allowed));
	}

	TEST(PinningRestrictsTheCallingThread) {
		const std::vector<int> allowed = ThreadPinning::allowedCpus();
		const auto groups = ThreadPinning::siblingGroups(allowed);
		if (groups.size() < 2) return;
		ThreadPinning::pinCallingThread(groups.front());
		CHECK(sorted(ThreadPinning::allowedCpus()) == sorted(groups.front()));
	}

	TEST(TwoThreadsCanBePinnedToDifferentGroups) {
		const auto groups = ThreadPinning::siblingGroups();
		if (groups.size() < 2) return;

		std::vector<int> first, second;
		std::thread left([&]() {
			ThreadPinning::pinCallingThread(groups[0]);
			first = ThreadPinning::allowedCpus();
		});
		std::thread right([&]() {
			ThreadPinning::pinCallingThread(groups[1]);
			second = ThreadPinning::allowedCpus();
		});
		left.join();
		right.join();

		CHECK(sorted(first) == sorted(groups[0]));
		CHECK(sorted(second) == sorted(groups[1]));
	}
}