<commit_before>
import os
import time

import aiohttp
import pytest

from peony import PeonyClient, exceptions, oauth

oauth2_keys = 'PEONY_CONSUMER_KEY', 'PEONY_CONSUMER_SECRET'

oauth1_keys = *oauth2_keys, 'PEONY_ACCESS_TOKEN', 'PEONY_ACCESS_TOKEN_SECRET'

keys_oauth = {1: oauth1_keys,
              2: oauth2_keys}

# test if the keys are in the environment variables
test_oauth = {i: all(key in os.environ for key in keys_oauth[i])
              for i in keys_oauth}

oauth2_creds = 'consumer_key', 'consumer_secret'
oauth1_creds = *oauth2_creds, 'access_token', 'access_token_secret'

creds_oauth = {1: oauth1_creds,
               2: oauth2_creds}

clients = {1: None,
           2: None}

headers = {1: oauth.OAuth1Headers,
           2: oauth.OAuth2Headers}

def client_oauth(key, event_loop):
    global clients

    if clients[key] is None:
        creds = {k: os.environ[envk]
                 for k, envk in zip(creds_oauth[key], keys_oauth[key])}

        clients[key] = PeonyClient(auth=headers[key], **creds)

    clients[key].loop = event_loop
    clients[key]._session = aiohttp.ClientSession()
    return clients[key]

@pytest.fixture
def oauth1_client(event_loop):
    return client_oauth(1, event_loop)

@pytest.fixture
def oauth2_client(event_loop):
    return client_oauth(2, event_loop)

def decorator_oauth(key):

    def oauth_decorator(func):

        # very dirty don't do this at home
        return pytest.mark.asyncio(
            pytest.mark.twitter(
                pytest.mark.timeout(1800)(
                    pytest.mark.skipif(not test_oauth[key],
                                       reason="no credentials found")(func)
                )
            )
        )

    return oauth_decorator

oauth1_decorator = decorator_oauth(1)
oauth2_decorator = decorator_oauth(2)

@oauth2_decorator
async def test_oauth2_get_token(oauth2_client):
    if 'Authorization' in oauth2_client.headers:
        del oauth2_client.headers['Authorization']

    await oauth2_client.headers.sign()

@oauth2_decorator
async def test_oauth2_request(oauth2_client):
    await oauth2_client.api.search.tweets.get(q="@twitter hello :)")

@pytest.mark.invalidate_token
@oauth2_decorator
async def test_oauth2_invalidate_token(oauth2_client):
    # make sure there is a token
    if 'Authorization' not in oauth2_client.headers:
        await oauth2_client.headers.sign()

    await oauth2_client.headers.invalidate_token()
    assert oauth2_client.headers.token is None

@oauth2_decorator
async def test_oauth2_bearer_token(oauth2_client):
    await oauth2_client.headers.sign()

    token = oauth2_client.headers.token

    client2 = PeonyClient("", "", bearer_token=token,
                          auth=oauth.OAuth2Headers)
    assert client2.headers.token == oauth2_client.headers.token

@oauth2_decorator
async def test_oauth2_nonexisting_endpoint(oauth2_client):
    with pytest.raises(exceptions.DoesNotExist):
        await oauth2_client.api.whereismytweet.get()

@oauth1_decorator
async def test_search(oauth1_client):
    await oauth1_client.api.search.tweets.get(q="@twitter hello :)")

@oauth1_decorator
async def test_search_with_max_id(oauth1_client):
    req = oauth1_client.api.search.tweets.get(q="nasa", count=20)
    all_tweets = set()
    async for tweets in req.iterator.with_max_id():
        # no duplicates
        assert not any(tweet.id in all_tweets for tweet in tweets.statuses)
        all_tweets |= set(tweet.id for tweet in tweets.statuses)

        if len(all_tweets) > 20:
            break

@oauth1_decorator
async def test_user_timeline(oauth1_client):
    req = oauth1_client.api.statuses.user_timeline.get(screen_name="twitter",
                                                       count=20)
    responses = req.iterator.with_max_id()

    all_tweets = set()
    async for tweets in responses:
        # no duplicates
        assert not any(tweet.id in all_tweets for tweet in tweets)
        all_tweets |= set(tweet.id for tweet in tweets)

        if len(all_tweets) > 20:
            break

@oauth1_decorator
async def test_iterator_incorrect_data(oauth1_client):
    req = oauth1_client.api.users.show.get(screen_name="twitter",
                                           include_entities=False)
    responses = req.iterator.with_max_id()

    async for twitter in responses:
        assert twitter == []

@oauth1_decorator
async def test_home_timeline(oauth1_client):
    await oauth1_client.api.statuses.home_timeline.get(count=20)

@oauth1_decorator
async def test_upload_media(oauth1_client, medias):
    media = await medias['lady_peony'].download()
    media = await oauth1_client.upload_media(media)

    await oauth1_client.api.statuses.update.post(status="",
                                                 media_ids=media.media_id)

@oauth1_decorator
async def test_upload_tweet(oauth1_client):
    status = "%d Living in the limelight the universal dream " \
             "for those who wish to seem" % time.time()
    await oauth1_client.api.statuses.update.post(status=status)

@oauth1_decorator
async def test_upload_tweet_with_media(oauth1_client, medias):
    data = await medias['seismic_waves'].download()
    media = await oauth1_client.upload_media(data)
    await oauth1_client.api.statuses.update.post(status="",
                                                 media_ids=media.media_id)

@oauth1_decorator
async def test_upload_tweet_with_media_chunked(oauth1_client, medias):
    for media in (medias[key] for key in ('pink_queen', 'bloom', 'video')):
        data = await media.download()
        media = await oauth1_client.upload_media(data, chunked=True)

        await oauth1_client.api.statuses.update.post(status="",
                                                     media_ids=media.media_id)

@oauth1_decorator
async def test_upload_tweet_with_media_from_url(oauth1_client, medias):
    for media in (medias[key] for key in ('pink_queen', 'bloom', 'video')):
        media = await oauth1_client.upload_media(media.url, chunked=True)

        await oauth1_client.api.statuses.update.post(status="",
                                                     media_ids=media.media_id)

@oauth1_decorator
async def test_direct_message(oauth1_client):
    user = await oauth1_client.user
    message = {
        'event': {
            'type': "message_create",
            'message_create': {
                'target': {'recipient_id': user.id},
                'message_data': {
                    'text': "test %d" % time.time(),
                    'quick_reply': {
                        'type': "options",
                        'options': [
                            {'label': "Hello",
                             'description': "Hello",
                             'metadata': "foo"},
                            {'label': "World",
                             'description': "World",
                             'metadata': "bar"}
                        ]
                    }
                }
            }
        }
    }
    await oauth1_client.api.direct_messages.events.new.post(_json=message)
<commit_msg>tests: expect HTTPNotFound on nonexisting endpoint (?)<commit_after>
import os
import time

import aiohttp
import pytest

from peony import PeonyClient, exceptions, oauth

oauth2_keys = 'PEONY_CONSUMER_KEY', 'PEONY_CONSUMER_SECRET'

oauth1_keys = *oauth2_keys, 'PEONY_ACCESS_TOKEN', 'PEONY_ACCESS_TOKEN_SECRET'

keys_oauth = {1: oauth1_keys,
              2: oauth2_keys}

# test if the keys are in the environment variables
test_oauth = {i: all(key in os.environ for key in keys_oauth[i])
              for i in keys_oauth}

oauth2_creds = 'consumer_key', 'consumer_secret'
oauth1_creds = *oauth2_creds, 'access_token', 'access_token_secret'

creds_oauth = {1: oauth1_creds,
               2: oauth2_creds}

clients = {1: None,
           2: None}

headers = {1: oauth.OAuth1Headers,
           2: oauth.OAuth2Headers}

def client_oauth(key, event_loop):
    global clients

    if clients[key] is None:
        creds = {k: os.environ[envk]
                 for k, envk in zip(creds_oauth[key], keys_oauth[key])}

        clients[key] = PeonyClient(auth=headers[key], **creds)

    clients[key].loop = event_loop
    clients[key]._session = aiohttp.ClientSession()
    return clients[key]

@pytest.fixture
def oauth1_client(event_loop):
    return client_oauth(1, event_loop)

@pytest.fixture
def oauth2_client(event_loop):
    return client_oauth(2, event_loop)

def decorator_oauth(key):

    def oauth_decorator(func):

        # very dirty don't do this at home
        return pytest.mark.asyncio(
            pytest.mark.twitter(
                pytest.mark.timeout(1800)(
                    pytest.mark.skipif(not test_oauth[key],
                                       reason="no credentials found")(func)
                )
            )
        )

    return oauth_decorator

oauth1_decorator = decorator_oauth(1)
oauth2_decorator = decorator_oauth(2)

@oauth2_decorator
async def test_oauth2_get_token(oauth2_client):
    if 'Authorization' in oauth2_client.headers:
        del oauth2_client.headers['Authorization']

    await oauth2_client.headers.sign()

@oauth2_decorator
async def test_oauth2_request(oauth2_client):
    await oauth2_client.api.search.tweets.get(q="@twitter hello :)")

@pytest.mark.invalidate_token
@oauth2_decorator
async def test_oauth2_invalidate_token(oauth2_client):
    # make sure there is a token
    if 'Authorization' not in oauth2_client.headers:
        await oauth2_client.headers.sign()

    await oauth2_client.headers.invalidate_token()
    assert oauth2_client.headers.token is None

@oauth2_decorator
async def test_oauth2_bearer_token(oauth2_client):
    await oauth2_client.headers.sign()

    token = oauth2_client.headers.token

    client2 = PeonyClient("", "", bearer_token=token,
                          auth=oauth.OAuth2Headers)
    assert client2.headers.token == oauth2_client.headers.token

@oauth2_decorator
async def test_oauth2_nonexisting_endpoint(oauth2_client):
    with pytest.raises(exceptions.HTTPNotFound):
        await oauth2_client.api.whereismytweet.get()

@oauth1_decorator
async def test_search(oauth1_client):
    await oauth1_client.api.search.tweets.get(q="@twitter hello :)")

@oauth1_decorator
async def test_search_with_max_id(oauth1_client):
    req = oauth1_client.api.search.tweets.get(q="nasa", count=20)
    all_tweets = set()
    async for tweets in req.iterator.with_max_id():
        # no duplicates
        assert not any(tweet.id in all_tweets for tweet in tweets.statuses)
        all_tweets |= set(tweet.id for tweet in tweets.statuses)

        if len(all_tweets) > 20:
            break

@oauth1_decorator
async def test_user_timeline(oauth1_client):
    req = oauth1_client.api.statuses.user_timeline.get(screen_name="twitter",
                                                       count=20)
    responses = req.iterator.with_max_id()

    all_tweets = set()
    async for tweets in responses:
        # no duplicates
        assert not any(tweet.id in all_tweets for tweet in tweets)
        all_tweets |= set(tweet.id for tweet in tweets)

        if len(all_tweets) > 20:
            break

@oauth1_decorator
async def test_iterator_incorrect_data(oauth1_client):
    req = oauth1_client.api.users.show.get(screen_name="twitter",
                                           include_entities=False)
    responses = req.iterator.with_max_id()

    async for twitter in responses:
        assert twitter == []

@oauth1_decorator
async def test_home_timeline(oauth1_client):
    await oauth1_client.api.statuses.home_timeline.get(count=20)

@oauth1_decorator
async def test_upload_media(oauth1_client, medias):
    media = await medias['lady_peony'].download()
    media = await oauth1_client.upload_media(media)

    await oauth1_client.api.statuses.update.post(status="",
                                                 media_ids=media.media_id)

@oauth1_decorator
async def test_upload_tweet(oauth1_client):
    status = "%d Living in the limelight the universal dream " \
             "for those who wish to seem" % time.time()
    await oauth1_client.api.statuses.update.post(status=status)

@oauth1_decorator
async def test_upload_tweet_with_media(oauth1_client, medias):
    data = await medias['seismic_waves'].download()
    media = await oauth1_client.upload_media(data)
    await oauth1_client.api.statuses.update.post(status="",
                                                 media_ids=media.media_id)

@oauth1_decorator
async def test_upload_tweet_with_media_chunked(oauth1_client, medias):
    for media in (medias[key] for key in ('pink_queen', 'bloom', 'video')):
        data = await media.download()
        media = await oauth1_client.upload_media(data, chunked=True)

        await oauth1_client.api.statuses.update.post(status="",
                                                     media_ids=media.media_id)

@oauth1_decorator
async def test_upload_tweet_with_media_from_url(oauth1_client, medias):
    for media in (medias[key] for key in ('pink_queen', 'bloom', 'video')):
        media = await oauth1_client.upload_media(media.url, chunked=True)

        await oauth1_client.api.statuses.update.post(status="",
                                                     media_ids=media.media_id)

@oauth1_decorator
async def test_direct_message(oauth1_client):
    user = await oauth1_client.user
    message = {
        'event': {
            'type': "message_create",
            'message_create': {
                'target': {'recipient_id': user.id},
                'message_data': {
                    'text': "test %d" % time.time(),
                    'quick_reply': {
                        'type': "options",
                        'options': [
                            {'label': "Hello",
                             'description': "Hello",
                             'metadata': "foo"},
                            {'label': "World",
                             'description': "World",
                             'metadata': "bar"}
                        ]
                    }
                }
            }
        }
    }
    await oauth1_client.api.direct_messages.events.new.post(_json=message)
<|endoftext|>


<commit_before>import json
import logging
import os
import tempfile
import unittest

import pybel
from pybel.canonicalize import postpend_location, decanonicalize_node
from pybel.constants import GOCC_LATEST, FUNCTION
from tests import constants
from tests.constants import test_bel, test_bel_4, mock_bel_resources

log = logging.getLogger('pybel')

pd_path = os.path.expanduser('~/dev/bms/aetionomy/parkinsons.bel')
small_corpus_path = os.path.expanduser('~/dev/bms/selventa/small_corpus.bel')

class TestCanonicalizeHelper(unittest.TestCase):
    def test_postpend_location_failure(self):
        with self.assertRaises(ValueError):
            postpend_location('', dict(name='failure'))

    def test_decanonicalize_node_failure(self):
        with self.assertRaises(ValueError):
            class NotGraph:
                node = None

            x = NotGraph()
            x.node = {'test_node': {FUNCTION: 'nope'}}

            decanonicalize_node(x, 'test_node')

    def test_members(self):
        s = 'p(SFAM:"ETS Family") hasMembers list (p(HGNC:GABPA),p(HGNC:NRF1))'
        #: this string is failing
        self.fail('need to write tests against the canonicalization and decanonicalization of this statement')

class TestCanonicalize(unittest.TestCase):
    def setUp(self):
        self.dir = tempfile.mkdtemp()
        self.path = os.path.join(self.dir, 'test.bel')

    def tearDown(self):
        os.remove(self.path)
        os.rmdir(self.dir)

    def canonicalize_helper(self, test_path):
        original = pybel.from_path(test_path)

        with open(self.path, 'w') as f:
            pybel.to_bel(original, f)

        reloaded = pybel.from_path(self.path)

        original.namespace_url['GOCC'] = GOCC_LATEST

        self.assertEqual(original.document, reloaded.document)
        self.assertEqual(original.namespace_owl, reloaded.namespace_owl)
        self.assertEqual(original.namespace_url, reloaded.namespace_url)
        self.assertEqual(original.annotation_url, reloaded.annotation_url)
        self.assertEqual(original.annotation_list, reloaded.annotation_list)

        self.assertEqual(set(original.nodes()), set(reloaded.nodes()))
        self.assertEqual(set(original.edges()), set(reloaded.edges()))

        # Really test everything is exactly the same, down to the edge data

        fmt = "Nodes with problem: {}, {}.\nOld Data:\n{}\nNew Data:\{}"
        for u, v, d in original.edges_iter(data=True):
            if d['relation'] == 'hasMember':
                continue

            for d1 in original.edge[u][v].values():
                x = False

                for d2 in reloaded.edge[u][v].values():
                    if set(d1.keys()) == set(d2.keys()) and all(d1[k] == d2[k] for k in d1):
                        x = True

                self.assertTrue(x, msg=fmt.format(u, v, json.dumps(original.edge[u][v], indent=2, sort_keys=True),
                                                  json.dumps(reloaded.edge[u][v], indent=2, sort_keys=True)))

    @mock_bel_resources
    def test_canonicalize_1(self, mock_get):
        self.canonicalize_helper(test_bel)

    @mock_bel_resources
    @constants.mock_parse_owl_ontospy
    @constants.mock_parse_owl_pybel
    def test_canonicalize_4(self, m1, m2, m3):
        self.canonicalize_helper(test_bel_4)

    @unittest.skipUnless(os.path.exists(small_corpus_path), 'Small Corpus Missing')
    def test_small_corpus(self):
        self.maxDiff = None
        self.canonicalize_helper(small_corpus_path)

    @unittest.skipUnless(os.path.exists(pd_path), 'PD Test File Missing')
    def test_parkinsons(self):
        self.maxDiff = None
        self.canonicalize_helper(pd_path)
<commit_msg>Remove testing for members<commit_after>import json
import logging
import os
import tempfile
import unittest

import pybel
from pybel.canonicalize import postpend_location, decanonicalize_node
from pybel.constants import GOCC_LATEST, FUNCTION
from tests import constants
from tests.constants import test_bel, test_bel_4, mock_bel_resources

log = logging.getLogger('pybel')

pd_path = os.path.expanduser('~/dev/bms/aetionomy/parkinsons.bel')
small_corpus_path = os.path.expanduser('~/dev/bms/selventa/small_corpus.bel')

class TestCanonicalizeHelper(unittest.TestCase):
    def test_postpend_location_failure(self):
        with self.assertRaises(ValueError):
            postpend_location('', dict(name='failure'))

    def test_decanonicalize_node_failure(self):
        with self.assertRaises(ValueError):
            class NotGraph:
                node = None

            x = NotGraph()
            x.node = {'test_node': {FUNCTION: 'nope'}}

            decanonicalize_node(x, 'test_node')

class TestCanonicalize(unittest.TestCase):
    def setUp(self):
        self.dir = tempfile.mkdtemp()
        self.path = os.path.join(self.dir, 'test.bel')

    def tearDown(self):
        os.remove(self.path)
        os.rmdir(self.dir)

    def canonicalize_helper(self, test_path):
        original = pybel.from_path(test_path)

        with open(self.path, 'w') as f:
            pybel.to_bel(original, f)

        reloaded = pybel.from_path(self.path)

        original.namespace_url['GOCC'] = GOCC_LATEST

        self.assertEqual(original.document, reloaded.document)
        self.assertEqual(original.namespace_owl, reloaded.namespace_owl)
        self.assertEqual(original.namespace_url, reloaded.namespace_url)
        self.assertEqual(original.annotation_url, reloaded.annotation_url)
        self.assertEqual(original.annotation_list, reloaded.annotation_list)

        self.assertEqual(set(original.nodes()), set(reloaded.nodes()))
        self.assertEqual(set(original.edges()), set(reloaded.edges()))

        # Really test everything is exactly the same, down to the edge data

        fmt = "Nodes with problem: {}, {}.\nOld Data:\n{}\nNew Data:\{}"
        for u, v, d in original.edges_iter(data=True):
            if d['relation'] == 'hasMember':
                continue

            for d1 in original.edge[u][v].values():
                x = False

                for d2 in reloaded.edge[u][v].values():
                    if set(d1.keys()) == set(d2.keys()) and all(d1[k] == d2[k] for k in d1):
                        x = True

                self.assertTrue(x, msg=fmt.format(u, v, json.dumps(original.edge[u][v], indent=2, sort_keys=True),
                                                  json.dumps(reloaded.edge[u][v], indent=2, sort_keys=True)))

    @mock_bel_resources
    def test_canonicalize_1(self, mock_get):
        self.canonicalize_helper(test_bel)

    @mock_bel_resources
    @constants.mock_parse_owl_ontospy
    @constants.mock_parse_owl_pybel
    def test_canonicalize_4(self, m1, m2, m3):
        self.canonicalize_helper(test_bel_4)

    @unittest.skipUnless(os.path.exists(small_corpus_path), 'Small Corpus Missing')
    def test_small_corpus(self):
        self.maxDiff = None
        self.canonicalize_helper(small_corpus_path)

    @unittest.skipUnless(os.path.exists(pd_path), 'PD Test File Missing')
    def test_parkinsons(self):
        self.maxDiff = None
        self.canonicalize_helper(pd_path)
<|endoftext|>


<commit_before>import pytest

from code.weather import set_screen_color

def test_set_screen_color():
    assert set_screen_color(25) == [0, 0, 155]
<commit_msg>another change<commit_after>import pytest

from weather import set_screen_color

def test_set_screen_color():
    assert set_screen_color(25) == [0, 0, 155]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
@author: Federico Cerchiari <federicocerchiari@gmail.com>
"""
import unittest
from copy import copy

from tempy.widgets import TempyTable, TempyList
from tempy.tags import Table, Tr, Td, Ul, Ol, Li

from tempy.exceptions import WidgetDataError

class TestTempyTable(unittest.TestCase):

    def setUp(self):
        self.data = [[x*y for x in range(1, 11)] for y in range(15)]

    def verify_content(self, table):
        # Check table content
        self.assertTrue(0 in table.body[0][4])
        self.assertTrue(1 in table.body[1][0])
        self.assertTrue(2 in table.body[1][1])
        self.assertTrue(16 in table.body[8][1])

    def test_empty_creation(self):
        table = TempyTable()
        self.assertFalse(table.body)
        # Future non-regression, TempyTable should remain a Table Tag
        self.assertIsInstance(table, Table)

    def test_skeleton_creation(self):
        table = TempyTable(rows=15, cols=10)
        self.assertTrue(table.body)
        # Check table sizes
        self.assertEqual(len(table.body), 15)
        self.assertEqual(len(table.body[0]), 10)

    def test_caption(self):
        table = TempyTable(caption='Test Table')
        self.assertTrue('Test Table' in table.caption)

    def test_init_from_data(self):
        table = TempyTable(data=self.data)
        self.assertEqual(len(table.body), 15)
        self.assertEqual(len(table.body[0]), 10)
        self.verify_content(table)

    def test_populate(self):
        # test no_resize
        table = TempyTable().populate(self.data)
        # Check table sizes
        self.assertEqual(len(table.body), 15)
        self.assertEqual(len(table.body[0]), 10)

        # test add row
        new_data = copy(self.data)
        new_data.append(list(range(1, 11)))
        table.populate(new_data)
        self.assertEqual(len(table.body), 16)

        # test raise with no resize
        new_data[0].append('test')
        with self.assertRaises(WidgetDataError):
            table.populate(new_data)

        # test resize
        table.populate(new_data, resize_x=True)
        self.assertEqual(len(table.body), 16)
        self.assertEqual(len(table.body[0]), 11)
        self.assertEqual(len(table.body[1]), 11)
        self.assertTrue('test' in table.body[0][10])

        # test not force
        for row in new_data:
            try:
                row[10] = 'test1'
            except:
                row.append('test1')
        table.populate(new_data, force=False)
        self.assertEqual(len(table.body), 16)
        self.assertEqual(len(table.body[0]), 11)
        self.assertEqual(len(table.body[1]), 11)
        self.assertFalse('test1' in table.body[0][10])
        self.assertTrue('test1' in table.body[1][10])

        # test force
        table.populate(new_data, force=True)
        self.assertTrue('test1' in table.body[0][10])

        # test non normalize:
        new_data[3].append('test2')
        table.populate(new_data, normalize=False, resize_x=True)
        self.assertTrue('test2' in table.body[3][11])
        with self.assertRaises(IndexError):
            table.body[0][11]

    def test_pop_row(self):
        table = TempyTable(data=self.data)

        # test pop last
        r = table.pop_row()
        self.assertEqual(r, self.data[-1])

        # test pop get tags
        r = table.pop_row(tags=True)
        test_row = Tr()(Td()(c) for c in self.data[-2])
        for cell, t_cell in zip(r, test_row):
            self.assertEqual(cell, t_cell)

        # test pop by index
        r = table.pop_row(0)
        self.assertEqual(r, self.data[0])

    def test_pop_cell(self):
        table = TempyTable(data=self.data)

        # test pop last
        r = table.pop_cell()
        self.assertEqual(r, self.data[-1][-1])

        # test pop get tags
        r = table.pop_cell(tags=True)
        test_cell = Td()(self.data[-2][-1])
        self.assertEqual(r, test_cell)

        # test pop by index row
        r = table.pop_cell(0)
        self.assertEqual(r, self.data[0][-1])

        # test pop by index row andcol
        r = table.pop_cell(0, 0)
        self.assertEqual(r, self.data[0][0])

class TestTempyList(unittest.TestCase):

    def test_create_empty(self):
        li = TempyList()
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 0)

        li = TempyList(Ol)
        self.assertIsInstance(li, Ol)
        self.assertEqual(len(li), 0)

        li = TempyList('Ol')
        self.assertIsInstance(li, Ol)
        self.assertEqual(len(li), 0)

    def test_populate_empty(self):
        li = TempyList()
        li.populate([1, 2, 3])
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)

    def test_create_full(self):
        li = TempyList(struct=[1, 2, 3])
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)
        self.assertTrue(1 in li[0])

        li = TempyList(struct={1, 2, 3})
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)

        li = TempyList(struct={1: None, 2: None, 3: None, '_typ': Ol})
        self.assertIsInstance(li, Ol)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)

    def test_populate_recursive(self):
        li = TempyList()
        li.populate({1: None, 2: ['a', 'b', 'c'], 3: {'test': [1, 2, 3]}})
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)
        self.assertIsInstance(li[1][1], Ul)
        self.assertEqual(len(li[1][1]), 3)
        self.assertIsInstance(li[1][1][0], Li)
        self.assertTrue('a' in li[1][1][0])

if __name__ == '__main__':
    unittest.main()
<commit_msg>dicts are ordered by default from Python3.5?<commit_after># -*- coding: utf-8 -*-
"""
@author: Federico Cerchiari <federicocerchiari@gmail.com>
"""
import unittest
from copy import copy

from tempy.widgets import TempyTable, TempyList
from tempy.tags import Table, Tr, Td, Ul, Ol, Li

from tempy.exceptions import WidgetDataError

class TestTempyTable(unittest.TestCase):

    def setUp(self):
        self.data = [[x*y for x in range(1, 11)] for y in range(15)]

    def verify_content(self, table):
        # Check table content
        self.assertTrue(0 in table.body[0][4])
        self.assertTrue(1 in table.body[1][0])
        self.assertTrue(2 in table.body[1][1])
        self.assertTrue(16 in table.body[8][1])

    def test_empty_creation(self):
        table = TempyTable()
        self.assertFalse(table.body)
        # Future non-regression, TempyTable should remain a Table Tag
        self.assertIsInstance(table, Table)

    def test_skeleton_creation(self):
        table = TempyTable(rows=15, cols=10)
        self.assertTrue(table.body)
        # Check table sizes
        self.assertEqual(len(table.body), 15)
        self.assertEqual(len(table.body[0]), 10)

    def test_caption(self):
        table = TempyTable(caption='Test Table')
        self.assertTrue('Test Table' in table.caption)

    def test_init_from_data(self):
        table = TempyTable(data=self.data)
        self.assertEqual(len(table.body), 15)
        self.assertEqual(len(table.body[0]), 10)
        self.verify_content(table)

    def test_populate(self):
        # test no_resize
        table = TempyTable().populate(self.data)
        # Check table sizes
        self.assertEqual(len(table.body), 15)
        self.assertEqual(len(table.body[0]), 10)

        # test add row
        new_data = copy(self.data)
        new_data.append(list(range(1, 11)))
        table.populate(new_data)
        self.assertEqual(len(table.body), 16)

        # test raise with no resize
        new_data[0].append('test')
        with self.assertRaises(WidgetDataError):
            table.populate(new_data)

        # test resize
        table.populate(new_data, resize_x=True)
        self.assertEqual(len(table.body), 16)
        self.assertEqual(len(table.body[0]), 11)
        self.assertEqual(len(table.body[1]), 11)
        self.assertTrue('test' in table.body[0][10])

        # test not force
        for row in new_data:
            try:
                row[10] = 'test1'
            except:
                row.append('test1')
        table.populate(new_data, force=False)
        self.assertEqual(len(table.body), 16)
        self.assertEqual(len(table.body[0]), 11)
        self.assertEqual(len(table.body[1]), 11)
        self.assertFalse('test1' in table.body[0][10])
        self.assertTrue('test1' in table.body[1][10])

        # test force
        table.populate(new_data, force=True)
        self.assertTrue('test1' in table.body[0][10])

        # test non normalize:
        new_data[3].append('test2')
        table.populate(new_data, normalize=False, resize_x=True)
        self.assertTrue('test2' in table.body[3][11])
        with self.assertRaises(IndexError):
            table.body[0][11]

    def test_pop_row(self):
        table = TempyTable(data=self.data)

        # test pop last
        r = table.pop_row()
        self.assertEqual(r, self.data[-1])

        # test pop get tags
        r = table.pop_row(tags=True)
        test_row = Tr()(Td()(c) for c in self.data[-2])
        for cell, t_cell in zip(r, test_row):
            self.assertEqual(cell, t_cell)

        # test pop by index
        r = table.pop_row(0)
        self.assertEqual(r, self.data[0])

    def test_pop_cell(self):
        table = TempyTable(data=self.data)

        # test pop last
        r = table.pop_cell()
        self.assertEqual(r, self.data[-1][-1])

        # test pop get tags
        r = table.pop_cell(tags=True)
        test_cell = Td()(self.data[-2][-1])
        self.assertEqual(r, test_cell)

        # test pop by index row
        r = table.pop_cell(0)
        self.assertEqual(r, self.data[0][-1])

        # test pop by index row andcol
        r = table.pop_cell(0, 0)
        self.assertEqual(r, self.data[0][0])

class TestTempyList(unittest.TestCase):

    def test_create_empty(self):
        li = TempyList()
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 0)

        li = TempyList(Ol)
        self.assertIsInstance(li, Ol)
        self.assertEqual(len(li), 0)

        li = TempyList('Ol')
        self.assertIsInstance(li, Ol)
        self.assertEqual(len(li), 0)

    def test_populate_empty(self):
        li = TempyList()
        li.populate([1, 2, 3])
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)

    def test_create_full(self):
        li = TempyList(struct=[1, 2, 3])
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)
        self.assertTrue(1 in li[0])

        li = TempyList(struct={1, 2, 3})
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)

        li = TempyList(struct={1: None, 2: None, 3: None, '_typ': Ol})
        self.assertIsInstance(li, Ol)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)

    def test_populate_recursive(self):
        li = TempyList()
        li.populate({1: None, 2: ['a', 'b', 'c'], 3: {'test': [1, 2, 3]}})
        self.assertIsInstance(li, Ul)
        self.assertEqual(len(li), 3)
        self.assertIsInstance(li[0], Li)
        self.assertIsInstance(li[1][1], Ul)
        self.assertEqual(len(li[1][1]), 3)
        self.assertIsInstance(li[1][1][0], Li)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  Copyright (C) 2013  Matt Jeffery <matt@clan.se>
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

import sys
import unittest2
from mock import patch

from semetric.apiclient import SemetricAPI
from semetric.apiclient.entity.base import Entity
from semetric.apiclient.entity.artist import Artist
from semetric.apiclient.entity.chart import Chart
from .consts import APIKEY

# Base string type for Python3
if sys.version_info >= (3,): # pragma: no cover
    basestring = str

class TestSemetricAPI(unittest2.TestCase):

    def test_search_request_artist(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = []
            api.search(Artist, name="lady gaga")

        api_mock.assert_called_once_with("artist", q="lady gaga")

    def test_get_request_artist(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = []
            api.get(Artist, id="foo")

        api_mock.assert_called_once_with("artist/foo")

    def test_search_request_not_implemented(self):

        api = SemetricAPI(APIKEY)
        with self.assertRaises(NotImplementedError) as exc:
            api.search(Entity, name="foo")

    def test_search_request_entity(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = []
            api.get(Entity, id="foo")

        api_mock.assert_called_once_with("entity/foo")

    def test_get_chart_with_country(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = ""
            api.get(Chart, id="0695f0bba6144dfaa390e9b9f017ceab", country="CA")

        api_mock.assert_called_once_with("chart/0695f0bba6144dfaa390e9b9f017ceab", country="CA")

    def test_get_chart_without_country(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = ""
            api.get(Chart, id="0695f0bba6144dfaa390e9b9f017ceab")

        api_mock.assert_called_once_with("chart/0695f0bba6144dfaa390e9b9f017ceab", country="ALL")

<commit_msg>use sane artist IDs for testing<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  Copyright (C) 2013  Matt Jeffery <matt@clan.se>
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

import sys
import unittest2
from mock import patch

from semetric.apiclient import SemetricAPI
from semetric.apiclient.entity.base import Entity
from semetric.apiclient.entity.artist import Artist
from semetric.apiclient.entity.chart import Chart
from .consts import APIKEY

# Base string type for Python3
if sys.version_info >= (3,): # pragma: no cover
    basestring = str

class TestSemetricAPI(unittest2.TestCase):

    def test_search_request_artist(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = []
            api.search(Artist, name="lady gaga")

        api_mock.assert_called_once_with("artist", q="lady gaga")

    def test_get_request_artist(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = []
            api.get(Artist, id="87e0e4ccc7f4415cbd6ba60ad49943b6")

        api_mock.assert_called_once_with("artist/87e0e4ccc7f4415cbd6ba60ad49943b6")

    def test_search_request_not_implemented(self):

        api = SemetricAPI(APIKEY)
        with self.assertRaises(NotImplementedError) as exc:
            api.search(Entity, name="foo")

    def test_search_request_entity(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = []
            api.get(Entity, id="foo")

        api_mock.assert_called_once_with("entity/foo")

    def test_get_chart_with_country(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = ""
            api.get(Chart, id="0695f0bba6144dfaa390e9b9f017ceab", country="CA")

        api_mock.assert_called_once_with("chart/0695f0bba6144dfaa390e9b9f017ceab", country="CA")

    def test_get_chart_without_country(self):

        api = SemetricAPI(APIKEY)

        # Make the api response
        with patch.object(api.client, 'request', autospec=True) as api_mock:
            api_mock.return_value = ""
            api.get(Chart, id="0695f0bba6144dfaa390e9b9f017ceab")

        api_mock.assert_called_once_with("chart/0695f0bba6144dfaa390e9b9f017ceab", country="ALL")

<|endoftext|>


<commit_before>"""Copyright 2013 Google Inc. All Rights Reserved.

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

import pinject

class VersionTest(unittest.TestCase):

  def testVersion(self):
    self.assertEqual(type(pinject.__version__), str)
    self.assertRegex(pinject.__version__,
                     r'^\d+\.\d+\.\d+$')

if __name__ == '__main__':
  unittest.main()
<commit_msg>Compitable with Python 2<commit_after>"""Copyright 2013 Google Inc. All Rights Reserved.

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

import pinject

class VersionTest(unittest.TestCase):

  def testVersion(self):
    self.assertEqual(type(pinject.__version__), str)
    self.assertRegexpMatches(pinject.__version__,
                             r'^\d+\.\d+\.\d+$')

if __name__ == '__main__':
  unittest.main()
<|endoftext|>


<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-

import unittest
from nose.plugins.attrib import attr

import sys
import os
import PyQt4
from PyQt4.QtGui import QApplication, QFileDialog, QWidget, QPushButton, QVBoxLayout
from PyQt4.QtTest import QTest
from PyQt4.QtCore import Qt

import imtools.sample_data

class MyTestCase(unittest.TestCase):
    def setUp(self):
        from PyQt4.QtGui import QApplication, QFileDialog
        # self.qapp = QApplication(sys.argv)

    @attr('interactive')
    def test_something(self):
        self.assertEqual(True, False)

    @attr('interactive')
    def test_visualization(self):

        import PyQt4
        from PyQt4.QtGui import QApplication, QFileDialog
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.gui import TeigenWidget
        import imtools.show_segmentation_qt
        app = QApplication(sys.argv)
        cfg = {"bool": True, "int":5, 'str': 'strdrr', 'vs':[1.0, 2.5, 7]}
        captions = {"int": "toto je int"}
        cw = imtools.show_segmentation_qt.ShowSegmentationWidget(None)
        cw.show()
        app.exec_()

    @attr('interactive')
    def test_showsegmentation_andclose(self):

        import PyQt4
        from PyQt4.QtGui import QApplication, QFileDialog
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.gui import TeigenWidget
        import imtools.show_segmentation_qt
        app = QApplication(sys.argv)
        cfg = {"bool": True, "int":5, 'str': 'strdrr', 'vs':[1.0, 2.5, 7]}
        captions = {"int": "toto je int"}
        cw = imtools.show_segmentation_qt.ShowSegmentationWidget(None)

        cw.show()
        cw.close()
        # app.exec_()

    @attr('interactive')
    def test_show_segmentation_qt_widget(self):
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.generators.cylindersqt import CylindersWidget
        import imtools.show_segmentation_qt as ssqt
        # app = QtGui.QApplication(sys.argv)
        app = QApplication(sys.argv)
        sw = ssqt.ShowSegmentationWidget(None, show_load_interface=True)
        sw.show()
        app.exec_()

    # def test_show_segmentation_qt_widget_2(self):
        # import imtools.show_segmentation_qt as ssqt
        # app = QApplication(sys.argv)
        # slab = {"la": 5, "sdfa":7}
        # sw = ssqt.SelectLabelWidget(slab=slab)
        # sw.show()
        # app.exec_()

    @attr('interactive')
    def test_show_qtwidget(self):
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.generators.cylindersqt import CylindersWidget
        import imtools.show_segmentation_qt as ssqt
        app = QApplication(sys.argv)
        qw = QWidget()
        button = QPushButton('Test')
        layout = QVBoxLayout()
        layout.addWidget(button)
        # sw = ssqt.ShowSegmentationWidget(None, qapp=app)
        qw.show()
        app.exec_()

    # @attr('interactive')
    def test_show_segmentation_qt_widget_hidden_buttons(self):
        # = np.zeros([10, 10, 10])
        import imtools
        import imtools.sample_data
        # imtools.sam
        # imtools.sample_data.get_sample_data("sliver_training_001")
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.generators.cylindersqt import CylindersWidget
        import imtools.show_segmentation_qt as ssqt
        app = QApplication(sys.argv)
        # app = QApplication([])

        # if "TRAVIS" in os.environ:
        #     app.setGraphicsSystem("openvg")
        # sw = ssqt.ShowSegmentationWidget(None, show_buttons=False)
        sw = ssqt.ShowSegmentationWidget(None, show_load_interface=True)
        self.assertIn("add_data_file", sw.ui_buttons.keys())
        sw.show()
        # app.exec_()
        sw.close()
        sw.deleteLater()
        sw = None
        app.quit()
        app.deleteLater()
        # app.quit()
        # app.exit()

    @attr('interactive')
    def test_show_donut(self):
        """
        creates VTK file from input data
        :return:
        """
        datap = imtools.sample_data.donut()

        segmentation = datap['segmentation']
        voxelsize_mm = datap['voxelsize_mm']

        import imtools.show_segmentation_qt as ssqt
        import gc
        app = QApplication(sys.argv)
        # app.setGraphicsSystem("openvg")
        sw = ssqt.ShowSegmentationWidget(None, show_load_button=True, show_load_interface=True)
        sw.smoothing = False
        sw.add_data(segmentation, voxelsize_mm=voxelsize_mm)
        # QTest.mouseClick(sw.ui_buttons['Show volume'], Qt.LeftButton)
        # sw.add_vtk_file("~/projects/imtools/mesh.vtk")
        sw.show()
        app.exec_()
        # sw.close()
        # sw.deleteLater()

    @attr('interactive')
    def test_show_donut_with_zerosize_label(self):
        """
        creates VTK file from input data
        :return:
        """
        datap = imtools.sample_data.donut()

        segmentation = datap['segmentation']
        voxelsize_mm = datap['voxelsize_mm']

        import imtools.show_segmentation_qt as ssqt
        import gc
        app = QApplication(sys.argv)
        # app.setGraphicsSystem("openvg")
        sw = ssqt.ShowSegmentationWidget(None, show_load_button=True, show_load_interface=True)
        sw.smoothing = False
        datap["slab"]["empty"]=17
        sw.add_data(segmentation, voxelsize_mm=voxelsize_mm, slab=datap["slab"])
        # QTest.mouseClick(sw.ui_buttons['Show volume'], Qt.LeftButton)
        # sw.add_vtk_file("~/projects/imtools/mesh.vtk")
        sw.show()
        app.exec_()
        # sw.close()
        # sw.deleteLater()

    # @attr('interactive')
    def test_add_data_and_show(self):
        """
        creates VTK file from input data
        :return:
        """
        datap = imtools.sample_data.donut()

        segmentation = datap['segmentation']
        voxelsize_mm = datap['voxelsize_mm']

        import imtools.show_segmentation_qt as ssqt
        import gc
        app = QApplication(sys.argv)
        # app.setGraphicsSystem("openvg")
        sw = ssqt.ShowSegmentationWidget(None, show_load_button=True)
        sw.smoothing = False
        sw.add_data(segmentation, voxelsize_mm=voxelsize_mm)
        QTest.mouseClick(sw.ui_buttons['Show volume'], Qt.LeftButton)
        # sw.add_vtk_file("~/projects/imtools/mesh.vtk")
        sw.show()
        # app.exec_(exec_)
        sw.close()
        sw.deleteLater()

        sw = None
        app.quit()
        app.deleteLater()
        # self.qapp.quit()
        # self.qapp.deleteLater()
        # gc.collect()

        # app = QApplication(sys.argv)
        # app.quit()
        # self.qapp.exit()
        # app.exec_()

if __name__ == '__main__':
    unittest.main()
<commit_msg>skip on travis<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-

import unittest
from nose.plugins.attrib import attr

import sys
import os
import PyQt4
from PyQt4.QtGui import QApplication, QFileDialog, QWidget, QPushButton, QVBoxLayout
from PyQt4.QtTest import QTest
from PyQt4.QtCore import Qt

import imtools.sample_data

class MyTestCase(unittest.TestCase):
    def setUp(self):
        from PyQt4.QtGui import QApplication, QFileDialog
        # self.qapp = QApplication(sys.argv)

    @attr('interactive')
    def test_something(self):
        self.assertEqual(True, False)

    @attr('interactive')
    def test_visualization(self):

        import PyQt4
        from PyQt4.QtGui import QApplication, QFileDialog
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.gui import TeigenWidget
        import imtools.show_segmentation_qt
        app = QApplication(sys.argv)
        cfg = {"bool": True, "int":5, 'str': 'strdrr', 'vs':[1.0, 2.5, 7]}
        captions = {"int": "toto je int"}
        cw = imtools.show_segmentation_qt.ShowSegmentationWidget(None)
        cw.show()
        app.exec_()

    @attr('interactive')
    def test_showsegmentation_andclose(self):

        import PyQt4
        from PyQt4.QtGui import QApplication, QFileDialog
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.gui import TeigenWidget
        import imtools.show_segmentation_qt
        app = QApplication(sys.argv)
        cfg = {"bool": True, "int":5, 'str': 'strdrr', 'vs':[1.0, 2.5, 7]}
        captions = {"int": "toto je int"}
        cw = imtools.show_segmentation_qt.ShowSegmentationWidget(None)

        cw.show()
        cw.close()
        # app.exec_()

    @attr('interactive')
    def test_show_segmentation_qt_widget(self):
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.generators.cylindersqt import CylindersWidget
        import imtools.show_segmentation_qt as ssqt
        # app = QtGui.QApplication(sys.argv)
        app = QApplication(sys.argv)
        sw = ssqt.ShowSegmentationWidget(None, show_load_interface=True)
        sw.show()
        app.exec_()

    # def test_show_segmentation_qt_widget_2(self):
        # import imtools.show_segmentation_qt as ssqt
        # app = QApplication(sys.argv)
        # slab = {"la": 5, "sdfa":7}
        # sw = ssqt.SelectLabelWidget(slab=slab)
        # sw.show()
        # app.exec_()

    @attr('interactive')
    def test_show_qtwidget(self):
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.generators.cylindersqt import CylindersWidget
        import imtools.show_segmentation_qt as ssqt
        app = QApplication(sys.argv)
        qw = QWidget()
        button = QPushButton('Test')
        layout = QVBoxLayout()
        layout.addWidget(button)
        # sw = ssqt.ShowSegmentationWidget(None, qapp=app)
        qw.show()
        app.exec_()

    # @attr('interactive')
    def test_show_segmentation_qt_widget_hidden_buttons(self):
        # = np.zeros([10, 10, 10])
        import imtools
        import imtools.sample_data
        # imtools.sam
        # imtools.sample_data.get_sample_data("sliver_training_001")
        # from teigen.dictwidgetqt import DictWidget
        # from teigen.generators.cylindersqt import CylindersWidget
        import imtools.show_segmentation_qt as ssqt
        app = QApplication(sys.argv)
        # app = QApplication([])

        # if "TRAVIS" in os.environ:
        #     app.setGraphicsSystem("openvg")
        # sw = ssqt.ShowSegmentationWidget(None, show_buttons=False)
        sw = ssqt.ShowSegmentationWidget(None, show_load_interface=True)
        self.assertIn("add_data_file", sw.ui_buttons.keys())
        sw.show()
        # app.exec_()
        sw.close()
        sw.deleteLater()
        sw = None
        app.quit()
        app.deleteLater()
        # app.quit()
        # app.exit()

    @attr('interactive')
    def test_show_donut(self):
        """
        creates VTK file from input data
        :return:
        """
        datap = imtools.sample_data.donut()

        segmentation = datap['segmentation']
        voxelsize_mm = datap['voxelsize_mm']

        import imtools.show_segmentation_qt as ssqt
        import gc
        app = QApplication(sys.argv)
        # app.setGraphicsSystem("openvg")
        sw = ssqt.ShowSegmentationWidget(None, show_load_button=True, show_load_interface=True)
        sw.smoothing = False
        sw.add_data(segmentation, voxelsize_mm=voxelsize_mm)
        # QTest.mouseClick(sw.ui_buttons['Show volume'], Qt.LeftButton)
        # sw.add_vtk_file("~/projects/imtools/mesh.vtk")
        sw.show()
        app.exec_()
        # sw.close()
        # sw.deleteLater()

    @attr('interactive')
    def test_show_donut_with_zerosize_label(self):
        """
        creates VTK file from input data
        :return:
        """
        datap = imtools.sample_data.donut()

        segmentation = datap['segmentation']
        voxelsize_mm = datap['voxelsize_mm']

        import imtools.show_segmentation_qt as ssqt
        import gc
        app = QApplication(sys.argv)
        # app.setGraphicsSystem("openvg")
        sw = ssqt.ShowSegmentationWidget(None, show_load_button=True, show_load_interface=True)
        sw.smoothing = False
        datap["slab"]["empty"]=17
        sw.add_data(segmentation, voxelsize_mm=voxelsize_mm, slab=datap["slab"])
        # QTest.mouseClick(sw.ui_buttons['Show volume'], Qt.LeftButton)
        # sw.add_vtk_file("~/projects/imtools/mesh.vtk")
        sw.show()
        app.exec_()
        # sw.close()
        # sw.deleteLater()

    # @attr('interactive')
    @unittest.skipIf(os.environ.get("TRAVIS", True), "Skip on Travis-CI")
    def test_add_data_and_show(self):
        """
        creates VTK file from input data and show and quit
        :return:
        """
        datap = imtools.sample_data.donut()

        segmentation = datap['segmentation']
        voxelsize_mm = datap['voxelsize_mm']

        import imtools.show_segmentation_qt as ssqt
        import gc
        app = QApplication(sys.argv)
        # app.setGraphicsSystem("openvg")
        sw = ssqt.ShowSegmentationWidget(None, show_load_button=True)
        sw.smoothing = False
        sw.add_data(segmentation, voxelsize_mm=voxelsize_mm)
        QTest.mouseClick(sw.ui_buttons['Show volume'], Qt.LeftButton)
        # sw.add_vtk_file("~/projects/imtools/mesh.vtk")
        sw.show()
        # app.exec_(exec_)
        sw.close()
        sw.deleteLater()

        sw = None
        app.quit()
        app.deleteLater()
        # self.qapp.quit()
        # self.qapp.deleteLater()
        # gc.collect()

        # app = QApplication(sys.argv)
        # app.quit()
        # self.qapp.exit()
        # app.exec_()

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before># vim: ts=4 sw=4 et ai:
# -*- coding: utf8 -*-
"""This module implements all contexts for state handling during uploads and
downloads, the main interface to which being the TftpContext base class.

The concept is simple. Each context object represents a single upload or
download, and the state object in the context object represents the current
state of that transfer. The state object has a handle() method that expects
the next packet in the transfer, and returns a state object until the transfer
is complete, at which point it returns None. That is, unless there is a fatal
error, in which case a TftpException is returned instead."""

from .TftpShared import *
from .TftpPacketTypes import *
from .TftpPacketFactory import TftpPacketFactory
from .TftpStates import *
from . import compat
import socket
import time
import sys
import os
import logging

log = logging.getLogger('tftpy.TftpContext')

###############################################################################
# Utility classes
###############################################################################

class TftpMetrics(object):
    """A class representing metrics of the transfer."""
    def __init__(self):
        # Bytes transferred
        self.bytes = 0
        # Bytes re-sent
        self.resent_bytes = 0
        # Duplicate packets received
        self.dups = {}
        self.dupcount = 0
        # Times
        self.start_time = 0
        self.end_time = 0
        self.duration = 0
        # Rates
        self.bps = 0
        self.kbps = 0
        # Generic errors
        self.errors = 0

    def compute(self):
        # Compute transfer time
        self.duration = self.end_time - self.start_time
        if self.duration == 0:
            self.duration = 1
        log.debug("TftpMetrics.compute: duration is %s", self.duration)
        self.bps = (self.bytes * 8.0) / self.duration
        self.kbps = self.bps / 1024.0
        log.debug("TftpMetrics.compute: kbps is %s", self.kbps)
        for key in self.dups:
            self.dupcount += self.dups[key]

    def add_dup(self, pkt):
        """This method adds a dup for a packet to the metrics."""
        log.debug("Recording a dup of %s", pkt)
        s = str(pkt)
        if s in self.dups:
            self.dups[s] += 1
        else:
            self.dups[s] = 1
        tftpassert(self.dups[s] < MAX_DUPS, "Max duplicates reached")

###############################################################################
# Context classes
###############################################################################

class TftpContext(object):
    """The base class of the contexts."""

    def __init__(self, host, port, timeout, retries=DEF_TIMEOUT_RETRIES, localip = ""):
        """Constructor for the base context, setting shared instance
        variables."""
        self.file_to_transfer = None
        self.fileobj = None
        self.options = None
        self.packethook = None
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        if localip != "":
            self.sock.bind((localip, 0))
        self.sock.settimeout(timeout)
        self.timeout = timeout
        self.retries = retries
        self.state = None
        self.next_block = 0
        self.factory = TftpPacketFactory()
        # Note, setting the host will also set self.address, as it's a property.
        self.host = host
        self.port = port
        # The port associated with the TID
        self.tidport = None
        # Metrics
        self.metrics = TftpMetrics()
        # Fluag when the transfer is pending completion.
        self.pending_complete = False
        # Time when this context last received any traffic.
        # FIXME: does this belong in metrics?
        self.last_update = 0
        # The last packet we sent, if applicable, to make resending easy.
        self.last_pkt = None
        # Count the number of retry attempts.
        self.retry_count = 0

    def getBlocksize(self):
        """Fetch the current blocksize for this session."""
        return int(self.options.get('blksize', 512))

    def __del__(self):
        """Simple destructor to try to call housekeeping in the end method if
        not called explicitely. Leaking file descriptors is not a good
        thing."""
        self.end()

    def checkTimeout(self, now):
        """Compare current time with last_update time, and raise an exception
        if we're over the timeout time."""
        log.debug("checking for timeout on session %s", self)
        if now - self.last_update > self.timeout:
            raise TftpTimeout("Timeout waiting for traffic")

    def start(self):
        raise NotImplementedError("Abstract method")

    def end(self, close_fileobj=True):
        """Perform session cleanup, since the end method should always be
        called explicitely by the calling code, this works better than the
        destructor.
        Set close_fileobj to False so fileobj can be returned open."""
        log.debug("in TftpContext.end - closing socket")
        self.sock.close()
        if close_fileobj and self.fileobj is not None and not self.fileobj.closed:
            log.debug("self.fileobj is open - closing")
            self.fileobj.close()

    def gethost(self):
        "Simple getter method for use in a property."
        return self.__host

    def sethost(self, host):
        """Setter method that also sets the address property as a result
        of the host that is set."""
        self.__host = host
        self.address = socket.gethostbyname(host)

    host = property(gethost, sethost)

    def setNextBlock(self, block):
        if block >= 2 ** 16:
            log.debug("Block number rollover to 0 again")
            block = 0
        self.__eblock = block

    def getNextBlock(self):
        return self.__eblock

    next_block = property(getNextBlock, setNextBlock)

    def cycle(self):
        """Here we wait for a response from the server after sending it
        something, and dispatch appropriate action to that response."""
        try:
            (buffer, (raddress, rport)) = self.sock.recvfrom(MAX_BLKSIZE)
        except socket.timeout:
            log.warning("Timeout waiting for traffic, retrying...")
            raise TftpTimeout("Timed-out waiting for traffic")

        # Ok, we've received a packet. Log it.
        log.debug("Received %d bytes from %s:%s",
                        len(buffer), raddress, rport)
        # And update our last updated time.
        self.last_update = time.time()

        # Decode it.
        recvpkt = self.factory.parse(buffer)

        # Check for known "connection".
        if raddress != self.address:
            log.warning("Received traffic from %s, expected host %s. Discarding"
                        % (raddress, self.host))

        if self.tidport and self.tidport != rport:
            log.warning("Received traffic from %s:%s but we're "
                        "connected to %s:%s. Discarding."
                        % (raddress, rport,
                        self.host, self.tidport))

        # If there is a packethook defined, call it. We unconditionally
        # pass all packets, it's up to the client to screen out different
        # kinds of packets. This way, the client is privy to things like
        # negotiated options.
        if self.packethook:
            self.packethook(recvpkt)

        # And handle it, possibly changing state.
        self.state = self.state.handle(recvpkt, raddress, rport)
        # If we didn't throw any exceptions here, reset the retry_count to
        # zero.
        self.retry_count = 0

class TftpContextServer(TftpContext):
    """The context for the server."""
    def __init__(self,
                 host,
                 port,
                 timeout,
                 root,
                 dyn_file_func=None,
                 upload_open=None,
                 retries=DEF_TIMEOUT_RETRIES):
        TftpContext.__init__(self,
                             host,
                             port,
                             timeout,
                             retries
                             )
        # At this point we have no idea if this is a download or an upload. We
        # need to let the start state determine that.
        self.state = TftpStateServerStart(self)

        self.root = root
        self.dyn_file_func = dyn_file_func
        self.upload_open = upload_open

    def __str__(self):
        return "%s:%s %s" % (self.host, self.port, self.state)

    def start(self, buffer):
        """Start the state cycle. Note that the server context receives an
        initial packet in its start method. Also note that the server does not
        loop on cycle(), as it expects the TftpServer object to manage
        that."""
        log.debug("In TftpContextServer.start")
        self.metrics.start_time = time.time()
        log.debug("Set metrics.start_time to %s", self.metrics.start_time)
        # And update our last updated time.
        self.last_update = time.time()

        pkt = self.factory.parse(buffer)
        log.debug("TftpContextServer.start() - factory returned a %s", pkt)

        # Call handle once with the initial packet. This should put us into
        # the download or the upload state.
        self.state = self.state.handle(pkt,
                                       self.host,
                                       self.port)

    def end(self):
        """Finish up the context."""
        TftpContext.end(self)
        self.metrics.end_time = time.time()
        log.debug("Set metrics.end_time to %s", self.metrics.end_time)
        self.metrics.compute()

class TftpContextClientUpload(TftpContext):
    """The upload context for the client during an upload.
    Note: If input is a hyphen, then we will use stdin."""
    def __init__(self,
                 host,
                 port,
                 filename,
                 input,
                 options,
                 packethook,
                 timeout,
                 retries=DEF_TIMEOUT_RETRIES,
                 localip = ""):
        TftpContext.__init__(self,
                             host,
                             port,
                             timeout,
                             retries,
                             localip)
        self.file_to_transfer = filename
        self.options = options
        self.packethook = packethook
        # If the input object has a read() function,
        # assume it is file-like.
        if hasattr(input, 'read'):
            self.fileobj = input
        elif input == '-':
            self.fileobj = compat.binary_stdin()
        else:
            self.fileobj = open(input, "rb")

        log.debug("TftpContextClientUpload.__init__()")
        log.debug("file_to_transfer = %s, options = %s" %
            (self.file_to_transfer, self.options))

    def __str__(self):
        return "%s:%s %s" % (self.host, self.port, self.state)

    def start(self):
        log.info("Sending tftp upload request to %s" % self.host)
        log.info("    filename -> %s" % self.file_to_transfer)
        log.info("    options -> %s" % self.options)

        self.metrics.start_time = time.time()
        log.debug("Set metrics.start_time to %s" % self.metrics.start_time)

        # FIXME: put this in a sendWRQ method?
        pkt = TftpPacketWRQ()
        pkt.filename = self.file_to_transfer
        pkt.mode = "octet" # FIXME - shouldn't hardcode this
        pkt.options = self.options
        self.sock.sendto(pkt.encode().buffer, (self.host, self.port))
        self.next_block = 1
        self.last_pkt = pkt
        # FIXME: should we centralize sendto operations so we can refactor all
        # saving of the packet to the last_pkt field?

        self.state = TftpStateSentWRQ(self)

        while self.state:
            try:
                log.debug("State is %s" % self.state)
                self.cycle()
            except TftpTimeout as err:
                log.error(str(err))
                self.retry_count += 1
                if self.retry_count >= self.retries:
                    log.debug("hit max retries, giving up")
                    raise
                else:
                    log.warning("resending last packet")
                    self.state.resendLast()

    def end(self):
        """Finish up the context."""
        TftpContext.end(self)
        self.metrics.end_time = time.time()
        log.debug("Set metrics.end_time to %s" % self.metrics.end_time)
        self.metrics.compute()

class TftpContextClientDownload(TftpContext):
    """The download context for the client during a download.
    Note: If output is a hyphen, then the output will be sent to stdout."""
    def __init__(self,
                 host,
                 port,
                 filename,
                 output,
                 options,
                 packethook,
                 timeout,
                 retries=DEF_TIMEOUT_RETRIES,
                 localip = ""):
        TftpContext.__init__(self,
                             host,
                             port,
                             timeout,
                             retries,
                             localip)
        # FIXME: should we refactor setting of these params?
        self.file_to_transfer = filename
        self.options = options
        self.packethook = packethook
        self.filelike_fileobj = False
        # If the output object has a write() function,
        # assume it is file-like.
        if hasattr(output, 'write'):
            self.fileobj = output
            self.filelike_fileobj = True
        # If the output filename is -, then use stdout
        elif output == '-':
            self.fileobj = sys.stdout
            self.filelike_fileobj = True
        else:
            self.fileobj = open(output, "wb")

        log.debug("TftpContextClientDownload.__init__()")
        log.debug("file_to_transfer = %s, options = %s" %
            (self.file_to_transfer, self.options))

    def __str__(self):
        return "%s:%s %s" % (self.host, self.port, self.state)

    def start(self):
        """Initiate the download."""
        log.info("Sending tftp download request to %s" % self.host)
        log.info("    filename -> %s" % self.file_to_transfer)
        log.info("    options -> %s" % self.options)

        self.metrics.start_time = time.time()
        log.debug("Set metrics.start_time to %s" % self.metrics.start_time)

        # FIXME: put this in a sendRRQ method?
        pkt = TftpPacketRRQ()
        pkt.filename = self.file_to_transfer
        pkt.mode = "octet" # FIXME - shouldn't hardcode this
        pkt.options = self.options
        self.sock.sendto(pkt.encode().buffer, (self.host, self.port))
        self.next_block = 1
        self.last_pkt = pkt

        self.state = TftpStateSentRRQ(self)

        while self.state:
            try:
                log.debug("State is %s" % self.state)
                self.cycle()
            except TftpTimeout as err:
                log.error(str(err))
                self.retry_count += 1
                if self.retry_count >= self.retries:
                    log.debug("hit max retries, giving up")
                    raise
                else:
                    log.warning("resending last packet")
                    self.state.resendLast()
            except TftpFileNotFoundError as err:
                # If we received file not found, then we should not save the open
                # output file or we'll be left with a size zero file. Delete it,
                # if it exists.
                log.error("Received File not found error")
                if self.fileobj is not None and not self.filelike_fileobj:
                    if os.path.exists(self.fileobj.name):
                        log.debug("unlinking output file of %s", self.fileobj.name)
                        os.unlink(self.fileobj.name)

                raise

    def end(self):
        """Finish up the context."""
        TftpContext.end(self, not self.filelike_fileobj)
        self.metrics.end_time = time.time()
        log.debug("Set metrics.end_time to %s" % self.metrics.end_time)
        self.metrics.compute()
<commit_msg>Fix typos<commit_after># vim: ts=4 sw=4 et ai:
# -*- coding: utf8 -*-
"""This module implements all contexts for state handling during uploads and
downloads, the main interface to which being the TftpContext base class.

The concept is simple. Each context object represents a single upload or
download, and the state object in the context object represents the current
state of that transfer. The state object has a handle() method that expects
the next packet in the transfer, and returns a state object until the transfer
is complete, at which point it returns None. That is, unless there is a fatal
error, in which case a TftpException is returned instead."""

from .TftpShared import *
from .TftpPacketTypes import *
from .TftpPacketFactory import TftpPacketFactory
from .TftpStates import *
from . import compat
import socket
import time
import sys
import os
import logging

log = logging.getLogger('tftpy.TftpContext')

###############################################################################
# Utility classes
###############################################################################

class TftpMetrics(object):
    """A class representing metrics of the transfer."""
    def __init__(self):
        # Bytes transferred
        self.bytes = 0
        # Bytes re-sent
        self.resent_bytes = 0
        # Duplicate packets received
        self.dups = {}
        self.dupcount = 0
        # Times
        self.start_time = 0
        self.end_time = 0
        self.duration = 0
        # Rates
        self.bps = 0
        self.kbps = 0
        # Generic errors
        self.errors = 0

    def compute(self):
        # Compute transfer time
        self.duration = self.end_time - self.start_time
        if self.duration == 0:
            self.duration = 1
        log.debug("TftpMetrics.compute: duration is %s", self.duration)
        self.bps = (self.bytes * 8.0) / self.duration
        self.kbps = self.bps / 1024.0
        log.debug("TftpMetrics.compute: kbps is %s", self.kbps)
        for key in self.dups:
            self.dupcount += self.dups[key]

    def add_dup(self, pkt):
        """This method adds a dup for a packet to the metrics."""
        log.debug("Recording a dup of %s", pkt)
        s = str(pkt)
        if s in self.dups:
            self.dups[s] += 1
        else:
            self.dups[s] = 1
        tftpassert(self.dups[s] < MAX_DUPS, "Max duplicates reached")

###############################################################################
# Context classes
###############################################################################

class TftpContext(object):
    """The base class of the contexts."""

    def __init__(self, host, port, timeout, retries=DEF_TIMEOUT_RETRIES, localip = ""):
        """Constructor for the base context, setting shared instance
        variables."""
        self.file_to_transfer = None
        self.fileobj = None
        self.options = None
        self.packethook = None
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        if localip != "":
            self.sock.bind((localip, 0))
        self.sock.settimeout(timeout)
        self.timeout = timeout
        self.retries = retries
        self.state = None
        self.next_block = 0
        self.factory = TftpPacketFactory()
        # Note, setting the host will also set self.address, as it's a property.
        self.host = host
        self.port = port
        # The port associated with the TID
        self.tidport = None
        # Metrics
        self.metrics = TftpMetrics()
        # Fluag when the transfer is pending completion.
        self.pending_complete = False
        # Time when this context last received any traffic.
        # FIXME: does this belong in metrics?
        self.last_update = 0
        # The last packet we sent, if applicable, to make resending easy.
        self.last_pkt = None
        # Count the number of retry attempts.
        self.retry_count = 0

    def getBlocksize(self):
        """Fetch the current blocksize for this session."""
        return int(self.options.get('blksize', 512))

    def __del__(self):
        """Simple destructor to try to call housekeeping in the end method if
        not called explicitly. Leaking file descriptors is not a good
        thing."""
        self.end()

    def checkTimeout(self, now):
        """Compare current time with last_update time, and raise an exception
        if we're over the timeout time."""
        log.debug("checking for timeout on session %s", self)
        if now - self.last_update > self.timeout:
            raise TftpTimeout("Timeout waiting for traffic")

    def start(self):
        raise NotImplementedError("Abstract method")

    def end(self, close_fileobj=True):
        """Perform session cleanup, since the end method should always be
        called explicitly by the calling code, this works better than the
        destructor.
        Set close_fileobj to False so fileobj can be returned open."""
        log.debug("in TftpContext.end - closing socket")
        self.sock.close()
        if close_fileobj and self.fileobj is not None and not self.fileobj.closed:
            log.debug("self.fileobj is open - closing")
            self.fileobj.close()

    def gethost(self):
        "Simple getter method for use in a property."
        return self.__host

    def sethost(self, host):
        """Setter method that also sets the address property as a result
        of the host that is set."""
        self.__host = host
        self.address = socket.gethostbyname(host)

    host = property(gethost, sethost)

    def setNextBlock(self, block):
        if block >= 2 ** 16:
            log.debug("Block number rollover to 0 again")
            block = 0
        self.__eblock = block

    def getNextBlock(self):
        return self.__eblock

    next_block = property(getNextBlock, setNextBlock)

    def cycle(self):
        """Here we wait for a response from the server after sending it
        something, and dispatch appropriate action to that response."""
        try:
            (buffer, (raddress, rport)) = self.sock.recvfrom(MAX_BLKSIZE)
        except socket.timeout:
            log.warning("Timeout waiting for traffic, retrying...")
            raise TftpTimeout("Timed-out waiting for traffic")

        # Ok, we've received a packet. Log it.
        log.debug("Received %d bytes from %s:%s",
                        len(buffer), raddress, rport)
        # And update our last updated time.
        self.last_update = time.time()

        # Decode it.
        recvpkt = self.factory.parse(buffer)

        # Check for known "connection".
        if raddress != self.address:
            log.warning("Received traffic from %s, expected host %s. Discarding"
                        % (raddress, self.host))

        if self.tidport and self.tidport != rport:
            log.warning("Received traffic from %s:%s but we're "
                        "connected to %s:%s. Discarding."
                        % (raddress, rport,
                        self.host, self.tidport))

        # If there is a packethook defined, call it. We unconditionally
        # pass all packets, it's up to the client to screen out different
        # kinds of packets. This way, the client is privy to things like
        # negotiated options.
        if self.packethook:
            self.packethook(recvpkt)

        # And handle it, possibly changing state.
        self.state = self.state.handle(recvpkt, raddress, rport)
        # If we didn't throw any exceptions here, reset the retry_count to
        # zero.
        self.retry_count = 0

class TftpContextServer(TftpContext):
    """The context for the server."""
    def __init__(self,
                 host,
                 port,
                 timeout,
                 root,
                 dyn_file_func=None,
                 upload_open=None,
                 retries=DEF_TIMEOUT_RETRIES):
        TftpContext.__init__(self,
                             host,
                             port,
                             timeout,
                             retries
                             )
        # At this point we have no idea if this is a download or an upload. We
        # need to let the start state determine that.
        self.state = TftpStateServerStart(self)

        self.root = root
        self.dyn_file_func = dyn_file_func
        self.upload_open = upload_open

    def __str__(self):
        return "%s:%s %s" % (self.host, self.port, self.state)

    def start(self, buffer):
        """Start the state cycle. Note that the server context receives an
        initial packet in its start method. Also note that the server does not
        loop on cycle(), as it expects the TftpServer object to manage
        that."""
        log.debug("In TftpContextServer.start")
        self.metrics.start_time = time.time()
        log.debug("Set metrics.start_time to %s", self.metrics.start_time)
        # And update our last updated time.
        self.last_update = time.time()

        pkt = self.factory.parse(buffer)
        log.debug("TftpContextServer.start() - factory returned a %s", pkt)

        # Call handle once with the initial packet. This should put us into
        # the download or the upload state.
        self.state = self.state.handle(pkt,
                                       self.host,
                                       self.port)

    def end(self):
        """Finish up the context."""
        TftpContext.end(self)
        self.metrics.end_time = time.time()
        log.debug("Set metrics.end_time to %s", self.metrics.end_time)
        self.metrics.compute()

class TftpContextClientUpload(TftpContext):
    """The upload context for the client during an upload.
    Note: If input is a hyphen, then we will use stdin."""
    def __init__(self,
                 host,
                 port,
                 filename,
                 input,
                 options,
                 packethook,
                 timeout,
                 retries=DEF_TIMEOUT_RETRIES,
                 localip = ""):
        TftpContext.__init__(self,
                             host,
                             port,
                             timeout,
                             retries,
                             localip)
        self.file_to_transfer = filename
        self.options = options
        self.packethook = packethook
        # If the input object has a read() function,
        # assume it is file-like.
        if hasattr(input, 'read'):
            self.fileobj = input
        elif input == '-':
            self.fileobj = compat.binary_stdin()
        else:
            self.fileobj = open(input, "rb")

        log.debug("TftpContextClientUpload.__init__()")
        log.debug("file_to_transfer = %s, options = %s" %
            (self.file_to_transfer, self.options))

    def __str__(self):
        return "%s:%s %s" % (self.host, self.port, self.state)

    def start(self):
        log.info("Sending tftp upload request to %s" % self.host)
        log.info("    filename -> %s" % self.file_to_transfer)
        log.info("    options -> %s" % self.options)

        self.metrics.start_time = time.time()
        log.debug("Set metrics.start_time to %s" % self.metrics.start_time)

        # FIXME: put this in a sendWRQ method?
        pkt = TftpPacketWRQ()
        pkt.filename = self.file_to_transfer
        pkt.mode = "octet" # FIXME - shouldn't hardcode this
        pkt.options = self.options
        self.sock.sendto(pkt.encode().buffer, (self.host, self.port))
        self.next_block = 1
        self.last_pkt = pkt
        # FIXME: should we centralize sendto operations so we can refactor all
        # saving of the packet to the last_pkt field?

        self.state = TftpStateSentWRQ(self)

        while self.state:
            try:
                log.debug("State is %s" % self.state)
                self.cycle()
            except TftpTimeout as err:
                log.error(str(err))
                self.retry_count += 1
                if self.retry_count >= self.retries:
                    log.debug("hit max retries, giving up")
                    raise
                else:
                    log.warning("resending last packet")
                    self.state.resendLast()

    def end(self):
        """Finish up the context."""
        TftpContext.end(self)
        self.metrics.end_time = time.time()
        log.debug("Set metrics.end_time to %s" % self.metrics.end_time)
        self.metrics.compute()

class TftpContextClientDownload(TftpContext):
    """The download context for the client during a download.
    Note: If output is a hyphen, then the output will be sent to stdout."""
    def __init__(self,
                 host,
                 port,
                 filename,
                 output,
                 options,
                 packethook,
                 timeout,
                 retries=DEF_TIMEOUT_RETRIES,
                 localip = ""):
        TftpContext.__init__(self,
                             host,
                             port,
                             timeout,
                             retries,
                             localip)
        # FIXME: should we refactor setting of these params?
        self.file_to_transfer = filename
        self.options = options
        self.packethook = packethook
        self.filelike_fileobj = False
        # If the output object has a write() function,
        # assume it is file-like.
        if hasattr(output, 'write'):
            self.fileobj = output
            self.filelike_fileobj = True
        # If the output filename is -, then use stdout
        elif output == '-':
            self.fileobj = sys.stdout
            self.filelike_fileobj = True
        else:
            self.fileobj = open(output, "wb")

        log.debug("TftpContextClientDownload.__init__()")
        log.debug("file_to_transfer = %s, options = %s" %
            (self.file_to_transfer, self.options))

    def __str__(self):
        return "%s:%s %s" % (self.host, self.port, self.state)

    def start(self):
        """Initiate the download."""
        log.info("Sending tftp download request to %s" % self.host)
        log.info("    filename -> %s" % self.file_to_transfer)
        log.info("    options -> %s" % self.options)

        self.metrics.start_time = time.time()
        log.debug("Set metrics.start_time to %s" % self.metrics.start_time)

        # FIXME: put this in a sendRRQ method?
        pkt = TftpPacketRRQ()
        pkt.filename = self.file_to_transfer
        pkt.mode = "octet" # FIXME - shouldn't hardcode this
        pkt.options = self.options
        self.sock.sendto(pkt.encode().buffer, (self.host, self.port))
        self.next_block = 1
        self.last_pkt = pkt

        self.state = TftpStateSentRRQ(self)

        while self.state:
            try:
                log.debug("State is %s" % self.state)
                self.cycle()
            except TftpTimeout as err:
                log.error(str(err))
                self.retry_count += 1
                if self.retry_count >= self.retries:
                    log.debug("hit max retries, giving up")
                    raise
                else:
                    log.warning("resending last packet")
                    self.state.resendLast()
            except TftpFileNotFoundError as err:
                # If we received file not found, then we should not save the open
                # output file or we'll be left with a size zero file. Delete it,
                # if it exists.
                log.error("Received File not found error")
                if self.fileobj is not None and not self.filelike_fileobj:
                    if os.path.exists(self.fileobj.name):
                        log.debug("unlinking output file of %s", self.fileobj.name)
                        os.unlink(self.fileobj.name)

                raise

    def end(self):
        """Finish up the context."""
        TftpContext.end(self, not self.filelike_fileobj)
        self.metrics.end_time = time.time()
        log.debug("Set metrics.end_time to %s" % self.metrics.end_time)
        self.metrics.compute()
<|endoftext|>