#
# This file is licensed under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from test_support.testconfig import TestConfig
from smv import SmvApp

import pyspark
from pyspark.context import SparkContext
from pyspark.sql import *

import os, shutil

class SmvBaseTest(unittest.TestCase):
    # DataDir value is deprecated. Use tmpDataDir instead
    DataDir = "./target/data"
    PytestDir = "./target/pytest"

    @classmethod
    def smvAppInitArgs(cls):
        return ['-m', 'None']

    @classmethod
    def setUpClass(cls):
        cls.sparkContext = TestConfig.sparkContext()
        cls.sqlContext = TestConfig.sqlContext()
        cls.sparkContext.setLogLevel("ERROR")

        import random;
        callback_server_port = random.randint(20000, 65535)

        args = TestConfig.smv_args() + cls.smvAppInitArgs() + ['--cbs-port', str(callback_server_port), '--data-dir', cls.tmpDataDir()]
        cls.smvApp = SmvApp.createInstance(args, cls.sparkContext, cls.sqlContext)

        cls.mkTmpTestDir()

    def setUp(self):
        """Patch for Python 2.6 without using unittest
        """
        cls = self.__class__
        if not hasattr(cls, 'smvApp'):
            cls.sparkContext = TestConfig.sparkContext()
            cls.sqlContext = TestConfig.sqlContext()
            cls.sparkContext.setLogLevel("ERROR")

            import random;
            callback_server_port = random.randint(20000, 65535)

            args = TestConfig.smv_args() + cls.smvAppInitArgs() + ['--cbs-port', str(callback_server_port)]
            cls.smvApp = SmvApp.createInstance(args, cls.sparkContext, cls.sqlContext)

    @classmethod
    def createDF(cls, schema, data):
        return cls.smvApp.createDF(schema, data)

    @classmethod
    def df(cls, fqn):
        return cls.smvApp.runModule("mod:" + fqn)

    def should_be_same(self, expected, result):
        """Asserts that the two dataframes contain the same data, ignoring order
        """

        # Since Python sort can't handle null values in DF, use DF's orderBy to sort
        def sort_collect(df):
            return df.coalesce(1).orderBy(*(df.columns)).collect()

        self.assertEqual(expected.columns, result.columns)
        self.assertEqual(sort_collect(expected), sort_collect(result))

    @classmethod
    def tmpTestDir(cls):
        return cls.PytestDir + "/" + cls.__name__

    @classmethod
    def tmpDataDir(cls):
        return cls.tmpTestDir() + "/data"

    @classmethod
    def tmpInputDir(cls):
        return cls.tmpDataDir() + "/input"

    @classmethod
    def mkTmpTestDir(cls):
        shutil.rmtree(cls.tmpTestDir(), ignore_errors=True)
        os.makedirs(cls.tmpTestDir())

    def createTempInputFile(self, baseName, fileContents = "xxx"):
        """create a temp file in the input data dir with the given contents"""
        import os
        fullPath = self.tmpInputDir() + "/" + baseName
        directory = os.path.dirname(fullPath)
        if not os.path.exists(directory):
            os.makedirs(directory)

        f = open(fullPath, "w")
        f.write(fileContents)
        f.close()

@@ -2812,178 +2812,492 @@
     
-return cls.PytestDir + %22/%22 + cls.__name__%0A%0A    @classmethod%0A    def tmpDataDir(cls):%0A        return cls.tmpTestDir() + %22/data%22%0A%0A    @classmethod%0A    def tmpInputDir(cls):
+%22%22%22Temporary directory for each test to put the files it creates. Automatically cleaned up.%22%22%22%0A        return cls.PytestDir + %22/%22 + cls.__name__%0A%0A    @classmethod%0A    def tmpDataDir(cls):%0A        %22%22%22Temporary directory for each test to put the data it creates. Automatically cleaned up.%22%22%22%0A        return cls.tmpTestDir() + %22/data%22%0A%0A    @classmethod%0A    def tmpInputDir(cls):%0A        %22%22%22Temporary directory for each test to put the input files it creates. Automatically cleaned up.%22%22%22
 %0A   

add docstring to SmvBaseTest tmp directory methods