from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorboardX import SummaryWriter
import numpy as np
import pytest
import unittest
import tensorboardX.beholder as beholder_lib
import tensorboardX.beholder.file_system_tools as fio
from collections import namedtuple


class BeholderTest(unittest.TestCase):
    def test_beholder(self):
        LOG_DIRECTORY = '/tmp/beholder-demo'
        tensor_and_name = namedtuple('tensor_and_name', 'tensor, name')
        fake_param = [tensor_and_name(np.random.randn(128, 768, 3), 'test' + str(i)) for i in range(5)]
        arrays = [tensor_and_name(np.random.randn(128, 768, 3), 'test' + str(i)) for i in range(5)]
        beholder = beholder_lib.Beholder(logdir=LOG_DIRECTORY)
        beholder.update(
            trainable=fake_param,
            arrays=arrays,
            frame=np.random.randn(128, 128),
        )

    def test_beholder_video(self):
        LOG_DIRECTORY = '/tmp/beholder-demo-recording'
        tensor_and_name = namedtuple('tensor_and_name', 'tensor, name')
        fake_param = [tensor_and_name(np.random.randn(128, 768, 3), 'test' + str(i)) for i in range(5)]
        arrays = [tensor_and_name(np.random.randn(128, 768, 3), 'test' + str(i)) for i in range(5)]
        beholder = beholder_lib.Beholder(logdir=LOG_DIRECTORY)
        pkl = fio.read_pickle(LOG_DIRECTORY + '/plugins/beholder/config.pkl')
        pkl['is_recording'] = True
        fio.write_pickle(pkl, LOG_DIRECTORY + '/plugins/beholder/config.pkl')
        for i in range(3):
            if i == 2:
                pkl = fio.read_pickle(LOG_DIRECTORY + '/plugins/beholder/config.pkl')
                pkl['is_recording'] = False
                fio.write_pickle(pkl, LOG_DIRECTORY + '/plugins/beholder/config.pkl')
            beholder.update(
                trainable=fake_param,
                arrays=arrays,
                frame=np.random.randn(128, 128),
            )

@@ -107,47 +107,8 @@
 on%0A%0A
-from tensorboardX import SummaryWriter%0A
 impo
@@ -126,22 +126,8 @@
  np%0A
-import pytest%0A
 impo
@@ -237,16 +237,32 @@
  as fio%0A
+import tempfile%0A
 from col
@@ -288,16 +288,30 @@
 edtuple%0A
+import shutil%0A
 %0A%0Aclass 
@@ -355,67 +355,204 @@
 def 
-test_beholder(self):%0A        LOG_DIRECTORY = '/tmp/
+setUp(self):%0A        self.LOG_DIRECTORY = tempfile.mkdtemp()%0A%0A    def tearDown(self):%0A        print(%22removing%22, self.LOG_DIRECTORY)%0A        shutil.rmtree(self.LOG_DIRECTORY)%0A%0A    def test_
 beholder
 -dem
@@ -551,14 +551,15 @@
 lder
--demo'
+(self):
 %0A   
@@ -871,32 +871,37 @@
 Beholder(logdir=
+self.
 LOG_DIRECTORY)%0A 
@@ -1104,38 +1104,26 @@
 Y = 
-'/tmp/beholder-demo-recording'
+self.LOG_DIRECTORY
 %0A   
@@ -1443,16 +1443,21 @@
 (logdir=
+self.
 LOG_DIRE

fix flaky test