• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2024 The Chromium Authors
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import argparse
6import pathlib
7from unittest import mock
8
9from crossbench.network.traffic_shaping.ts_proxy import (TsProxyProcess,
10                                                         TsProxyServer,
11                                                         TsProxyTrafficShaper)
12from tests import test_helper
13from tests.crossbench.base import BaseCrossbenchTestCase
14
15
16class TsProxyBaseTestCase(BaseCrossbenchTestCase):
17
18  def setUp(self) -> None:
19    super().setUp()
20    self.ts_proxy_path = pathlib.Path("/chrome/tsproxy/tsproxy.py")
21    self.fs.create_file(self.ts_proxy_path, st_size=100)
22    # Avoid dealing with fcntl for testing.
23    patcher = mock.patch.object(
24        TsProxyProcess, "_setup_non_blocking_io", return_value=None)
25    self.addCleanup(patcher.stop)
26    patcher.start()
27
28
29class TsProxyTestCase(TsProxyBaseTestCase):
30
31  def test_ts_proxy_traffic_shaper_no_tsproxy(self):
32    with self.assertRaises(RuntimeError):
33      TsProxyTrafficShaper(self.platform)
34
35  def test_ts_proxy_traffic_shaper_default(self):
36    ts_proxy = TsProxyTrafficShaper(self.platform, self.ts_proxy_path)
37    self.assertFalse(ts_proxy.is_running)
38
39
40class TsProxyServerTestCase(TsProxyBaseTestCase):
41
42  def test_construct_invalid(self):
43    with self.assertRaises(argparse.ArgumentTypeError):
44      TsProxyServer(pathlib.Path("does/not/exist"))
45
46  def test_basic_instance(self):
47    server = TsProxyServer(self.ts_proxy_path)
48    self.assertFalse(server.is_running)
49
50    with self.assertRaises(AssertionError):
51      server.set_traffic_settings()
52    with self.assertRaises(AssertionError):
53      _ = server.socks_proxy_port
54    self.assertIsNone(server.stop())
55
56  def test_basic_instance_http_port(self):
57    server = TsProxyServer(self.ts_proxy_path, http_port=8080)
58    self.assertFalse(server.is_running)
59    with self.assertRaises(AssertionError):
60      _ = server.socks_proxy_port
61    self.assertIsNone(server.stop())
62
63  def test_ports(self):
64    with self.assertRaises(ValueError):
65      TsProxyServer(self.ts_proxy_path, https_port=400)
66    with self.assertRaises(ValueError):
67      TsProxyServer(self.ts_proxy_path, http_port=400, https_port=400)
68    with self.assertRaises(argparse.ArgumentTypeError):
69      TsProxyServer(self.ts_proxy_path, http_port=-400, https_port=400)
70    with self.assertRaises(argparse.ArgumentTypeError):
71      TsProxyServer(self.ts_proxy_path, http_port=400, https_port=-400)
72
73  def test_start_server(self):
74    server = TsProxyServer(self.ts_proxy_path)
75
76    proc = mock.Mock()
77    proc.configure_mock(**{
78        "poll.return_value": None,
79        "communicate.return_value": (None, None)
80    })
81    proc.stdout = mock.Mock()
82    proc.stdout.configure_mock(**{
83        "readline.return_value":
84            "Started Socks5 proxy server on 127.0.0.1:43210"
85    })
86    proc.stderr = mock.Mock()
87
88    def popen_mock(cmd, *args, **kwargs):
89      self.assertEqual(cmd[1], self.ts_proxy_path)
90      self.assertEqual(cmd[2], "--port=0")
91      del args, kwargs
92      return proc
93
94    with mock.patch("subprocess.Popen", side_effect=popen_mock) as popen:
95      self.assertFalse(server.is_running)
96      with server:
97        self.assertTrue(server.is_running)
98        self.assertEqual(server.socks_proxy_port, 43210)
99        proc.stdout.readline.assert_called_once()
100        # Set return value for exit command.
101        proc.stdout.readline.return_value = "OK"
102      self.assertFalse(server.is_running)
103
104    popen.assert_called_once()
105    proc.stdin.write.assert_called_with("exit\n")
106
107
108if __name__ == "__main__":
109  test_helper.run_pytest(__file__)
110