def setup_method(self):
     """Setup the test environment."""
     # pylint: disable=attribute-defined-outside-init
     self.elasticsearch1 = Elasticsearch("endpoint:9201")
     self.elasticsearch2 = Elasticsearch("endpoint:9202")
     self.cluster1 = ec.ElasticsearchCluster(self.elasticsearch1,
                                             None,
                                             dry_run=False)
     self.cluster2 = ec.ElasticsearchCluster(self.elasticsearch2,
                                             None,
                                             dry_run=False)
     self.clusters = [self.cluster1, self.cluster2]
 def test_no_call_to_freeze_write_in_dry_run(self):
     """Test that when dry run is enabled, call to write to cluster index to freeze write is not placed."""
     self.elasticsearch1.index = mock.Mock(return_value=True)
     self.elasticsearch2.delete = mock.Mock(return_value=True)
     cluster1 = ec.ElasticsearchCluster(self.elasticsearch1,
                                        None,
                                        dry_run=True)
     cluster2 = ec.ElasticsearchCluster(self.elasticsearch2,
                                        None,
                                        dry_run=True)
     reason = Reason("test", "test_user", "test_host", task_id="T111222")
     elasticsearch_clusters = ec.ElasticsearchClusters([cluster1, cluster2],
                                                       None, None,
                                                       ["eqiad", "codfw"])
     with elasticsearch_clusters.frozen_writes(reason):
         assert not self.elasticsearch1.index.called
         assert not self.elasticsearch2.delete.called
def test_cluster_settings_are_unchanged_when_stopped_replication_is_dry_run():
    """Check that cluster routing in dry run mode is truly safe."""
    elasticsearch = Elasticsearch("endpoint:9200")
    elasticsearch.cluster.put_settings = mock.Mock(return_value=True)
    elasticsearch_cluster = ec.ElasticsearchCluster(elasticsearch,
                                                    None,
                                                    dry_run=True)
    with elasticsearch_cluster.stopped_replication():
        assert not elasticsearch.cluster.put_settings.called
def test_get_nodes_wraps_exceptions():
    """Get nodes should wrap exceptions from elasticsearch client."""
    elasticsearch = mock.Mock()
    elasticsearch.nodes.info = mock.Mock(
        side_effect=TransportError(500, "test"))
    remote = mock.Mock()

    cluster = ec.ElasticsearchCluster(elasticsearch, remote)
    with pytest.raises(ec.ElasticsearchClusterError):
        cluster.get_nodes()
def mock_node_info(values):
    """Creates a list of ElasticsearchCluster which will return the given node info."""
    clusters = []
    port = 9200
    for nodes in values:
        elasticsearch = Elasticsearch(f"localhost:{port}")
        port += 1
        elasticsearch.nodes.info = mock.Mock(return_value={"nodes": nodes})
        cluster = ec.ElasticsearchCluster(elasticsearch, None, dry_run=False)
        clusters.append(cluster)
    return clusters
    def test_frozen_writes_write_to_index(self):
        """Test that elasticsearch write to index is called to freeze writes."""
        self.elasticsearch1.index = mock.Mock(return_value=True)
        self.elasticsearch2.index = mock.Mock(return_value=True)
        self.elasticsearch1.delete = mock.Mock(return_value=True)
        self.elasticsearch2.delete = mock.Mock(return_value=True)
        reason = Reason("test", "test_user", "test_host", task_id="T111222")
        cluster1 = ec.ElasticsearchCluster(self.elasticsearch1,
                                           None,
                                           dry_run=False)
        cluster2 = ec.ElasticsearchCluster(self.elasticsearch2,
                                           None,
                                           dry_run=False)
        elasticsearch_clusters = ec.ElasticsearchClusters([cluster1, cluster2],
                                                          None, None,
                                                          ["eqiad", "codfw"])
        with elasticsearch_clusters.frozen_writes(reason):
            assert self.elasticsearch1.index.called
            assert self.elasticsearch2.index.called

        assert self.elasticsearch1.delete.called
        assert self.elasticsearch2.delete.called