salt/tests/unit/tokens/test_localfs.py

# -*- coding: utf-8 -*-
'''
Tests the localfs tokens interface.
'''
from __future__ import absolute_import, print_function, unicode_literals

import os

import salt.exceptions
import salt.tokens.localfs
import salt.utils.files

from tests.support.unit import TestCase, skipIf
from tests.support.helpers import with_tempdir
from tests.support.mock import NO_MOCK, NO_MOCK_REASON, patch


class CalledWith(object):

    def __init__(self, func, called_with=None):
        self.func = func
        if called_with is None:
            self.called_with = []
        else:
            self.called_with = called_with

    def __call__(self, *args, **kwargs):
        self.called_with.append((args, kwargs))
        return self.func(*args, **kwargs)


@skipIf(NO_MOCK, NO_MOCK_REASON)
class WriteTokenTest(TestCase):

    @with_tempdir()
    def test_write_token(self, tmpdir):
        '''
        Validate tokens put in place with an atomic move
        '''
        opts = {
            'token_dir': tmpdir
        }
        fopen = CalledWith(salt.utils.files.fopen)
        rename = CalledWith(os.rename)
        with patch('salt.utils.files.fopen', fopen), patch('os.rename', rename):
            tdata = salt.tokens.localfs.mk_token(opts, {})
        assert 'token' in tdata
        t_path = os.path.join(tmpdir, tdata['token'])
        temp_t_path = '{}.tmp'.format(t_path)
        assert len(fopen.called_with) == 1, len(fopen.called_with)
        assert fopen.called_with == [
            ((temp_t_path, 'w+b'), {})
        ], fopen.called_with
        assert len(rename.called_with) == 1, len(rename.called_with)
        assert rename.called_with == [
            ((temp_t_path, t_path), {})
        ], rename.called_with


class TestLocalFS(TestCase):
    def setUp(self):
        # Default expected data
        self.expected_data = {'this': 'is', 'some': 'token data'}

    @with_tempdir()
    def test_get_token_should_return_token_if_exists(self, tempdir):
        opts = {'token_dir': tempdir}
        tok = salt.tokens.localfs.mk_token(
            opts=opts,
            tdata=self.expected_data,
        )['token']
        actual_data = salt.tokens.localfs.get_token(opts=opts, tok=tok)
        self.assertDictEqual(self.expected_data, actual_data)

    @with_tempdir()
    def test_get_token_should_raise_SaltDeserializationError_if_token_file_is_empty(self, tempdir):
        opts = {'token_dir': tempdir}
        tok = salt.tokens.localfs.mk_token(
            opts=opts,
            tdata=self.expected_data,
        )['token']
        with salt.utils.files.fopen(os.path.join(tempdir, tok), 'w') as f:
            f.truncate()
        with self.assertRaises(salt.exceptions.SaltDeserializationError) as e:
            salt.tokens.localfs.get_token(opts=opts, tok=tok)

    @with_tempdir()
    def test_get_token_should_raise_SaltDeserializationError_if_token_file_is_malformed(self, tempdir):
        opts = {'token_dir': tempdir}
        tok = salt.tokens.localfs.mk_token(
            opts=opts,
            tdata=self.expected_data,
        )['token']
        with salt.utils.files.fopen(os.path.join(tempdir, tok), 'w') as f:
            f.truncate()
            f.write('this is not valid msgpack data')
        with self.assertRaises(salt.exceptions.SaltDeserializationError) as e:
            salt.tokens.localfs.get_token(opts=opts, tok=tok)