laywerrobot/lib/python3.6/site-packages/smart_open/tests/test_smart_open.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2015 Radim Rehurek <me@radimrehurek.com>
#
# This code is distributed under the terms and conditions
# from the MIT License (MIT).

import io
import unittest
import logging
import tempfile
import os
import sys
import hashlib

import boto3
import mock
from moto import mock_s3
import responses
import gzip
import six

import smart_open
from smart_open import smart_open_lib

logger = logging.getLogger(__name__)

PY2 = sys.version_info[0] == 2
CURR_DIR = os.path.abspath(os.path.dirname(__file__))


class ParseUriTest(unittest.TestCase):
    """
    Test ParseUri class.

    """
    def test_scheme(self):
        """Do URIs schemes parse correctly?"""
        # supported schemes
        for scheme in ("s3", "s3n", "hdfs", "file", "http", "https"):
            parsed_uri = smart_open_lib._parse_uri(scheme + "://mybucket/mykey")
            self.assertEqual(parsed_uri.scheme, scheme)

        # unsupported scheme => NotImplementedError
        self.assertRaises(NotImplementedError, smart_open_lib._parse_uri, "foobar://mybucket/mykey")

        # unknown scheme => default_scheme
        parsed_uri = smart_open_lib._parse_uri("blah blah")
        self.assertEqual(parsed_uri.scheme, "file")

    def test_s3_uri(self):
        """Do S3 URIs parse correctly?"""
        # correct uri without credentials
        parsed_uri = smart_open_lib._parse_uri("s3://mybucket/mykey")
        self.assertEqual(parsed_uri.scheme, "s3")
        self.assertEqual(parsed_uri.bucket_id, "mybucket")
        self.assertEqual(parsed_uri.key_id, "mykey")
        self.assertEqual(parsed_uri.access_id, None)
        self.assertEqual(parsed_uri.access_secret, None)

        # correct uri, key contains slash
        parsed_uri = smart_open_lib._parse_uri("s3://mybucket/mydir/mykey")
        self.assertEqual(parsed_uri.scheme, "s3")
        self.assertEqual(parsed_uri.bucket_id, "mybucket")
        self.assertEqual(parsed_uri.key_id, "mydir/mykey")
        self.assertEqual(parsed_uri.access_id, None)
        self.assertEqual(parsed_uri.access_secret, None)

        # correct uri with credentials
        parsed_uri = smart_open_lib._parse_uri("s3://ACCESSID456:acces/sse_cr-et@mybucket/mykey")
        self.assertEqual(parsed_uri.scheme, "s3")
        self.assertEqual(parsed_uri.bucket_id, "mybucket")
        self.assertEqual(parsed_uri.key_id, "mykey")
        self.assertEqual(parsed_uri.access_id, "ACCESSID456")
        self.assertEqual(parsed_uri.access_secret, "acces/sse_cr-et")

        # correct uri, contains credentials
        parsed_uri = smart_open_lib._parse_uri("s3://accessid:access/secret@mybucket/mykey")
        self.assertEqual(parsed_uri.scheme, "s3")
        self.assertEqual(parsed_uri.bucket_id, "mybucket")
        self.assertEqual(parsed_uri.key_id, "mykey")
        self.assertEqual(parsed_uri.access_id, "accessid")
        self.assertEqual(parsed_uri.access_secret, "access/secret")

        # incorrect uri - only two '@' in uri are allowed
        self.assertRaises(RuntimeError, smart_open_lib._parse_uri, "s3://access_id@access_secret@mybucket@port/mykey")

    def test_webhdfs_uri(self):
        """Do webhdfs URIs parse correctly"""
        # valid uri, no query
        parsed_uri = smart_open_lib._parse_uri("webhdfs://host:port/path/file")
        self.assertEqual(parsed_uri.scheme, "webhdfs")
        self.assertEqual(parsed_uri.uri_path, "host:port/webhdfs/v1/path/file")

        # valid uri, with query
        parsed_uri = smart_open_lib._parse_uri("webhdfs://host:port/path/file?query_part_1&query_part2")
        self.assertEqual(parsed_uri.scheme, "webhdfs")
        self.assertEqual(parsed_uri.uri_path, "host:port/webhdfs/v1/path/file?query_part_1&query_part2")


class SmartOpenHttpTest(unittest.TestCase):
    """
    Test reading from HTTP connections in various ways.

    """
    @responses.activate
    def test_http_read(self):
        """Does http read method work correctly"""
        responses.add(responses.GET, "http://127.0.0.1/index.html",
                      body='line1\nline2', stream=True)
        smart_open_object = smart_open.smart_open("http://127.0.0.1/index.html")
        self.assertEqual(smart_open_object.read().decode("utf-8"), "line1\nline2")

    @responses.activate
    def test_https_readline(self):
        """Does https readline method work correctly"""
        responses.add(responses.GET, "https://127.0.0.1/index.html",
                      body='line1\nline2', stream=True)
        smart_open_object = smart_open.smart_open("https://127.0.0.1/index.html")
        self.assertEqual(smart_open_object.readline().decode("utf-8"), "line1\n")

    @responses.activate
    def test_http_pass(self):
        """Does http authentication work correctly"""
        responses.add(responses.GET, "http://127.0.0.1/index.html",
                      body='line1\nline2', stream=True)
        _ = smart_open.smart_open("http://127.0.0.1/index.html", user='me', password='pass')
        self.assertEqual(len(responses.calls), 1)
        actual_request = responses.calls[0].request
        self.assertTrue('Authorization' in actual_request.headers)
        self.assertTrue(actual_request.headers['Authorization'].startswith('Basic '))

    @responses.activate
    @unittest.skipIf(six.PY2, 'gzip support for Py2 is not implemented yet')
    def test_http_gz(self):
        """Can open gzip via http?"""
        fpath = os.path.join(CURR_DIR, 'test_data/crlf_at_1k_boundary.warc.gz')
        with open(fpath, 'rb') as infile:
            data = infile.read()

        with gzip.GzipFile(fpath) as fin:
            expected_hash = hashlib.md5(fin.read()).hexdigest()

        responses.add(responses.GET, "http://127.0.0.1/data.gz", body=data, stream=True)
        smart_open_object = smart_open.smart_open("http://127.0.0.1/data.gz?some_param=some_val")

        m = hashlib.md5(smart_open_object.read())
        # decompress the gzip and get the same md5 hash
        self.assertEqual(m.hexdigest(), expected_hash)

    @responses.activate
    @unittest.skipIf(six.PY2, 'gzip support for Py2 is not implemented yet')
    def test_http_gz_noquerystring(self):
        """Can open gzip via http?"""
        fpath = os.path.join(CURR_DIR, 'test_data/crlf_at_1k_boundary.warc.gz')
        with open(fpath, 'rb') as infile:
            data = infile.read()

        with gzip.GzipFile(fpath) as fin:
            expected_hash = hashlib.md5(fin.read()).hexdigest()

        responses.add(responses.GET, "http://127.0.0.1/data.gz", body=data, stream=True)
        smart_open_object = smart_open.smart_open("http://127.0.0.1/data.gz")

        m = hashlib.md5(smart_open_object.read())
        # decompress the gzip and get the same md5 hash
        self.assertEqual(m.hexdigest(), expected_hash)

    @responses.activate
    def test_http_bz2(self):
        """Can open bz2 via http?"""
        test_string = b'Hello World Compressed.'
        #
        # TODO: why are these tests writing to temporary files?  We can do the
        # bz2 compression in memory.
        #
        with tempfile.NamedTemporaryFile('wb', suffix='.bz2', delete=False) as infile:
            test_file = infile.name

        with smart_open.smart_open(test_file, 'wb') as outfile:
            outfile.write(test_string)

        with open(test_file, 'rb') as infile:
            compressed_data = infile.read()

        if os.path.isfile(test_file):
            os.unlink(test_file)

        responses.add(responses.GET, "http://127.0.0.1/data.bz2",
                      body=compressed_data, stream=True)
        smart_open_object = smart_open.smart_open("http://127.0.0.1/data.bz2")

        # decompress the gzip and get the same md5 hash
        self.assertEqual(smart_open_object.read(), test_string)


class SmartOpenReadTest(unittest.TestCase):
    """
    Test reading from files under various schemes.

    """

    def test_shortcut(self):
        fpath = os.path.join(CURR_DIR, 'test_data/crime-and-punishment.txt')
        with mock.patch('io.open') as mock_open:
            smart_open.smart_open(fpath, 'r').read()
        mock_open.assert_called_with(fpath, 'r')

    def test_open_with_keywords(self):
        """This test captures Issue #142."""
        fpath = os.path.join(CURR_DIR, 'test_data/cp852.tsv.txt')
        with open(fpath, 'rb') as fin:
            expected = fin.read().decode('cp852')
        with smart_open.smart_open(fpath, encoding='cp852') as fin:
            actual = fin.read()
        self.assertEqual(expected, actual)

    def test_open_with_keywords_explicit_r(self):
        fpath = os.path.join(CURR_DIR, 'test_data/cp852.tsv.txt')
        with open(fpath, 'rb') as fin:
            expected = fin.read().decode('cp852')
        with smart_open.smart_open(fpath, mode='r', encoding='cp852') as fin:
            actual = fin.read()
        self.assertEqual(expected, actual)

    @unittest.skipUnless(
        smart_open_lib.PATHLIB_SUPPORT,
        "do not test pathlib support if pathlib or backport are not available")
    def test_open_and_read_pathlib_path(self):
        """If ``pathlib.Path`` is available we should be able to open and read."""
        from smart_open.smart_open_lib import pathlib

        fpath = os.path.join(CURR_DIR, 'test_data/cp852.tsv.txt')
        with open(fpath, 'rb') as fin:
            expected = fin.read().decode('cp852')
        with smart_open.smart_open(pathlib.Path(fpath), mode='r', encoding='cp852') as fin:
            actual = fin.read()
        self.assertEqual(expected, actual)

    @mock_s3
    def test_read_never_returns_none(self):
        """read should never return None."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')

        test_string = u"ветер по морю гуляет..."
        with smart_open.smart_open("s3://mybucket/mykey", "wb") as fout:
            fout.write(test_string.encode('utf8'))

        r = smart_open.smart_open("s3://mybucket/mykey", "rb")
        self.assertEqual(r.read(), test_string.encode("utf-8"))
        self.assertEqual(r.read(), b"")
        self.assertEqual(r.read(), b"")

    @mock_s3
    def test_readline(self):
        """Does readline() return the correct file content?"""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')
        test_string = u"hello žluťoučký world!\nhow are you?".encode('utf8')
        with smart_open.smart_open("s3://mybucket/mykey", "wb") as fout:
            fout.write(test_string)

        reader = smart_open.smart_open("s3://mybucket/mykey", "rb")
        self.assertEqual(reader.readline(), u"hello žluťoučký world!\n".encode("utf-8"))

    @mock_s3
    def test_readline_iter(self):
        """Does __iter__ return the correct file content?"""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')
        lines = [u"всем привет!\n", u"что нового?"]
        with smart_open.smart_open("s3://mybucket/mykey", "wb") as fout:
            fout.write("".join(lines).encode("utf-8"))

        reader = smart_open.smart_open("s3://mybucket/mykey", "rb")

        actual_lines = [l.decode("utf-8") for l in reader]
        self.assertEqual(2, len(actual_lines))
        self.assertEqual(lines[0], actual_lines[0])
        self.assertEqual(lines[1], actual_lines[1])

    @mock_s3
    def test_readline_eof(self):
        """Does readline() return empty string on EOF?"""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')
        with smart_open.smart_open("s3://mybucket/mykey", "wb"):
            pass

        reader = smart_open.smart_open("s3://mybucket/mykey", "rb")

        self.assertEqual(reader.readline(), b"")
        self.assertEqual(reader.readline(), b"")
        self.assertEqual(reader.readline(), b"")

    @mock_s3
    def test_s3_iter_lines(self):
        """Does s3_iter_lines give correct content?"""
        # create fake bucket and fake key
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')
        test_string = u"hello žluťoučký world!\nhow are you?".encode('utf8')
        with smart_open.smart_open("s3://mybucket/mykey", "wb") as fin:
            fin.write(test_string)

        # call s3_iter_lines and check output
        reader = smart_open.smart_open("s3://mybucket/mykey", "rb")
        output = list(reader)
        self.assertEqual(b''.join(output), test_string)

    # TODO: add more complex test for file://
    @mock.patch('io.open')
    def test_file(self, mock_smart_open):
        """Is file:// line iterator called correctly?"""
        prefix = "file://"
        full_path = '/tmp/test.txt'
        read_mode = "rb"
        smart_open_object = smart_open.smart_open(prefix+full_path, read_mode)
        smart_open_object.__iter__()
        # called with the correct path?
        mock_smart_open.assert_called_with(full_path, read_mode)

        full_path = '/tmp/test#hash##more.txt'
        read_mode = "rb"
        smart_open_object = smart_open.smart_open(prefix+full_path, read_mode)
        smart_open_object.__iter__()
        # called with the correct path?
        mock_smart_open.assert_called_with(full_path, read_mode)

        full_path = 'aa#aa'
        read_mode = "rb"
        smart_open_object = smart_open.smart_open(full_path, read_mode)
        smart_open_object.__iter__()
        # called with the correct path?
        mock_smart_open.assert_called_with(full_path, read_mode)

        short_path = "~/tmp/test.txt"
        full_path = os.path.expanduser(short_path)

        smart_open_object = smart_open.smart_open(prefix+short_path, read_mode, errors='strict')
        smart_open_object.__iter__()
        # called with the correct expanded path?
        mock_smart_open.assert_called_with(full_path, read_mode, errors='strict')

    # couldn't find any project for mocking up HDFS data
    # TODO: we want to test also a content of the files, not just fnc call params
    @mock.patch('smart_open.hdfs.subprocess')
    def test_hdfs(self, mock_subprocess):
        """Is HDFS line iterator called correctly?"""
        mock_subprocess.PIPE.return_value = "test"
        smart_open_object = smart_open.smart_open("hdfs:///tmp/test.txt")
        smart_open_object.__iter__()
        # called with the correct params?
        mock_subprocess.Popen.assert_called_with(["hdfs", "dfs", "-cat", "/tmp/test.txt"], stdout=mock_subprocess.PIPE)

        # second possibility of schema
        smart_open_object = smart_open.smart_open("hdfs://tmp/test.txt")
        smart_open_object.__iter__()
        mock_subprocess.Popen.assert_called_with(["hdfs", "dfs", "-cat", "/tmp/test.txt"], stdout=mock_subprocess.PIPE)

    @responses.activate
    def test_webhdfs(self):
        """Is webhdfs line iterator called correctly"""
        responses.add(responses.GET, "http://127.0.0.1:8440/webhdfs/v1/path/file",
                      body='line1\nline2', stream=True)
        smart_open_object = smart_open.smart_open("webhdfs://127.0.0.1:8440/path/file")
        iterator = iter(smart_open_object)
        self.assertEqual(next(iterator).decode("utf-8"), "line1\n")
        self.assertEqual(next(iterator).decode("utf-8"), "line2")

    @responses.activate
    def test_webhdfs_encoding(self):
        """Is HDFS line iterator called correctly?"""
        input_url = "webhdfs://127.0.0.1:8440/path/file"
        actual_url = 'http://127.0.0.1:8440/webhdfs/v1/path/file'
        text = u'не для меня прийдёт весна, не для меня дон разольётся'
        body = text.encode('utf-8')
        responses.add(responses.GET, actual_url, body=body, stream=True)

        actual = smart_open.smart_open(input_url, encoding='utf-8').read()
        self.assertEqual(text, actual)

    @responses.activate
    def test_webhdfs_read(self):
        """Does webhdfs read method work correctly"""
        responses.add(responses.GET, "http://127.0.0.1:8440/webhdfs/v1/path/file",
                      body='line1\nline2', stream=True)
        smart_open_object = smart_open.smart_open("webhdfs://127.0.0.1:8440/path/file")
        self.assertEqual(smart_open_object.read().decode("utf-8"), "line1\nline2")

    @mock_s3
    def test_s3_iter_moto(self):
        """Are S3 files iterated over correctly?"""
        # a list of strings to test with
        expected = [b"*" * 5 * 1024**2] + [b'0123456789'] * 1024 + [b"test"]

        # create fake bucket and fake key
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')

        with smart_open.smart_open("s3://mybucket/mykey", "wb", s3_min_part_size=5 * 1024**2) as fout:
            # write a single huge line (=full multipart upload)
            fout.write(expected[0] + b'\n')

            # write lots of small lines
            for lineno, line in enumerate(expected[1:-1]):
                fout.write(line + b'\n')

            # ...and write the last line too, no newline at the end
            fout.write(expected[-1])

        # connect to fake s3 and read from the fake key we filled above
        smart_open_object = smart_open.smart_open("s3://mybucket/mykey")
        output = [line.rstrip(b'\n') for line in smart_open_object]
        self.assertEqual(output, expected)

        # same thing but using a context manager
        with smart_open.smart_open("s3://mybucket/mykey") as smart_open_object:
            output = [line.rstrip(b'\n') for line in smart_open_object]
            self.assertEqual(output, expected)

    @mock_s3
    def test_s3_read_moto(self):
        """Are S3 files read correctly?"""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')

        # write some bogus key so we can check it below
        content = u"hello wořld\nhow are you?".encode('utf8')
        with smart_open.smart_open("s3://mybucket/mykey", "wb") as fout:
            fout.write(content)

        smart_open_object = smart_open.smart_open("s3://mybucket/mykey")
        self.assertEqual(content[:6], smart_open_object.read(6))
        self.assertEqual(content[6:14], smart_open_object.read(8))  # ř is 2 bytes

        self.assertEqual(content[14:], smart_open_object.read())  # read the rest

    @unittest.skip('seek functionality for S3 currently disabled because of Issue #152')
    @mock_s3
    def test_s3_seek_moto(self):
        """Does seeking in S3 files work correctly?"""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')

        # write some bogus key so we can check it below
        content = u"hello wořld\nhow are you?".encode('utf8')
        with smart_open.smart_open("s3://mybucket/mykey", "wb") as fout:
            fout.write(content)

        smart_open_object = smart_open.smart_open("s3://mybucket/mykey")
        self.assertEqual(content[:6], smart_open_object.read(6))
        self.assertEqual(content[6:14], smart_open_object.read(8))  # ř is 2 bytes

        smart_open_object.seek(0)
        self.assertEqual(content, smart_open_object.read())  # no size given => read whole file

        smart_open_object.seek(0)
        self.assertEqual(content, smart_open_object.read(-1))  # same thing


class SmartOpenS3KwargsTest(unittest.TestCase):
    @mock.patch('boto3.Session')
    def test_no_kwargs(self, mock_session):
        smart_open.smart_open('s3://mybucket/mykey')
        mock_session.assert_called_with(profile_name=None)
        mock_session.return_value.resource.assert_called_with('s3')

    @mock.patch('boto3.Session')
    def test_credentials(self, mock_session):
        smart_open.smart_open('s3://access_id:access_secret@mybucket/mykey')
        mock_session.assert_called_with(profile_name=None)
        mock_session.return_value.resource.assert_called_with(
            's3', aws_access_key_id='access_id', aws_secret_access_key='access_secret'
        )

    @mock.patch('boto3.Session')
    def test_profile(self, mock_session):
        smart_open.smart_open('s3://mybucket/mykey', profile_name='my_credentials')
        mock_session.assert_called_with(profile_name='my_credentials')
        mock_session.return_value.resource.assert_called_with('s3')

    @mock.patch('boto3.Session')
    def test_host(self, mock_session):
        smart_open.smart_open("s3://access_id:access_secret@mybucket/mykey", host='aa.domain.com')
        mock_session.return_value.resource.assert_called_with(
            's3', aws_access_key_id='access_id', aws_secret_access_key='access_secret',
            endpoint_url='http://aa.domain.com'
        )

    @mock.patch('boto3.Session')
    def test_s3_upload(self, mock_session):
        smart_open.smart_open("s3://bucket/key", 'wb', s3_upload={
            'ServerSideEncryption': 'AES256',
            'ContentType': 'application/json'
        })

        # Locate the s3.Object instance (mock)
        s3_resource = mock_session.return_value.resource.return_value
        s3_object = s3_resource.Object.return_value

        # Check that `initiate_multipart_upload` was called
        # with the desired args
        s3_object.initiate_multipart_upload.assert_called_with(
            ServerSideEncryption='AES256',
            ContentType='application/json'
        )

    def test_session_read_mode(self):
        """
        Read stream should use a custom boto3.Session
        """
        session = boto3.Session()
        session.resource = mock.MagicMock()

        smart_open.smart_open('s3://bucket/key', s3_session=session)
        session.resource.assert_called_with('s3')

    def test_session_write_mode(self):
        """
        Write stream should use a custom boto3.Session
        """
        session = boto3.Session()
        session.resource = mock.MagicMock()

        smart_open.smart_open('s3://bucket/key', 'wb', s3_session=session)
        session.resource.assert_called_with('s3')


class SmartOpenTest(unittest.TestCase):
    """
    Test reading and writing from/into files.

    """
    @mock.patch('smart_open.smart_open_lib.boto')
    def test_file_mode_mock(self, mock_boto):
        """Are file:// open modes passed correctly?"""
        as_text = u'куда идём мы с пятачком - большой большой секрет'
        as_bytes = as_text.encode('utf-8')

        # incorrect file mode
        self.assertRaises(
            NotImplementedError, smart_open.smart_open, "s3://bucket/key", "x"
        )

        # correct read modes
        #
        # We always open files in binary mode first, but engage
        # encoders/decoders as necessary.  Instead of checking how the file
        # _initially_ got opened, we now also check the end result: if the
        # contents got decoded correctly.
        #
        with mock.patch('io.open', mock.Mock(return_value=io.StringIO(as_text))) as mock_open:
            with smart_open.smart_open("blah", "r", encoding='utf-8') as fin:
                self.assertEqual(fin.read(), as_text)
                mock_open.assert_called_with("blah", "r", encoding='utf-8')

        with mock.patch('io.open', mock.Mock(return_value=io.BytesIO(as_bytes))) as mock_open:
            with smart_open.smart_open("blah", "rb") as fin:
                self.assertEqual(fin.read(), as_bytes)
                mock_open.assert_called_with("blah", "rb")

        short_path = "~/blah"
        full_path = os.path.expanduser(short_path)
        with mock.patch('io.open', mock.Mock(return_value=io.BytesIO(as_bytes))) as mock_open:
            with smart_open.smart_open(short_path, "rb") as fin:
                mock_open.assert_called_with(full_path, "rb")

        # correct write modes, incorrect scheme
        self.assertRaises(NotImplementedError, smart_open.smart_open, "hdfs:///blah.txt", "wb+")
        self.assertRaises(NotImplementedError, smart_open.smart_open, "http:///blah.txt", "w")
        self.assertRaises(NotImplementedError, smart_open.smart_open, "s3://bucket/key", "wb+")

        # correct write mode, correct file:// URI
        with mock.patch('io.open', mock.Mock(return_value=io.StringIO(as_text))) as mock_open:
            with smart_open.smart_open("blah", "w", encoding='utf-8') as fout:
                mock_open.assert_called_with("blah", "w", encoding='utf-8')
                fout.write(as_text)

        with mock.patch('io.open', mock.Mock(return_value=io.StringIO(as_text))) as mock_open:
            with smart_open.smart_open("/some/file.txt", "w", encoding='utf-8') as fout:
                mock_open.assert_called_with("/some/file.txt", "w", encoding='utf-8')
                fout.write(as_text)

        with mock.patch('io.open', mock.Mock(return_value=io.StringIO(as_text))) as mock_open:
            with smart_open.smart_open("/some/file.txt", "w+", encoding='utf-8') as fout:
                mock_open.assert_called_with("/some/file.txt", "w+", encoding='utf-8')
                fout.write(as_text)

        with mock.patch('io.open', mock.Mock(return_value=io.BytesIO(as_bytes))) as mock_open:
            with smart_open.smart_open("/some/file.txt", "wb+") as fout:
                mock_open.assert_called_with("/some/file.txt", "wb+")
                fout.write(as_bytes)

    @mock.patch('boto3.Session')
    def test_s3_mode_mock(self, mock_session):
        """Are s3:// open modes passed correctly?"""

        # correct write mode, correct s3 URI
        smart_open.smart_open("s3://mybucket/mykey", "w", host='s3.amazonaws.com')
        mock_session.return_value.resource.assert_called_with(
            's3', endpoint_url='http://s3.amazonaws.com'
        )

    @mock.patch('smart_open.hdfs.subprocess')
    def test_hdfs(self, mock_subprocess):
        """Is HDFS write called correctly"""
        smart_open_object = smart_open.smart_open("hdfs:///tmp/test.txt", 'wb')
        smart_open_object.write("test")
        # called with the correct params?
        mock_subprocess.Popen.assert_called_with(
            ["hdfs", "dfs", "-put", "-f", "-", "/tmp/test.txt"], stdin=mock_subprocess.PIPE
        )

        # second possibility of schema
        smart_open_object = smart_open.smart_open("hdfs://tmp/test.txt", 'wb')
        smart_open_object.write("test")
        mock_subprocess.Popen.assert_called_with(
            ["hdfs", "dfs", "-put", "-f", "-", "/tmp/test.txt"], stdin=mock_subprocess.PIPE
        )

    @mock_s3
    def test_s3_modes_moto(self):
        """Do s3:// open modes work correctly?"""
        # fake bucket and key
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')
        test_string = b"second test"

        # correct write mode, correct s3 URI
        with smart_open.smart_open("s3://mybucket/newkey", "wb") as fout:
            logger.debug('fout: %r', fout)
            fout.write(test_string)

        logger.debug("write successfully completed")

        output = list(smart_open.smart_open("s3://mybucket/newkey", "rb"))

        self.assertEqual(output, [test_string])

    @mock_s3
    def test_s3_metadata_write(self):
        # Read local file fixture
        path = os.path.join(CURR_DIR, 'test_data/crime-and-punishment.txt.gz')
        data = ""
        with smart_open.smart_open(path, 'rb') as fd:
            data = fd.read()

        # Create a test bucket
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='mybucket')

        # Write data, with multipart_upload options
        write_stream = smart_open.smart_open(
            's3://mybucket/crime-and-punishment.txt.gz', 'wb',
            s3_upload={
                'ContentType': 'text/plain',
                'ContentEncoding': 'gzip'
            }
        )
        with write_stream as fout:
            fout.write(data)

        key = s3.Object('mybucket', 'crime-and-punishment.txt.gz')
        self.assertIn('text/plain', key.content_type)
        self.assertEqual(key.content_encoding, 'gzip')

    @mock_s3
    def test_write_bad_encoding_strict(self):
        """Should abort on encoding error."""
        text = u'欲しい気持ちが成長しすぎて'

        with self.assertRaises(UnicodeEncodeError):
            with tempfile.NamedTemporaryFile('wb', delete=True) as infile:
                with smart_open.smart_open(infile.name, 'w', encoding='koi8-r',
                                           errors='strict') as fout:
                    fout.write(text)

    @mock_s3
    def test_write_bad_encoding_replace(self):
        """Should replace characters that failed to encode."""
        text = u'欲しい気持ちが成長しすぎて'
        expected = u'?' * len(text)

        with tempfile.NamedTemporaryFile('wb', delete=True) as infile:
            with smart_open.smart_open(infile.name, 'w', encoding='koi8-r',
                                       errors='replace') as fout:
                fout.write(text)
            with smart_open.smart_open(infile.name, 'r', encoding='koi8-r') as fin:
                actual = fin.read()

        self.assertEqual(expected, actual)


class WebHdfsWriteTest(unittest.TestCase):
    """
    Test writing into webhdfs files.

    """

    @responses.activate
    def test_initialize_write(self):
        def request_callback(_):
            resp_body = ""
            headers = {'location': 'http://127.0.0.1:8440/file'}
            return 307, headers, resp_body

        responses.add_callback(responses.PUT, "http://127.0.0.1:8440/webhdfs/v1/path/file", callback=request_callback)
        responses.add(responses.PUT, "http://127.0.0.1:8440/file", status=201)
        smart_open.smart_open("webhdfs://127.0.0.1:8440/path/file", 'wb')

        assert len(responses.calls) == 2
        path, params = responses.calls[0].request.url.split("?")
        assert path == "http://127.0.0.1:8440/webhdfs/v1/path/file"
        assert params == "overwrite=True&op=CREATE" or params == "op=CREATE&overwrite=True"
        assert responses.calls[1].request.url == "http://127.0.0.1:8440/file"

    @responses.activate
    def test_write(self):
        def request_callback(_):
            resp_body = ""
            headers = {'location': 'http://127.0.0.1:8440/file'}
            return 307, headers, resp_body

        responses.add_callback(responses.PUT, "http://127.0.0.1:8440/webhdfs/v1/path/file", callback=request_callback)
        responses.add(responses.PUT, "http://127.0.0.1:8440/file", status=201)
        smart_open_object = smart_open.smart_open("webhdfs://127.0.0.1:8440/path/file", 'wb')

        def write_callback(request):
            assert request.body == u"žluťoučký koníček".encode('utf8')
            headers = {}
            return 200, headers, ""

        test_string = u"žluťoučký koníček".encode('utf8')
        responses.add_callback(responses.POST, "http://127.0.0.1:8440/webhdfs/v1/path/file", callback=request_callback)
        responses.add_callback(responses.POST, "http://127.0.0.1:8440/file", callback=write_callback)
        smart_open_object.write(test_string)
        smart_open_object.close()

        assert len(responses.calls) == 4
        assert responses.calls[2].request.url == "http://127.0.0.1:8440/webhdfs/v1/path/file?op=APPEND"
        assert responses.calls[3].request.url == "http://127.0.0.1:8440/file"


class CompressionFormatTest(unittest.TestCase):
    """
    Test that compression
    """

    TEXT = 'Hello'

    def write_read_assertion(self, test_file):
        with smart_open.smart_open(test_file, 'wb') as fout:  # 'b' for binary, needed on Windows
            fout.write(self.TEXT.encode('utf8'))

        with smart_open.smart_open(test_file, 'rb') as fin:
            self.assertEqual(fin.read().decode('utf8'), self.TEXT)

        if os.path.isfile(test_file):
            os.unlink(test_file)

    def test_open_gz(self):
        """Can open gzip?"""
        fpath = os.path.join(CURR_DIR, 'test_data/crlf_at_1k_boundary.warc.gz')
        with smart_open.smart_open(fpath) as infile:
            data = infile.read()
        m = hashlib.md5(data)
        assert m.hexdigest() == '18473e60f8c7c98d29d65bf805736a0d', \
            'Failed to read gzip'

    def test_write_read_gz(self):
        """Can write and read gzip?"""
        with tempfile.NamedTemporaryFile('wb', suffix='.gz', delete=False) as infile:
            test_file_name = infile.name
        self.write_read_assertion(test_file_name)

    def test_write_read_bz2(self):
        """Can write and read bz2?"""
        with tempfile.NamedTemporaryFile('wb', suffix='.bz2', delete=False) as infile:
            test_file_name = infile.name
        self.write_read_assertion(test_file_name)


class MultistreamsBZ2Test(unittest.TestCase):
    """
    Test that multistream bzip2 compressed files can be read.

    """

    # note: these tests are derived from the Python 3.x tip bz2 tests.

    TEXT_LINES = [
        b'root:x:0:0:root:/root:/bin/bash\n',
        b'bin:x:1:1:bin:/bin:\n',
        b'daemon:x:2:2:daemon:/sbin:\n',
        b'adm:x:3:4:adm:/var/adm:\n',
        b'lp:x:4:7:lp:/var/spool/lpd:\n',
        b'sync:x:5:0:sync:/sbin:/bin/sync\n',
        b'shutdown:x:6:0:shutdown:/sbin:/sbin/shutdown\n',
        b'halt:x:7:0:halt:/sbin:/sbin/halt\n',
        b'mail:x:8:12:mail:/var/spool/mail:\n',
        b'news:x:9:13:news:/var/spool/news:\n',
        b'uucp:x:10:14:uucp:/var/spool/uucp:\n',
        b'operator:x:11:0:operator:/root:\n',
        b'games:x:12:100:games:/usr/games:\n',
        b'gopher:x:13:30:gopher:/usr/lib/gopher-data:\n',
        b'ftp:x:14:50:FTP User:/var/ftp:/bin/bash\n',
        b'nobody:x:65534:65534:Nobody:/home:\n',
        b'postfix:x:100:101:postfix:/var/spool/postfix:\n',
        b'niemeyer:x:500:500::/home/niemeyer:/bin/bash\n',
        b'postgres:x:101:102:PostgreSQL Server:/var/lib/pgsql:/bin/bash\n',
        b'mysql:x:102:103:MySQL server:/var/lib/mysql:/bin/bash\n',
        b'www:x:103:104::/var/www:/bin/false\n',
    ]

    TEXT = b''.join(TEXT_LINES)

    DATA = \
        b'BZh91AY&SY.\xc8N\x18\x00\x01>_\x80\x00\x10@\x02\xff\xf0\x01\x07n\x00?\xe7\xff\xe00\x01\x99\xaa\x00' \
        b'\xc0\x03F\x86\x8c#&\x83F\x9a\x03\x06\xa6\xd0\xa6\x93M\x0fQ\xa7\xa8\x06\x804hh\x12$\x11\xa4i4\xf14S' \
        b'\xd2<Q\xb5\x0fH\xd3\xd4\xdd\xd5\x87\xbb\xf8\x94\r\x8f\xafI\x12\xe1\xc9\xf8/E\x00pu\x89\x12]\xc9' \
        b'\xbbDL\nQ\x0e\t1\x12\xdf\xa0\xc0\x97\xac2O9\x89\x13\x94\x0e\x1c7\x0ed\x95I\x0c\xaaJ\xa4\x18L\x10' \
        b'\x05#\x9c\xaf\xba\xbc/\x97\x8a#C\xc8\xe1\x8cW\xf9\xe2\xd0\xd6M\xa7\x8bXa<e\x84t\xcbL\xb3\xa7\xd9' \
        b'\xcd\xd1\xcb\x84.\xaf\xb3\xab\xab\xad`n}\xa0lh\tE,\x8eZ\x15\x17VH>\x88\xe5\xcd9gd6\x0b\n\xe9\x9b' \
        b'\xd5\x8a\x99\xf7\x08.K\x8ev\xfb\xf7xw\xbb\xdf\xa1\x92\xf1\xdd|/";\xa2\xba\x9f\xd5\xb1#A\xb6\xf6' \
        b'\xb3o\xc9\xc5y\\\xebO\xe7\x85\x9a\xbc\xb6f8\x952\xd5\xd7"%\x89>V,\xf7\xa6z\xe2\x9f\xa3\xdf\x11' \
        b'\x11"\xd6E)I\xa9\x13^\xca\xf3r\xd0\x03U\x922\xf26\xec\xb6\xed\x8b\xc3U\x13\x9d\xc5\x170\xa4\xfa^' \
        b'\x92\xacDF\x8a\x97\xd6\x19\xfe\xdd\xb8\xbd\x1a\x9a\x19\xa3\x80ankR\x8b\xe5\xd83]\xa9\xc6\x08' \
        b'\x82f\xf6\xb9"6l$\xb8j@\xc0\x8a\xb0l1..\xbak\x83ls\x15\xbc\xf4\xc1\x13\xbe\xf8E\xb8\x9d\r\xa8\x9dk' \
        b'\x84\xd3n\xfa\xacQ\x07\xb1%y\xaav\xb4\x08\xe0z\x1b\x16\xf5\x04\xe9\xcc\xb9\x08z\x1en7.G\xfc]\xc9\x14' \
        b'\xe1B@\xbb!8`'

    def create_temp_bz2(self, streams=1):
        f = tempfile.NamedTemporaryFile('wb', suffix='.bz2', delete=False)
        f.write(self.DATA * streams)
        f.close()
        return f.name

    def cleanup_temp_bz2(self, test_file):
        if os.path.isfile(test_file):
            os.unlink(test_file)

    def test_can_read_multistream_bz2(self):
        if PY2:
            # this is a backport from Python 3
            from bz2file import BZ2File
        else:
            from bz2 import BZ2File

        test_file = self.create_temp_bz2(streams=5)
        with BZ2File(test_file) as bz2f:
            self.assertEqual(bz2f.read(), self.TEXT * 5)
        self.cleanup_temp_bz2(test_file)

    def test_python2_stdlib_bz2_cannot_read_multistream(self):
        # Multistream bzip is included in Python 3
        if not PY2:
            return
        import bz2

        test_file = self.create_temp_bz2(streams=5)
        bz2f = bz2.BZ2File(test_file)
        self.assertNotEqual(bz2f.read(), self.TEXT * 5)
        bz2f.close()
        self.cleanup_temp_bz2(test_file)

    def test_file_smart_open_can_read_multistream_bz2(self):
        test_file = self.create_temp_bz2(streams=5)
        with smart_open_lib.smart_open(test_file) as bz2f:
            self.assertEqual(bz2f.read(), self.TEXT * 5)
        self.cleanup_temp_bz2(test_file)


class S3OpenTest(unittest.TestCase):

    @mock_s3
    def test_r(self):
        """Reading a UTF string should work."""
        text = u"физкульт-привет!"

        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = s3.Object('bucket', 'key')
        key.put(Body=text.encode('utf-8'))

        with smart_open.smart_open('s3://bucket/key', "rb") as fin:
            self.assertEqual(fin.read(), text.encode('utf-8'))

        with smart_open.smart_open('s3://bucket/key', "r", encoding='utf-8') as fin:
            self.assertEqual(fin.read(), text)

    def test_bad_mode(self):
        """Bad mode should raise and exception."""
        uri = smart_open_lib._parse_uri("s3://bucket/key")
        self.assertRaises(NotImplementedError, smart_open.smart_open, uri, "x")

    @mock_s3
    def test_rw_encoding(self):
        """Should read and write text, respecting encodings, etc."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')

        key = "s3://bucket/key"
        text = u"расцветали яблони и груши"

        with smart_open.smart_open(key, "w", encoding="koi8-r") as fout:
            fout.write(text)

        with smart_open.smart_open(key, "r", encoding="koi8-r") as fin:
            self.assertEqual(text, fin.read())

        with smart_open.smart_open(key, "rb") as fin:
            self.assertEqual(text.encode("koi8-r"), fin.read())

        with smart_open.smart_open(key, "r", encoding="euc-jp") as fin:
            self.assertRaises(UnicodeDecodeError, fin.read)

        with smart_open.smart_open(key, "r", encoding="euc-jp", errors="replace") as fin:
            fin.read()

    @mock_s3
    def test_rw_gzip(self):
        """Should read/write gzip files, implicitly and explicitly."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = "s3://bucket/key.gz"

        text = u"не слышны в саду даже шорохи"
        with smart_open.smart_open(key, "wb") as fout:
            fout.write(text.encode("utf-8"))

        #
        # Check that what we've created is a gzip.
        #
        with smart_open.smart_open(key, "rb", ignore_extension=True) as fin:
            gz = gzip.GzipFile(fileobj=fin)
            self.assertEqual(gz.read().decode("utf-8"), text)

        #
        # We should be able to read it back as well.
        #
        with smart_open.smart_open(key, "rb") as fin:
            self.assertEqual(fin.read().decode("utf-8"), text)

    @mock_s3
    def test_gzip_write_mode(self):
        """Should always open in binary mode when writing through a codec."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        uri = smart_open_lib._parse_uri("s3://bucket/key.gz")

        with mock.patch('smart_open.smart_open_s3.open') as mock_open:
            smart_open.smart_open("s3://bucket/key.gz", "wb")
            mock_open.assert_called_with('bucket', 'key.gz', 'wb')

    @mock_s3
    def test_gzip_read_mode(self):
        """Should always open in binary mode when reading through a codec."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = "s3://bucket/key.gz"

        text = u"если-б я был султан и имел трёх жён, то тройной красотой был бы окружён"
        with smart_open.smart_open(key, "wb") as fout:
            fout.write(text.encode("utf-8"))

        with mock.patch('smart_open.smart_open_s3.open') as mock_open:
            smart_open.smart_open(key, "r")
            mock_open.assert_called_with('bucket', 'key.gz', 'rb')

    @mock_s3
    def test_read_encoding(self):
        """Should open the file with the correct encoding, explicit text read."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = "s3://bucket/key.txt"
        text = u'это знала ева, это знал адам, колеса любви едут прямо по нам'
        with smart_open.smart_open(key, 'wb') as fout:
            fout.write(text.encode('koi8-r'))
        with smart_open.smart_open(key, 'r', encoding='koi8-r') as fin:
            actual = fin.read()
        self.assertEqual(text, actual)

    @mock_s3
    def test_read_encoding_implicit_text(self):
        """Should open the file with the correct encoding, implicit text read."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = "s3://bucket/key.txt"
        text = u'это знала ева, это знал адам, колеса любви едут прямо по нам'
        with smart_open.smart_open(key, 'wb') as fout:
            fout.write(text.encode('koi8-r'))
        with smart_open.smart_open(key, encoding='koi8-r') as fin:
            actual = fin.read()
        self.assertEqual(text, actual)

    @mock_s3
    def test_write_encoding(self):
        """Should open the file for writing with the correct encoding."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = "s3://bucket/key.txt"
        text = u'какая боль, какая боль, аргентина - ямайка, 5-0'

        with smart_open.smart_open(key, 'w', encoding='koi8-r') as fout:
            fout.write(text)
        with smart_open.smart_open(key, encoding='koi8-r') as fin:
            actual = fin.read()
        self.assertEqual(text, actual)

    @mock_s3
    def test_write_bad_encoding_strict(self):
        """Should open the file for writing with the correct encoding."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = "s3://bucket/key.txt"
        text = u'欲しい気持ちが成長しすぎて'

        with self.assertRaises(UnicodeEncodeError):
            with smart_open.smart_open(key, 'w', encoding='koi8-r', errors='strict') as fout:
                fout.write(text)

    @mock_s3
    def test_write_bad_encoding_replace(self):
        """Should open the file for writing with the correct encoding."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = "s3://bucket/key.txt"
        text = u'欲しい気持ちが成長しすぎて'
        expected = u'?' * len(text)

        with smart_open.smart_open(key, 'w', encoding='koi8-r', errors='replace') as fout:
            fout.write(text)
        with smart_open.smart_open(key, encoding='koi8-r') as fin:
            actual = fin.read()
        self.assertEqual(expected, actual)

    @mock_s3
    def test_write_text_gzip(self):
        """Should open the file for writing with the correct encoding."""
        s3 = boto3.resource('s3')
        s3.create_bucket(Bucket='bucket')
        key = "s3://bucket/key.txt.gz"
        text = u'какая боль, какая боль, аргентина - ямайка, 5-0'

        with smart_open.smart_open(key, 'w', encoding='utf-8') as fout:
            fout.write(text)
        with smart_open.smart_open(key, 'r', encoding='utf-8') as fin:
            actual = fin.read()
        self.assertEqual(text, actual)


if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.DEBUG)
    unittest.main()