test_record_writer.py

from tensorboardX import SummaryWriter
import unittest
from tensorboardX.record_writer import S3RecordWriter, make_valid_tf_name, GCSRecordWriter
import os
import boto3
from moto import mock_s3

os.environ.setdefault("AWS_ACCESS_KEY_ID", "foobar_key")
os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "foobar_secret")


class RecordWriterTest(unittest.TestCase):
    @mock_s3
    def test_record_writer_s3(self):
        client = boto3.client('s3', region_name='us-east-1')
        client.create_bucket(Bucket='this')
        writer = S3RecordWriter('s3://this/is/apen')
        bucket, path = writer.bucket_and_path()
        assert bucket == 'this'
        assert path == 'is/apen'
        writer.write(bytes(42))
        writer.flush()

    def test_make_valid_tf_name(self):
        newname = make_valid_tf_name('$ave/&sound')
        assert newname == '._ave/_sound'

    def test_record_writer_gcs(self):
        pass
        # we don't have mock test, so expect error here. However,
        # Travis CI env won't raise exception for the following code,
        # so I commented it out.
        # with self.assertRaises(Exception):
        #   writer = GCSRecordWriter('gs://this/is/apen')
        #   writer.write(bytes(42))
        #   writer.flush()