def test_serialize(self):
req_matcher10 = RequestMatcher(policy=Policy.PRIORITIZE) \
.add_domain_re(r'^intranet$') \
.add_port(443)
req_matcher11 = RequestMatcher(policy=Policy.PRIORITIZE) \
.add_domain_re(r'^server-test$') \
.add_port(8080)
link1 = Link(interface="wlp5s0") \
.add_request_matcher(req_matcher10) \
.add_request_matcher(req_matcher11)
req_matcher2 = RequestMatcher(policy=Policy.FORBID) \
.add_domain_re(r'.com$')
link2 = Link(interface="enp3s0", weight=2) \
.add_request_matcher(req_matcher2)
balancer_req_matcher = RequestMatcher(policy=Policy.ALLOW) \
.add_port(443) \
.add_port(80) \
.add_port(8080)
balancer = Balancer(strategy=Strategy.LEAST_CONNECTIONS) \
.add_request_matcher(balancer_req_matcher) \
.add_link(link1) \
.add_link(link2)
server = Server() \
.set_balancer(balancer)
serialized = serialize(server)
deserialized = deserialize(serialized, Server)
self.assertEqual(serialized, serialize(deserialized))
save_server("test.json", server)
def test_should_give_high_priority_if_matcher_says_so(self):
link = Link()
matcher = RequestMatcher(Policy.PRIORITIZE)
matcher.add_domain_re(r'^.+\.com$')
link.add_request_matcher(matcher)
actual = link.get_request_priority_level(Request("google.com", 80))
self.assertEqual(actual, PriorityLevel.HIGH)
def test_should_forbid_request_if_domain_is_not_in_whitelist(self):
link = Link()
matcher = RequestMatcher(Policy.ALLOW)
matcher.add_domain_re(r'^.+\.com$')
link.add_request_matcher(matcher)
actual = link.get_request_priority_level(Request("google.fr", 80))
self.assertEqual(actual, PriorityLevel.FORBID)
def test_round_robin_should_return_next_one(self):
link_1 = Link()
link_2 = Link()
link_3 = Link()
links = [link_1, link_2, link_3]
actual = round_robin.get_next_link(links, last_link=link_1)
self.assertEqual(actual, link_2)
def test_should_almost_always_return_element(self):
link_1 = Link(weight=1)
link_2 = Link(weight=100000)
link_3 = Link(weight=1)
links = [link_1, link_2, link_3]
next_link = random_link.get_next_link(links)
self.assertEqual(next_link, link_2)
def test_round_robin_should_return_the_first_one_when_no_last_link(self):
link_1 = Link()
link_2 = Link()
link_3 = Link()
links = [link_1, link_2, link_3]
actual = round_robin.get_next_link(links)
self.assertEqual(actual, link_1)
def test_open_and_close_connection_link(self):
link = Link()
self.assertEqual(len(link.connections), 0)
socket_connection = link.open_connection("1")
self.assertEqual(len(link.connections), 1)
self.assertFalse(socket_connection._closed)
self.assertIsInstance(socket_connection, socks.socksocket)
link.close_connection("1")
self.assertEqual(len(link.connections), 0)
self.assertTrue(socket_connection._closed)
def test_should_return_random_idx(self):
link_1 = Link()
link_2 = Link()
link_3 = Link()
links = [link_1, link_2, link_3]
next_linkA = random_link.get_next_link(links)
next_linkB = random_link.get_next_link(links)
next_linkC = random_link.get_next_link(links)
next_linkD = random_link.get_next_link(links)
next_linkE = random_link.get_next_link(links)
self.assertFalse(
next_linkA == next_linkB == next_linkC == next_linkD == next_linkE)
def test_least_connections_should_return_the_link_with_least_connections(
self):
link_1 = Link()
link_2 = Link()
link_3 = Link()
link_1.open_connection("1")
link_3.open_connection("1")
links = [link_1, link_2, link_3]
actual = least_connections.get_next_link(links)
self.assertEqual(actual, link_2)
def test_should_return_next_link_in_round_robin(self):
expected_first_link = Link(domain="Link1", port=1080)
expected_second_link = Link(domain="Link2", port=1081)
expected_third_link = Link(domain="Link3", port=1082)
balancer = Balancer() \
.set_strategy(Strategy.ROUND_ROBIN) \
.add_link(expected_first_link) \
.add_link(expected_second_link) \
.add_link(expected_third_link)
actual_first_link = balancer.get_next_link(Request('test', 80))
self.assertEqual(actual_first_link, expected_first_link)
actual_second_link = balancer.get_next_link(Request('test', 80))
self.assertEqual(actual_second_link, expected_second_link)
actual_third_link = balancer.get_next_link(Request('test', 80))
self.assertEqual(actual_third_link, expected_third_link)
actual_first_link = balancer.get_next_link(Request('test', 80))
self.assertEqual(expected_first_link, actual_first_link)
def test_should_allow_request(self):
expected_link = Link(domain="Link1", port=1080)
matcher = RequestMatcher(policy=Policy.ALLOW)\
.add_port(80)
balancer = Balancer() \
.set_strategy(Strategy.ROUND_ROBIN) \
.add_link(expected_link)\
.add_request_matcher(matcher)
should_be_link = balancer.get_next_link(Request('test', 80))
self.assertEqual(expected_link, should_be_link)
def test_should_block_request(self):
expected_link = Link(domain="Link1", port=1080)
matcher = RequestMatcher(policy=Policy.FORBID)\
.add_port(80)
balancer = Balancer() \
.set_strategy(Strategy.ROUND_ROBIN) \
.add_link(expected_link)\
.add_request_matcher(matcher)
should_be_none = balancer.get_next_link(Request('test', 80))
self.assertIsNone(should_be_none)
def test_should_give_normal_priority_if_prioritized_and_deprioritized_at_the_same_time(
self):
link = Link()
matcher = RequestMatcher(Policy.ALLOW)
matcher.add_domain_re(r'^.+\.com$')
link.add_request_matcher(matcher)
matcher = RequestMatcher(Policy.FORBID)
matcher.add_domain_re(r'^.+\.com$')
link.add_request_matcher(matcher)
actual = link.get_request_priority_level(Request("google.com", 80))
self.assertEqual(actual, PriorityLevel.FORBID)
def test_should_give_normal_priority_if_no_matcher_is_given(self):
link = Link()
actual = link.get_request_priority_level(Request("google.com", 80))
self.assertEqual(actual, PriorityLevel.NORMAL)