from unittest.mock import Mock
from unittest.mock import patch
import csv
import io

from labelaudit.logic.stream_to_s3 import StreamToS3


mock_key = 'foo'
mock_bucket = 'bar'
mock_dict_keys = ['first', 'second', 'third']
mock_display_header = ['First Col', 'Second Col', 'Third Col']


def mock_set_csv_writer(monkeypatch):
    monkeypatch.setattr(StreamToS3, 'set_csv_writer', Mock(return_value=None))


def test_display_column_heading_optional(monkeypatch):
    s = StreamToS3(mock_key, mock_bucket, mock_dict_keys)
    assert s.header == mock_dict_keys


def test_when_display_header_sent(monkeypatch):
    """Test when display header is sent, header uses that and not dict keys."""
    s = StreamToS3(mock_key, mock_bucket, mock_dict_keys, mock_display_header)
    assert s.header == mock_display_header


def test_clear_stream(monkeypatch):
    value = 'hi'
    monkeypatch.setattr(csv, 'DictWriter', Mock(return_value=True))
    monkeypatch.setattr(io, 'StringIO', Mock(return_value=value))
    s = StreamToS3(mock_key, mock_bucket, mock_dict_keys)
    s.stream = 'test'
    s.set_csv_writer(mock_dict_keys)
    assert s.stream == value


def test_do_not_refresh_stream(monkeypatch):
    monkeypatch.setattr(csv, 'DictWriter', Mock(return_value=True))
    # monkeypatch.setattr(io, 'StringIO', Mock(return_value='hi'))
    s = StreamToS3(mock_key, mock_bucket, mock_dict_keys)
    s.stream = 'test'
    s.set_csv_writer(mock_dict_keys, False)
    assert s.stream == 'test'


def test_write_column_headers_calls_resources(monkeypatch):
    mock_set_csv_writer(monkeypatch)
    with patch('csv.DictWriter') as MockClass:
        instance = MockClass.return_value
        instance.writeheader.return_value = 'foo'
        instance.write.return_value = 'bar'
        instance.getvalue.return_value = 'foobar'
        s = StreamToS3(mock_key, mock_bucket, mock_dict_keys)
        s.csv_writer = instance
        s.s3writer = instance
        s.stream = instance
        s.write_column_headers()
        assert s.csv_writer.writeheader() == 'foo'
        assert s.s3writer.write() == 'bar'


def test_enter_opens_s3_file(monkeypatch):
    with patch('smart_open.smart_open') as MockClass:
        instance = MockClass.return_value
        instance.smart_open.return_value = 'open'
        instance.close.return_value = 'close'
        with StreamToS3(mock_key, mock_bucket, mock_dict_keys) as s:
            s.stream = instance  # so it will close w/o error
            assert s.s3writer is instance


def test_write_batch_rows(monkeypatch):
    mock_set_csv_writer(monkeypatch)
    fake_rowset = [
        {'one': 1, 'two': 2, 'three': 3},
        {'one': 1, 'two': 2, 'three': 3}]
    with patch('csv.DictWriter') as MockClass:
        instance = MockClass.return_value
        instance.writerow.return_value = 'foo'
        instance.write.return_value = 'bar'
        instance.getvalue.return_value = 'foobar'
        s = StreamToS3(mock_key, mock_bucket, mock_dict_keys)
        s.csv_writer = instance
        s.s3writer = instance
        s.stream = instance
        s.write_batch_rows(fake_rowset)
        # write 3 rows to stream, write once to s3
        assert instance.writerow.call_count == 2
        assert instance.write.call_count == 1