<commit_before>import sys
import time
import unittest
import yaml

from mock import Mock, patch
from patroni.api import RestApiServer
from patroni.async_executor import AsyncExecutor
from patroni.etcd import Etcd
from patroni import Patroni, main as _main
from patroni.zookeeper import ZooKeeper
from six.moves import BaseHTTPServer
from test_etcd import Client, SleepException, etcd_read, etcd_write
from test_postgresql import Postgresql, psycopg2_connect
from test_zookeeper import MockKazooClient


def time_sleep(*args):
    raise SleepException()


@patch('time.sleep', Mock())
@patch('subprocess.call', Mock(return_value=0))
@patch('psycopg2.connect', psycopg2_connect)
@patch.object(Postgresql, 'write_pg_hba', Mock())
@patch.object(Postgresql, 'write_recovery_conf', Mock())
@patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
@patch.object(AsyncExecutor, 'run', Mock())
class TestPatroni(unittest.TestCase):

    def setUp(self):
        with patch.object(Client, 'machines') as mock_machines:
            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
            self.touched = False
            self.init_cancelled = False
            RestApiServer._BaseServer__is_shut_down = Mock()
            RestApiServer._BaseServer__shutdown_request = True
            RestApiServer.socket = 0
            with open('postgres0.yml', 'r') as f:
                config = yaml.load(f)
                self.p = Patroni(config)
                self.p.ha.dcs.client.write = etcd_write
                self.p.ha.dcs.client.read = etcd_read

    @patch('patroni.zookeeper.KazooClient', MockKazooClient())
    def test_get_dcs(self):
        self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
        self.assertRaises(Exception, self.p.get_dcs, '', {})

    @patch('time.sleep', Mock(side_effect=SleepException()))
    @patch.object(Etcd, 'delete_leader', Mock())
    @patch.object(Client, 'machines')
    def test_patroni_main(self, mock_machines):
        _main()
        sys.argv = ['patroni.py', 'postgres0.yml']

        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
        with patch.object(Patroni, 'run', Mock(side_effect=SleepException())):
            self.assertRaises(SleepException, _main)
        with patch.object(Patroni, 'run', Mock(side_effect=KeyboardInterrupt())):
            _main()

    @patch('time.sleep', Mock(side_effect=SleepException()))
    def test_run(self):
        self.p.ha.dcs.watch = time_sleep
        self.assertRaises(SleepException, self.p.run)

        self.p.ha.state_handler.is_leader = Mock(return_value=False)
        self.p.api.start = Mock()
        self.assertRaises(SleepException, self.p.run)
<commit_msg>Mock dcs.watch directly instead of using wraper<commit_after>import sys
import time
import unittest
import yaml

from mock import Mock, patch
from patroni.api import RestApiServer
from patroni.async_executor import AsyncExecutor
from patroni.etcd import Etcd
from patroni import Patroni, main as _main
from patroni.zookeeper import ZooKeeper
from six.moves import BaseHTTPServer
from test_etcd import Client, SleepException, etcd_read, etcd_write
from test_postgresql import Postgresql, psycopg2_connect
from test_zookeeper import MockKazooClient


@patch('time.sleep', Mock())
@patch('subprocess.call', Mock(return_value=0))
@patch('psycopg2.connect', psycopg2_connect)
@patch.object(Postgresql, 'write_pg_hba', Mock())
@patch.object(Postgresql, 'write_recovery_conf', Mock())
@patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
@patch.object(AsyncExecutor, 'run', Mock())
class TestPatroni(unittest.TestCase):

    def setUp(self):
        with patch.object(Client, 'machines') as mock_machines:
            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
            self.touched = False
            self.init_cancelled = False
            RestApiServer._BaseServer__is_shut_down = Mock()
            RestApiServer._BaseServer__shutdown_request = True
            RestApiServer.socket = 0
            with open('postgres0.yml', 'r') as f:
                config = yaml.load(f)
                self.p = Patroni(config)
                self.p.ha.dcs.client.write = etcd_write
                self.p.ha.dcs.client.read = etcd_read

    @patch('patroni.zookeeper.KazooClient', MockKazooClient())
    def test_get_dcs(self):
        self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
        self.assertRaises(Exception, self.p.get_dcs, '', {})

    @patch('time.sleep', Mock(side_effect=SleepException()))
    @patch.object(Etcd, 'delete_leader', Mock())
    @patch.object(Client, 'machines')
    def test_patroni_main(self, mock_machines):
        _main()
        sys.argv = ['patroni.py', 'postgres0.yml']

        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
        with patch.object(Patroni, 'run', Mock(side_effect=SleepException())):
            self.assertRaises(SleepException, _main)
        with patch.object(Patroni, 'run', Mock(side_effect=KeyboardInterrupt())):
            _main()

    @patch('time.sleep', Mock(side_effect=SleepException()))
    def test_run(self):
        self.p.ha.dcs.watch = Mock(side_effect=SleepException())
        self.assertRaises(SleepException, self.p.run)

        self.p.ha.state_handler.is_leader = Mock(return_value=False)
        self.p.api.start = Mock()
        self.assertRaises(SleepException, self.p.run)

{'fraction_non_alphanumeric': 0.125544267053701, 'fraction_numerical': 0.00544267053701016, 'mean_word_length': 3.728130360205832, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '4972957', 'n_tokens_mistral': 1836, 'n_tokens_neox': 1759, 'n_words': 330}