<commit_before>import pytest
from pytorch_lightning import Trainer
from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
from pytorch_lightning.testing_models.lm_test_module import LightningTestModel
from argparse import Namespace
from test_tube import Experiment
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from pytorch_lightning.utils.debugging import MisconfigurationException
from pytorch_lightning.root_module import memory
import numpy as np
import warnings
import torch
import os
import shutil

SEED = 2334
torch.manual_seed(SEED)
np.random.seed(SEED)

# ------------------------------------------------------------------------
# TESTS
# ------------------------------------------------------------------------

def test_cpu_model_with_amp():
    """
    Make sure model trains on CPU
    :return:
    """

    trainer_options = dict(
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4,
        use_amp=True
    )

    model, hparams = get_model()

    with pytest.raises(MisconfigurationException):
        run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)

def test_amp_gpu_ddp_slurm_managed():
    """
    Make sure DDP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])

    hparams = get_hparams()
    model = LightningTestModel(hparams)

    trainer_options = dict(
        progress_bar=True,
        max_nb_epochs=1,
        gpus=[0],
        distributed_backend='ddp',
        use_amp=True
    )

    save_dir = init_save_dir()

    # exp file to get meta
    exp = get_exp(False)
    exp.argparse(hparams)
    exp.save()

    # exp file to get weights
    checkpoint = ModelCheckpoint(save_dir)

    # add these to the trainer options
    trainer_options['checkpoint_callback'] = checkpoint
    trainer_options['experiment'] = exp

    # fit model
    trainer = Trainer(**trainer_options)
    trainer.is_slurm_managing_tasks = True
    result = trainer.fit(model)

    # correct result and ok accuracy
    assert result == 1, 'amp + ddp model failed to complete'

    # test root model address
    assert trainer.resolve_root_node_address('abc') == 'abc'
    assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
    assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
    assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'

    # test model loading
    pretrained_model = load_model(exp, save_dir, True)

    # test model preds
    run_prediction(model.test_dataloader, pretrained_model)

    if trainer.use_ddp:
        # on hpc this would work fine... but need to hack it for the purpose of the test
        trainer.model = pretrained_model
        trainer.optimizers = pretrained_model.configure_optimizers()

    # test HPC loading / saving
    trainer.hpc_save(save_dir, exp)
    trainer.hpc_load(save_dir, on_gpu=True)

    clear_save_dir()

def test_cpu_model():
    """
    Make sure model trains on CPU
    :return:
    """

    trainer_options = dict(
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4
    )

    model, hparams = get_model()

    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)

def test_all_features_cpu_model():
    """
    Test each of the trainer options
    :return:
    """

    trainer_options = dict(
        gradient_clip=1.0,
        overfit_pct=0.20,
        track_grad_norm=2,
        print_nan_grads=True,
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4
    )

    model, hparams = get_model()
    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)

def test_early_stopping_cpu_model():
    """
    Test each of the trainer options
    :return:
    """

    stopping = EarlyStopping()
    trainer_options = dict(
        early_stop_callback=stopping,
        gradient_clip=1.0,
        overfit_pct=0.20,
        track_grad_norm=2,
        print_nan_grads=True,
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4
    )

    model, hparams = get_model()
    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)

def test_single_gpu_model():
    """
    Make sure single GPU works (DP mode)
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
        return
    model, hparams = get_model()

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus=[0]
    )

    run_gpu_model_test(trainer_options, model, hparams)

def test_multi_gpu_model_dp():
    """
    Make sure DP works
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return
    model, hparams = get_model()
    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus='-1'
    )

    run_gpu_model_test(trainer_options, model, hparams)

    # test memory helper functions
    memory.get_gpu_memory_map()

def test_amp_gpu_dp():
    """
    Make sure DP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return
    model, hparams = get_model()
    trainer_options = dict(
        max_nb_epochs=1,
        gpus='0, 1',  # test init with gpu string
        distributed_backend='dp',
        use_amp=True
    )
    with pytest.raises(MisconfigurationException):
        run_gpu_model_test(trainer_options, model, hparams)

def test_multi_gpu_model_ddp():
    """
    Make sure DDP works
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
    model, hparams = get_model()
    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.2,
        gpus=[0, 1],
        distributed_backend='ddp'
    )

    run_gpu_model_test(trainer_options, model, hparams)

def test_amp_gpu_ddp():
    """
    Make sure DDP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])

    hparams = get_hparams()
    model = LightningTestModel(hparams)

    trainer_options = dict(
        progress_bar=True,
        max_nb_epochs=1,
        gpus=[0, 1],
        distributed_backend='ddp',
        use_amp=True
    )

    run_gpu_model_test(trainer_options, model, hparams)

def test_ddp_sampler_error():
    """
    Make sure DDP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])

    hparams = get_hparams()
    model = LightningTestModel(hparams, force_remove_distributed_sampler=True)

    exp = get_exp(True)
    exp.save()

    trainer = Trainer(
        experiment=exp,
        progress_bar=False,
        max_nb_epochs=1,
        gpus=[0, 1],
        distributed_backend='ddp',
        use_amp=True
    )

    with pytest.raises(MisconfigurationException):
        trainer.get_dataloaders(model)

    clear_save_dir()

# ------------------------------------------------------------------------
# UTILS
# ------------------------------------------------------------------------
def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
    save_dir = init_save_dir()

    # exp file to get meta
    exp = get_exp(False)
    exp.argparse(hparams)
    exp.save()

    # exp file to get weights
    checkpoint = ModelCheckpoint(save_dir)

    # add these to the trainer options
    trainer_options['checkpoint_callback'] = checkpoint
    trainer_options['experiment'] = exp

    # fit model
    trainer = Trainer(**trainer_options)
    result = trainer.fit(model)

    # correct result and ok accuracy
    assert result == 1, 'amp + ddp model failed to complete'

    # test model loading
    pretrained_model = load_model(exp, save_dir, on_gpu)

    # test model preds
    run_prediction(model.test_dataloader, pretrained_model)

    if trainer.use_ddp:
        # on hpc this would work fine... but need to hack it for the purpose of the test
        trainer.model = pretrained_model
        trainer.optimizers = pretrained_model.configure_optimizers()

    # test HPC loading / saving
    trainer.hpc_save(save_dir, exp)
    trainer.hpc_load(save_dir, on_gpu=on_gpu)

    clear_save_dir()

def get_hparams():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    hparams = Namespace(**{'drop_prob': 0.2,
                           'batch_size': 32,
                           'in_features': 28*28,
                           'learning_rate': 0.001*8,
                           'optimizer_name': 'adam',
                           'data_root': os.path.join(root_dir, 'mnist'),
                           'out_features': 10,
                           'hidden_dim': 1000})
    return hparams

def get_model():
    # set up model with these hyperparams
    hparams = get_hparams()
    model = LightningTemplateModel(hparams)

    return model, hparams

def get_exp(debug=True):
    # set up exp object without actually saving logs
    root_dir = os.path.dirname(os.path.realpath(__file__))
    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
    return exp

def init_save_dir():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    save_dir = os.path.join(root_dir, 'save_dir')

    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)

    os.makedirs(save_dir, exist_ok=True)

    return save_dir

def clear_save_dir():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    save_dir = os.path.join(root_dir, 'save_dir')
    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)

def load_model(exp, save_dir, on_gpu):

    # load trained model
    tags_path = exp.get_data_path(exp.name, exp.version)
    tags_path = os.path.join(tags_path, 'meta_tags.csv')

    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
    weights_dir = os.path.join(save_dir, checkpoints[0])

    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=on_gpu)

    assert trained_model is not None, 'loading model failed'

    return trained_model

def run_prediction(dataloader, trained_model):
    # run prediction on 1 batch
    for batch in dataloader:
        break

    x, y = batch
    x = x.view(x.size(0), -1)

    y_hat = trained_model(x)

    # acc
    labels_hat = torch.argmax(y_hat, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
    val_acc = torch.tensor(val_acc)
    val_acc = val_acc.item()

    print(val_acc)

    assert val_acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {val_acc})'

def assert_ok_acc(trainer):
    # this model should get 0.80+ acc
    acc = trainer.tng_tqdm_dic['val_acc']
    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'

if __name__ == '__main__':
    pytest.main([__file__])
<commit_msg>added cpu + amp error<commit_after>import pytest
from pytorch_lightning import Trainer
from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
from pytorch_lightning.testing_models.lm_test_module import LightningTestModel
from argparse import Namespace
from test_tube import Experiment
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from pytorch_lightning.utils.debugging import MisconfigurationException
from pytorch_lightning.root_module import memory
import numpy as np
import warnings
import torch
import os
import shutil

SEED = 2334
torch.manual_seed(SEED)
np.random.seed(SEED)

# ------------------------------------------------------------------------
# TESTS
# ------------------------------------------------------------------------
def test_early_stopping_cpu_model():
    """
    Test each of the trainer options
    :return:
    """

    stopping = EarlyStopping()
    trainer_options = dict(
        early_stop_callback=stopping,
        gradient_clip=1.0,
        overfit_pct=0.20,
        track_grad_norm=2,
        print_nan_grads=True,
        progress_bar=False,
        experiment=get_exp(),
        train_percent_check=0.1,
        val_percent_check=0.1
    )

    model, hparams = get_model()
    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)

def test_cpu_model_with_amp():
    """
    Make sure model trains on CPU
    :return:
    """

    trainer_options = dict(
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4,
        use_amp=True
    )

    model, hparams = get_model()

    with pytest.raises(MisconfigurationException):
        run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)

def test_amp_gpu_ddp_slurm_managed():
    """
    Make sure DDP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])

    hparams = get_hparams()
    model = LightningTestModel(hparams)

    trainer_options = dict(
        progress_bar=True,
        max_nb_epochs=1,
        gpus=[0],
        distributed_backend='ddp',
        use_amp=True
    )

    save_dir = init_save_dir()

    # exp file to get meta
    exp = get_exp(False)
    exp.argparse(hparams)
    exp.save()

    # exp file to get weights
    checkpoint = ModelCheckpoint(save_dir)

    # add these to the trainer options
    trainer_options['checkpoint_callback'] = checkpoint
    trainer_options['experiment'] = exp

    # fit model
    trainer = Trainer(**trainer_options)
    trainer.is_slurm_managing_tasks = True
    result = trainer.fit(model)

    # correct result and ok accuracy
    assert result == 1, 'amp + ddp model failed to complete'

    # test root model address
    assert trainer.resolve_root_node_address('abc') == 'abc'
    assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
    assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
    assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'

    # test model loading
    pretrained_model = load_model(exp, save_dir, True)

    # test model preds
    run_prediction(model.test_dataloader, pretrained_model)

    if trainer.use_ddp:
        # on hpc this would work fine... but need to hack it for the purpose of the test
        trainer.model = pretrained_model
        trainer.optimizers = pretrained_model.configure_optimizers()

    # test HPC loading / saving
    trainer.hpc_save(save_dir, exp)
    trainer.hpc_load(save_dir, on_gpu=True)

    clear_save_dir()

def test_cpu_model():
    """
    Make sure model trains on CPU
    :return:
    """

    trainer_options = dict(
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4
    )

    model, hparams = get_model()

    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)

def test_all_features_cpu_model():
    """
    Test each of the trainer options
    :return:
    """

    trainer_options = dict(
        gradient_clip=1.0,
        overfit_pct=0.20,
        track_grad_norm=2,
        print_nan_grads=True,
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4
    )

    model, hparams = get_model()
    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)

def test_single_gpu_model():
    """
    Make sure single GPU works (DP mode)
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
        return
    model, hparams = get_model()

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus=[0]
    )

    run_gpu_model_test(trainer_options, model, hparams)

def test_multi_gpu_model_dp():
    """
    Make sure DP works
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return
    model, hparams = get_model()
    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus='-1'
    )

    run_gpu_model_test(trainer_options, model, hparams)

    # test memory helper functions
    memory.get_gpu_memory_map()

def test_amp_gpu_dp():
    """
    Make sure DP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return
    model, hparams = get_model()
    trainer_options = dict(
        max_nb_epochs=1,
        gpus='0, 1',  # test init with gpu string
        distributed_backend='dp',
        use_amp=True
    )
    with pytest.raises(MisconfigurationException):
        run_gpu_model_test(trainer_options, model, hparams)

def test_multi_gpu_model_ddp():
    """
    Make sure DDP works
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
    model, hparams = get_model()
    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.2,
        gpus=[0, 1],
        distributed_backend='ddp'
    )

    run_gpu_model_test(trainer_options, model, hparams)

def test_amp_gpu_ddp():
    """
    Make sure DDP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])

    hparams = get_hparams()
    model = LightningTestModel(hparams)

    trainer_options = dict(
        progress_bar=True,
        max_nb_epochs=1,
        gpus=[0, 1],
        distributed_backend='ddp',
        use_amp=True
    )

    run_gpu_model_test(trainer_options, model, hparams)

def test_ddp_sampler_error():
    """
    Make sure DDP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])

    hparams = get_hparams()
    model = LightningTestModel(hparams, force_remove_distributed_sampler=True)

    exp = get_exp(True)
    exp.save()

    trainer = Trainer(
        experiment=exp,
        progress_bar=False,
        max_nb_epochs=1,
        gpus=[0, 1],
        distributed_backend='ddp',
        use_amp=True
    )

    with pytest.raises(MisconfigurationException):
        trainer.get_dataloaders(model)

    clear_save_dir()

# ------------------------------------------------------------------------
# UTILS
# ------------------------------------------------------------------------
def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
    save_dir = init_save_dir()

    # exp file to get meta
    exp = get_exp(False)
    exp.argparse(hparams)
    exp.save()

    # exp file to get weights
    checkpoint = ModelCheckpoint(save_dir)

    # add these to the trainer options
    trainer_options['checkpoint_callback'] = checkpoint
    trainer_options['experiment'] = exp

    # fit model
    trainer = Trainer(**trainer_options)
    result = trainer.fit(model)

    # correct result and ok accuracy
    assert result == 1, 'amp + ddp model failed to complete'

    # test model loading
    pretrained_model = load_model(exp, save_dir, on_gpu)

    # test model preds
    run_prediction(model.test_dataloader, pretrained_model)

    if trainer.use_ddp:
        # on hpc this would work fine... but need to hack it for the purpose of the test
        trainer.model = pretrained_model
        trainer.optimizers = pretrained_model.configure_optimizers()

    # test HPC loading / saving
    trainer.hpc_save(save_dir, exp)
    trainer.hpc_load(save_dir, on_gpu=on_gpu)

    clear_save_dir()

def get_hparams():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    hparams = Namespace(**{'drop_prob': 0.2,
                           'batch_size': 32,
                           'in_features': 28*28,
                           'learning_rate': 0.001*8,
                           'optimizer_name': 'adam',
                           'data_root': os.path.join(root_dir, 'mnist'),
                           'out_features': 10,
                           'hidden_dim': 1000})
    return hparams

def get_model():
    # set up model with these hyperparams
    hparams = get_hparams()
    model = LightningTemplateModel(hparams)

    return model, hparams

def get_exp(debug=True):
    # set up exp object without actually saving logs
    root_dir = os.path.dirname(os.path.realpath(__file__))
    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
    return exp

def init_save_dir():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    save_dir = os.path.join(root_dir, 'save_dir')

    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)

    os.makedirs(save_dir, exist_ok=True)

    return save_dir

def clear_save_dir():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    save_dir = os.path.join(root_dir, 'save_dir')
    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)

def load_model(exp, save_dir, on_gpu):

    # load trained model
    tags_path = exp.get_data_path(exp.name, exp.version)
    tags_path = os.path.join(tags_path, 'meta_tags.csv')

    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
    weights_dir = os.path.join(save_dir, checkpoints[0])

    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=on_gpu)

    assert trained_model is not None, 'loading model failed'

    return trained_model

def run_prediction(dataloader, trained_model):
    # run prediction on 1 batch
    for batch in dataloader:
        break

    x, y = batch
    x = x.view(x.size(0), -1)

    y_hat = trained_model(x)

    # acc
    labels_hat = torch.argmax(y_hat, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
    val_acc = torch.tensor(val_acc)
    val_acc = val_acc.item()

    print(val_acc)

    assert val_acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {val_acc})'

def assert_ok_acc(trainer):
    # this model should get 0.80+ acc
    acc = trainer.tng_tqdm_dic['val_acc']
    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'

if __name__ == '__main__':
    pytest.main([__file__])
<|endoftext|>


<commit_before>import pytest
from django.contrib.sites.models import Site

from controlled_vocabularies.models import Vocabulary, Term, Property
from . import factories

pytestmark = pytest.mark.django_db

LINK = "{}"

class TestVocabulary:
    def test_save(self):
        vocab = factories.VocabularyFactory.build()
        assert Vocabulary.objects.all().count() == 0
        vocab.save()
        assert Vocabulary.objects.all().count() == 1

    def test_unicode(self):
        vocab = factories.VocabularyFactory()
        assert unicode(vocab) == vocab.name

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        factories.VocabularyFactory(
            name=' spaces  here ',
            label=' Spaces here   too  ',
            maintainer=' Please   Fix Me ',
            maintainerEmail='  here@there.com ',
            definition='  Strip leading   and trailing, condense   others. '
        )
        vocab = Vocabulary.objects.get(pk=1)

        assert vocab.name == 'spaces here'
        assert vocab.label == 'Spaces here too'
        assert vocab.maintainer == 'Please Fix Me'
        assert vocab.maintainerEmail == 'here@there.com'
        assert vocab.definition == 'Strip leading and trailing, condense others.'

class TestTerm:
    def test_unicode(self):
        term = factories.TermFactory()
        assert unicode(term) == term.name

    def test_save(self):
        term = factories.TermFactory.build(vocab_list=factories.VocabularyFactory())
        assert Term.objects.all().count() == 0
        term.save()
        assert Term.objects.all().count() == 1

    def test_get_vocab(self):
        term = factories.TermFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'vocabulary',
            term.vocab_list.id,
            term.vocab_list
        )
        assert term.get_vocab() == expected

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        factories.TermFactory(
            name=' spaces  here ',
            label=' Spaces here   too  ',
        )
        term = Term.objects.get(pk=1)

        assert term.name == 'spaces here'
        assert term.label == 'Spaces here too'

class TestProperty:
    def test_save(self):
        prop = factories.PropertyFactory.build(term_key=factories.TermFactory())
        assert Property.objects.all().count() == 0
        prop.save()
        assert Property.objects.all().count() == 1

    def test_get_vocab(self):
        prop = factories.PropertyFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'vocabulary',
            prop.term_key.vocab_list.id,
            prop.term_key.vocab_list
        )
        assert prop.get_vocab() == expected

    def test_get_term(self):
        prop = factories.PropertyFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'term',
            prop.term_key.id,
            prop.term_key
        )
        assert prop.get_term() == expected

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        factories.PropertyFactory(
            property_name=' spaces  here ',
            label=' Spaces here   too  ',
        )
        prop = Property.objects.get(pk=1)

        assert prop.property_name == 'spaces here'
        assert prop.label == 'Spaces here too'
<commit_msg>Removed hardcoded primary keys.<commit_after>import pytest
from django.contrib.sites.models import Site

from controlled_vocabularies.models import Vocabulary, Term, Property
from . import factories

pytestmark = pytest.mark.django_db

LINK = "{}"

class TestVocabulary:
    def test_save(self):
        vocab = factories.VocabularyFactory.build()
        assert Vocabulary.objects.all().count() == 0
        vocab.save()
        assert Vocabulary.objects.all().count() == 1

    def test_unicode(self):
        vocab = factories.VocabularyFactory()
        assert unicode(vocab) == vocab.name

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        vocab = factories.VocabularyFactory(
            name=' spaces  here ',
            label=' Spaces here   too  ',
            maintainer=' Please   Fix Me ',
            maintainerEmail='  here@there.com ',
            definition='  Strip leading   and trailing, condense   others. '
        )
        updated_vocab = Vocabulary.objects.get(pk=vocab.pk)

        assert updated_vocab.name == 'spaces here'
        assert updated_vocab.label == 'Spaces here too'
        assert updated_vocab.maintainer == 'Please Fix Me'
        assert updated_vocab.maintainerEmail == 'here@there.com'
        assert updated_vocab.definition == 'Strip leading and trailing, condense others.'

class TestTerm:
    def test_unicode(self):
        term = factories.TermFactory()
        assert unicode(term) == term.name

    def test_save(self):
        term = factories.TermFactory.build(vocab_list=factories.VocabularyFactory())
        assert Term.objects.all().count() == 0
        term.save()
        assert Term.objects.all().count() == 1

    def test_get_vocab(self):
        term = factories.TermFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'vocabulary',
            term.vocab_list.id,
            term.vocab_list
        )
        assert term.get_vocab() == expected

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        term = factories.TermFactory(
            name=' spaces  here ',
            label=' Spaces here   too  ',
        )
        updated_term = Term.objects.get(pk=term.pk)

        assert updated_term.name == 'spaces here'
        assert updated_term.label == 'Spaces here too'

class TestProperty:
    def test_save(self):
        prop = factories.PropertyFactory.build(term_key=factories.TermFactory())
        assert Property.objects.all().count() == 0
        prop.save()
        assert Property.objects.all().count() == 1

    def test_get_vocab(self):
        prop = factories.PropertyFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'vocabulary',
            prop.term_key.vocab_list.id,
            prop.term_key.vocab_list
        )
        assert prop.get_vocab() == expected

    def test_get_term(self):
        prop = factories.PropertyFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'term',
            prop.term_key.id,
            prop.term_key
        )
        assert prop.get_term() == expected

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        prop = factories.PropertyFactory(
            property_name=' spaces  here ',
            label=' Spaces here   too  ',
        )
        updated_prop = Property.objects.get(pk=prop)

        assert updated_prop.property_name == 'spaces here'
        assert updated_prop.label == 'Spaces here too'
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import unittest

from coaster.sqlalchemy import BaseMixin, BaseNameMixin, BaseScopedNameMixin, BaseIdNameMixin, BaseScopedIdNameMixin
from sqlalchemy import create_engine, Column, Integer, Unicode, UniqueConstraint, ForeignKey
from sqlalchemy.orm import relationship, synonym, sessionmaker, scoped_session
from sqlalchemy.ext.declarative import declarative_base

engine = create_engine('sqlite://')
Session = scoped_session(sessionmaker(autocommit=False,
              autoflush=False,
              bind=engine))
Base = declarative_base(bind=engine)

# --- Models ------------------------------------------------------------------

class Container(BaseMixin, Base):
    __tablename__ = 'container'
    query = Session.query_property()

class UnnamedDocument(BaseMixin, Base):
    __tablename__ = 'unnamed_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)

    content = Column(Unicode(250))

class NamedDocument(BaseNameMixin, Base):
    __tablename__ = 'named_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)

    content = Column(Unicode(250))

class ScopedNamedDocument(BaseScopedNameMixin, Base):
    __tablename__ = 'scoped_named_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)
    parent = synonym('container')

    content = Column(Unicode(250))
    __table_args__ = (UniqueConstraint('name', 'container_id'),)

class IdNamedDocument(BaseIdNameMixin, Base):
    __tablename__ = 'id_named_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)

    content = Column(Unicode(250))

class ScopedIdNamedDocument(BaseScopedIdNameMixin, Base):
    __tablename__ = 'scoped_id_named_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)
    parent = synonym('container')

    content = Column(Unicode(250))
    __table_args__ = (UniqueConstraint('url_id', 'container_id'),)

# -- Tests --------------------------------------------------------------------

class TestCoasterModels(unittest.TestCase):
    def setUp(self):
        Base.metadata.create_all()
        self.session = Session()

    def tearDown(self):
        self.session.rollback()
        Base.metadata.drop_all()

    def make_container(self):
        c = Container()
        self.session.add(c)
        return c

    def test_container(self):
        c = self.make_container()
        self.assertEqual(c.id, None)
        self.session.commit()
        self.assertEqual(c.id, 1)

    def test_unnamed(self):
        c = self.make_container()
        d = UnnamedDocument(content="hello", container=c)
        self.session.add(d)
        self.session.commit()
        self.assertEqual(c.id, 1)
        self.assertEqual(d.id, 1)

    def test_named(self):
        """Named documents have globally unique names."""
        c1 = self.make_container()
        self.session.add(c1)
        d1 = NamedDocument(title="Hello", content="World", container=c1)
        self.session.add(d1)
        self.session.commit()
        self.assertEqual(d1.name, 'hello')

        c2 = self.make_container()
        d2 = NamedDocument(title="Hello", content="Again", container=c2)
        self.session.add(d2)
        self.session.commit()
        self.assertEqual(d2.name, 'hello1')

    def test_scoped_named(self):
        """Scoped named documents have names unique to their containers."""
        c1 = self.make_container()
        self.session.add(c1)
        d1 = ScopedNamedDocument(title="Hello", content="World", container=c1)
        self.session.add(d1)
        self.session.commit()
        self.assertEqual(d1.name, 'hello')

        d2 = ScopedNamedDocument(title="Hello", content="Again", container=c1)
        self.session.add(d2)
        self.session.commit()
        self.assertEqual(d2.name, 'hello1')

        c2 = self.make_container()
        d3 = ScopedNamedDocument(title="Hello", content="Once More", container=c2)
        self.session.add(d3)
        self.session.commit()
        self.assertEqual(d3.name, 'hello')

    def test_id_named(self):
        """Documents with a global id in the URL"""
        c1 = self.make_container()
        self.session.add(c1)
        d1 = IdNamedDocument(title="Hello", content="World", container=c1)
        self.session.add(d1)
        self.session.commit()
        self.assertEqual(d1.url_name, '1-hello')

        d2 = IdNamedDocument(title="Hello", content="Again", container=c1)
        self.session.add(d2)
        self.session.commit()
        self.assertEqual(d2.url_name, '2-hello')

        c2 = self.make_container()
        d3 = IdNamedDocument(title="Hello", content="Once More", container=c2)
        self.session.add(d3)
        self.session.commit()
        self.assertEqual(d3.url_name, '3-hello')

    def test_scoped_id_named(self):
        """Documents with a container-specifc id in the URL"""
        c1 = self.make_container()
        self.session.add(c1)
        d1 = ScopedIdNamedDocument(title="Hello", content="World", container=c1)
        self.session.add(d1)
        self.session.commit()
        self.assertEqual(d1.url_name, '1-hello')

        d2 = ScopedIdNamedDocument(title="Hello again", content="New name", container=c1)
        self.session.add(d2)
        self.session.commit()
        self.assertEqual(d2.url_name, '2-hello-again')

        c2 = self.make_container()
        d3 = ScopedIdNamedDocument(title="Hello", content="Once More", container=c2)
        self.session.add(d3)
        self.session.commit()
        self.assertEqual(d3.url_name, '1-hello')

        d4 = ScopedIdNamedDocument(title="Hello", content="Third", container=c1)
        self.session.add(d4)
        self.session.commit()
        self.assertEqual(d4.url_name, '3-hello')

if __name__ == '__main__':
    unittest.main()
<commit_msg>Timestamp tests.<commit_after># -*- coding: utf-8 -*-

import unittest

from datetime import datetime
from coaster.sqlalchemy import BaseMixin, BaseNameMixin, BaseScopedNameMixin, BaseIdNameMixin, BaseScopedIdNameMixin
from sqlalchemy import create_engine, Column, Integer, Unicode, UniqueConstraint, ForeignKey
from sqlalchemy.orm import relationship, synonym, sessionmaker, scoped_session
from sqlalchemy.ext.declarative import declarative_base

engine = create_engine('sqlite://')
Session = scoped_session(sessionmaker(autocommit=False,
              autoflush=False,
              bind=engine))
Base = declarative_base(bind=engine)

# --- Models ------------------------------------------------------------------

class Container(BaseMixin, Base):
    __tablename__ = 'container'
    query = Session.query_property()

    content = Column(Unicode(250))

class UnnamedDocument(BaseMixin, Base):
    __tablename__ = 'unnamed_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)

    content = Column(Unicode(250))

class NamedDocument(BaseNameMixin, Base):
    __tablename__ = 'named_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)

    content = Column(Unicode(250))

class ScopedNamedDocument(BaseScopedNameMixin, Base):
    __tablename__ = 'scoped_named_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)
    parent = synonym('container')

    content = Column(Unicode(250))
    __table_args__ = (UniqueConstraint('name', 'container_id'),)

class IdNamedDocument(BaseIdNameMixin, Base):
    __tablename__ = 'id_named_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)

    content = Column(Unicode(250))

class ScopedIdNamedDocument(BaseScopedIdNameMixin, Base):
    __tablename__ = 'scoped_id_named_document'
    query = Session.query_property()
    container_id = Column(Integer, ForeignKey('container.id'))
    container = relationship(Container)
    parent = synonym('container')

    content = Column(Unicode(250))
    __table_args__ = (UniqueConstraint('url_id', 'container_id'),)

# -- Tests --------------------------------------------------------------------

class TestCoasterModels(unittest.TestCase):
    def setUp(self):
        Base.metadata.create_all()
        self.session = Session()

    def tearDown(self):
        self.session.rollback()
        Base.metadata.drop_all()

    def make_container(self):
        c = Container()
        self.session.add(c)
        return c

    def test_container(self):
        c = self.make_container()
        self.assertEqual(c.id, None)
        self.session.commit()
        self.assertEqual(c.id, 1)

    def test_timestamp(self):
        now1 = datetime.utcnow()
        c = self.make_container()
        self.session.commit()
        u = c.updated_at
        now2 = datetime.utcnow()
        self.assertTrue(now1 < c.created_at)
        self.assertTrue(now2 > c.created_at)
        c.content = "updated"
        self.session.commit()
        self.assertTrue(c.updated_at > now2)
        self.assertTrue(c.updated_at > c.created_at)
        self.assertTrue(c.updated_at > u)

    def test_unnamed(self):
        c = self.make_container()
        d = UnnamedDocument(content="hello", container=c)
        self.session.add(d)
        self.session.commit()
        self.assertEqual(c.id, 1)
        self.assertEqual(d.id, 1)

    def test_named(self):
        """Named documents have globally unique names."""
        c1 = self.make_container()
        self.session.add(c1)
        d1 = NamedDocument(title="Hello", content="World", container=c1)
        self.session.add(d1)
        self.session.commit()
        self.assertEqual(d1.name, 'hello')

        c2 = self.make_container()
        d2 = NamedDocument(title="Hello", content="Again", container=c2)
        self.session.add(d2)
        self.session.commit()
        self.assertEqual(d2.name, 'hello1')

    def test_scoped_named(self):
        """Scoped named documents have names unique to their containers."""
        c1 = self.make_container()
        self.session.add(c1)
        d1 = ScopedNamedDocument(title="Hello", content="World", container=c1)
        self.session.add(d1)
        self.session.commit()
        self.assertEqual(d1.name, 'hello')

        d2 = ScopedNamedDocument(title="Hello", content="Again", container=c1)
        self.session.add(d2)
        self.session.commit()
        self.assertEqual(d2.name, 'hello1')

        c2 = self.make_container()
        d3 = ScopedNamedDocument(title="Hello", content="Once More", container=c2)
        self.session.add(d3)
        self.session.commit()
        self.assertEqual(d3.name, 'hello')

    def test_id_named(self):
        """Documents with a global id in the URL"""
        c1 = self.make_container()
        self.session.add(c1)
        d1 = IdNamedDocument(title="Hello", content="World", container=c1)
        self.session.add(d1)
        self.session.commit()
        self.assertEqual(d1.url_name, '1-hello')

        d2 = IdNamedDocument(title="Hello", content="Again", container=c1)
        self.session.add(d2)
        self.session.commit()
        self.assertEqual(d2.url_name, '2-hello')

        c2 = self.make_container()
        d3 = IdNamedDocument(title="Hello", content="Once More", container=c2)
        self.session.add(d3)
        self.session.commit()
        self.assertEqual(d3.url_name, '3-hello')

    def test_scoped_id_named(self):
        """Documents with a container-specifc id in the URL"""
        c1 = self.make_container()
        self.session.add(c1)
        d1 = ScopedIdNamedDocument(title="Hello", content="World", container=c1)
        self.session.add(d1)
        self.session.commit()
        self.assertEqual(d1.url_name, '1-hello')

        d2 = ScopedIdNamedDocument(title="Hello again", content="New name", container=c1)
        self.session.add(d2)
        self.session.commit()
        self.assertEqual(d2.url_name, '2-hello-again')

        c2 = self.make_container()
        d3 = ScopedIdNamedDocument(title="Hello", content="Once More", container=c2)
        self.session.add(d3)
        self.session.commit()
        self.assertEqual(d3.url_name, '1-hello')

        d4 = ScopedIdNamedDocument(title="Hello", content="Third", container=c1)
        self.session.add(d4)
        self.session.commit()
        self.assertEqual(d4.url_name, '3-hello')

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


<commit_before>
import shutil
from tempfile import mkdtemp

import pytest
import peewee
from vcstools import get_vcs_client

from kez.models import Project, Document, Repository
from kez.utils import evaluate_config_options
from .data import *

pathexists = os.path.exists
pathjoin = os.path.join

def test_vcs_checkout():
    tmp= mkdtemp(prefix="kez-test-")
    git = pathjoin(tmp, '.git')
    assert not pathexists(git)
    client = get_vcs_client("git", tmp)
    client.checkout(URL1)
    assert pathexists(git)
    shutil.rmtree(tmp)

def test_create_project_from_url(db, vcs_cache):
    query = list(Project.select())
    assert len(query) == 0
    project, repo = Project.from_url(URL1, vcs_cache, name="blog")
    assert project
    assert project.url == URL1
    assert project.vcs == "git"
    assert project.name == "blog"
    assert project.host == "github"
    project.save()
    query = list(Project.select())
    assert len(query) == 1
    assert query[0].url == URL1
    assert query[0].vcs == "git"
    assert query[0].name == "blog"
    assert query[0].host == "github"

def test_empty_and_non_empty_query():
    query = Project.select().where(Project.url=="doesnotexist")
    assert len(list(query)) == 0
    query = Project.select().where(Project.url==URL1)
    assert len(list(query)) == 1

def test_read_project_repository_config(vcs_cache):
    repo = Repository.instance("blog", vcs_cache)
    repo.checkout()
    cfg = repo.get_project_config()
    assert cfg.sections() == ['maths.averagehuman.org']
    assert cfg.get('maths.averagehuman.org', '__doctype__') == 'pelican'
    options, settings = evaluate_config_options(cfg, 'maths.averagehuman.org')
    assert options['doctype'] == 'pelican'
    assert settings['THEME'] == 'theme'

def test_process_project_repository(vcs_cache):
    repo = Repository.instance("blog", vcs_cache)
    repo.checkout()
    assert(len(list(Document.select()))) == 0
    repo.process()
    docs = list(Document.select())
    assert len(docs) == 1
    doc = docs[0]
    assert doc.name == "maths.averagehuman.org"
    assert doc.doctype == "pelican"

def test_delete_project(db):
    query = Project.select().where(Project.url==URL1)
    assert len(list(query)) == 1
    project = query[0]
    assert project.url == URL1
    dq = project.delete()
    assert isinstance(dq, peewee.DeleteQuery)
    dq.execute()
    query = Project.select().where(Project.url==URL1)
    assert len(list(query)) == 0

<commit_msg>update model test<commit_after>
import shutil
from tempfile import mkdtemp

import pytest
import peewee
from vcstools import get_vcs_client

from kez.models import Project, Document, Repository
from kez.utils import evaluate_config_options
from .data import *

pathexists = os.path.exists
pathjoin = os.path.join

def test_vcs_checkout():
    tmp= mkdtemp(prefix="kez-test-")
    git = pathjoin(tmp, '.git')
    assert not pathexists(git)
    client = get_vcs_client("git", tmp)
    client.checkout(URL1)
    assert pathexists(git)
    shutil.rmtree(tmp)

def test_create_project_from_url(db, vcs_cache):
    query = list(Project.select())
    assert len(query) == 0
    project, repo = Project.from_url(URL1, vcs_cache, name="blog")
    assert project
    assert project.url == URL1
    assert project.vcs == "git"
    assert project.name == "blog"
    assert project.host == "github"
    project.save()
    query = list(Project.select())
    assert len(query) == 1
    assert query[0].url == URL1
    assert query[0].vcs == "git"
    assert query[0].name == "blog"
    assert query[0].host == "github"

def test_empty_and_non_empty_query():
    query = Project.select().where(Project.url=="doesnotexist")
    assert len(list(query)) == 0
    query = Project.select().where(Project.url==URL1)
    assert len(list(query)) == 1

def test_read_project_repository_config(vcs_cache):
    repo = Repository.instance("blog", vcs_cache)
    repo.checkout()
    cfg = repo.get_project_config()
    assert cfg.sections() == ['maths.averagehuman.org']
    assert cfg.get('maths.averagehuman.org', '__doctype__') == 'pelican'
    options, settings = evaluate_config_options(cfg, 'maths.averagehuman.org')
    assert options['doctype'] == 'pelican'
    assert settings['THEME']

def test_process_project_repository(vcs_cache):
    repo = Repository.instance("blog", vcs_cache)
    repo.checkout()
    assert(len(list(Document.select()))) == 0
    repo.process()
    docs = list(Document.select())
    assert len(docs) == 1
    doc = docs[0]
    assert doc.name == "maths.averagehuman.org"
    assert doc.doctype == "pelican"

def test_delete_project(db):
    query = Project.select().where(Project.url==URL1)
    assert len(list(query)) == 1
    project = query[0]
    assert project.url == URL1
    dq = project.delete()
    assert isinstance(dq, peewee.DeleteQuery)
    dq.execute()
    query = Project.select().where(Project.url==URL1)
    assert len(list(query)) == 0

<|endoftext|>


<commit_before># $Id$
# vi:ts=4:et

# same as test_multi2.py, but enforce some debugging and strange API-calls

import os, sys, time
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO
import pycurl

urls = (
    "http://curl.haxx.se",
    "http://www.python.org",
    "http://pycurl.sourceforge.net",
    "http://pycurl.sourceforge.net/THIS_HANDLE_IS_CLOSED",
)

# init
m = pycurl.multi_init()
m.handles = []
for url in urls:
    c = pycurl.init()
    # save info in standard Python attributes
    c.url = url
    c.body = StringIO()
    c.debug = 0
    m.handles.append(c)
    # pycurl API calls
    c.setopt(pycurl.URL, c.url)
    c.setopt(pycurl.WRITEFUNCTION, c.body.write)
    m.add_handle(c)

# debug - close a handle
if 1:
    c = m.handles[3]
    c.debug = 1
    c.cleanup()

# get data
while 1:
    num_handles = m.perform()
    if num_handles == 0:
        break
    # currently no more I/O is pending, could do something in the meantime
    # (display a progress bar, etc.)
    time.sleep(0.01)

# close handles
for c in m.handles:
    # save info in standard Python attributes
    try:
        c.http_code = c.getinfo(pycurl.HTTP_CODE)
    except pycurl.error:
        # handle already closed - see debug above
        assert c.debug
        c.http_code = -1
    # pycurl API calls
    if 0:
        m.remove_handle(c)
        c.cleanup()
    else:
        # in the C API this is the wrong calling order, but pycurl
        # handles this automatically
        c.cleanup()
        m.remove_handle(c)
m.cleanup()

# print result
for c in m.handles:
    data = c.body.getvalue()
    if 0:
        print "**********", c.url, "**********"
        print data
    else:
        print "%-53s http_code %3d, %6d bytes" % (c.url, c.http_code, len(data))

<commit_msg>Updated.<commit_after># $Id$
# vi:ts=4:et

# same as test_multi2.py, but enforce some debugging and strange API-calls

import os, sys, time
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO
import pycurl

urls = (
    "http://curl.haxx.se",
    "http://www.python.org",
    "http://pycurl.sourceforge.net",
    "http://pycurl.sourceforge.net/THIS_HANDLE_IS_CLOSED",
)

# init
m = pycurl.multi_init()
m.handles = []
for url in urls:
    c = pycurl.init()
    # save info in standard Python attributes
    c.url = url
    c.body = StringIO()
    c.debug = 0
    m.handles.append(c)
    # pycurl API calls
    c.setopt(pycurl.URL, c.url)
    c.setopt(pycurl.WRITEFUNCTION, c.body.write)
    m.add_handle(c)

# debug - close a handle
if 1:
    c = m.handles[3]
    c.debug = 1
    c.cleanup()

# get data
while 1:
    num_handles = m.perform()
    if num_handles == 0:
        break
    # currently no more I/O is pending, could do something in the meantime
    # (display a progress bar, etc.)
    time.sleep(0.01)

# close handles
for c in m.handles:
    # save info in standard Python attributes
    try:
        c.http_code = c.getinfo(pycurl.HTTP_CODE)
    except pycurl.error:
        # handle already closed - see debug above
        assert c.debug
        c.http_code = -1
    # pycurl API calls
    if 0:
        m.remove_handle(c)
        c.cleanup()
    elif 0:
        # in the C API this is the wrong calling order, but pycurl
        # handles this automatically
        c.cleanup()
        m.remove_handle(c)
    else:
        # actually, remove_handle is called automatically on cleanup
        c.cleanup()
m.cleanup()

# print result
for c in m.handles:
    data = c.body.getvalue()
    if 0:
        print "**********", c.url, "**********"
        print data
    else:
        print "%-53s http_code %3d, %6d bytes" % (c.url, c.http_code, len(data))

<|endoftext|>


<commit_before>import os
import sys
import traceback
import StringIO
import unittest
import warnings

warnings.simplefilter("error")

from support import html5lib_test_files, TestData, convert, convertExpected
import html5lib
from html5lib import html5parser, treebuilders, constants

treeTypes = {"simpletree":treebuilders.getTreeBuilder("simpletree"),
             "DOM":treebuilders.getTreeBuilder("dom")}

#Try whatever etree implementations are avaliable from a list that are
#"supposed" to work
try:
    import xml.etree.ElementTree as ElementTree
    treeTypes['ElementTree'] = treebuilders.getTreeBuilder("etree", ElementTree, fullTree=True)
except ImportError:
    try:
        import elementtree.ElementTree as ElementTree
        treeTypes['ElementTree'] = treebuilders.getTreeBuilder("etree", ElementTree, fullTree=True)
    except ImportError:
        pass

try:
    import xml.etree.cElementTree as cElementTree
    treeTypes['cElementTree'] = treebuilders.getTreeBuilder("etree", cElementTree, fullTree=True)
except ImportError:
    try:
        import cElementTree
        treeTypes['cElementTree'] = treebuilders.getTreeBuilder("etree", cElementTree, fullTree=True)
    except ImportError:
        pass

try:
    try:
        import lxml.html as lxml
    except ImportError:
        import lxml.etree as lxml
    treeTypes['lxml'] = treebuilders.getTreeBuilder("lxml", lxml, fullTree=True)
except ImportError:
    pass

try:
    import BeautifulSoup
    treeTypes["beautifulsoup"] = treebuilders.getTreeBuilder("beautifulsoup", fullTree=True)
except ImportError:
    pass

#Try whatever dom implementations are avaliable from a list that are
#"supposed" to work
try:
    import pxdom
    treeTypes["pxdom"] = treebuilders.getTreeBuilder("dom", pxdom)
except ImportError:
    pass

#Run the parse error checks
checkParseErrors = False

#XXX - There should just be one function here but for some reason the testcase
#format differs from the treedump format by a single space character
def convertTreeDump(data):
    return "\n".join(convert(3)(data).split("\n")[1:])

import re
attrlist = re.compile(r"^(\s+)\w+=.*(\n\1\w+=.*)+",re.M)
def sortattrs(x):
  lines = x.group(0).split("\n")
  lines.sort()
  return "\n".join(lines)

class TestCase(unittest.TestCase):
    def runParserTest(self, innerHTML, input, expected, errors, treeClass):
        #XXX - move this out into the setup function
        #concatenate all consecutive character tokens into a single token
        p = html5parser.HTMLParser(tree = treeClass)

        try:
            if innerHTML:
                document = p.parseFragment(StringIO.StringIO(input), innerHTML)
            else:
                try:
                    document = p.parse(StringIO.StringIO(input))
                except constants.DataLossWarning:
                    sys.stderr.write("Test input causes known dataloss, skipping")
                    return
        except:
            errorMsg = "\n".join(["\n\nInput:", input, "\nExpected:", expected,
                                  "\nTraceback:", traceback.format_exc()])
            self.assertTrue(False, errorMsg)

        output = convertTreeDump(p.tree.testSerializer(document))
        output = attrlist.sub(sortattrs, output)

        expected = convertExpected(expected)
        expected = attrlist.sub(sortattrs, expected)
        errorMsg = "\n".join(["\n\nInput:", input, "\nExpected:", expected,
                              "\nReceived:", output])
        self.assertEquals(expected, output, errorMsg)
        errStr = ["Line: %i Col: %i %s %s"%(line, col,
                                         constants.E[errorcode], datavars) for
                  ((line,col), errorcode, datavars) in p.errors]
        errorMsg2 = "\n".join(["\n\nInput:", input,
                               "\nExpected errors (" + str(len(errors)) + "):\n" + "\n".join(errors),
                               "\nActual errors (" + str(len(p.errors)) + "):\n" + "\n".join(errStr)])
        if checkParseErrors:
            self.assertEquals(len(p.errors), len(errors), errorMsg2)

def buildTestSuite():
    sys.stdout.write('Testing tree builders '+ " ".join(treeTypes.keys()) + "\n")

    for treeName, treeCls in treeTypes.iteritems():
        files = html5lib_test_files('tree-construction')
        files = [f for f in files if
                 not f.split(".")[-2][-2:] in ("s9", "10", "11", "12")] #skip namespace tests for now
        for filename in files:
            testName = os.path.basename(filename).replace(".dat","")

            tests = TestData(filename, "data")

            for index, test in enumerate(tests):
                input, errors, innerHTML, expected = [test[key] for key in
                                                      'data', 'errors',
                                                      'document-fragment',
                                                      'document']
                if errors:
                    errors = errors.split("\n")
                def testFunc(self, innerHTML=innerHTML, input=input,
                    expected=expected, errors=errors, treeCls=treeCls):
                    return self.runParserTest(innerHTML, input, expected, errors, treeCls)
                setattr(TestCase, "test_%s_%d_%s" % (testName,index+1,treeName),
                     testFunc)

    return unittest.TestLoader().loadTestsFromTestCase(TestCase)

def main():
    # the following is temporary while the unit tests for parse errors are
    # still in flux
    if '-p' in sys.argv: # suppress check for parse errors
        sys.argv.remove('-p')
        global checkParseErrors
        checkParseErrors = False
    buildTestSuite()
    try:
        unittest.main()
    except SystemExit:
	    pass

    # f = open("graph.dot", "w")
    # f.write(str(g))

    # print g.nodes.keys()

if __name__ == "__main__":
    print sys.argv
    main()
<commit_msg>Fix test_parser.py to have working single-test runs, also add *.pyc to ignore list.<commit_after>import os
import sys
import traceback
import StringIO
import unittest
import warnings

warnings.simplefilter("error")

from support import html5lib_test_files, TestData, convert, convertExpected
import html5lib
from html5lib import html5parser, treebuilders, constants

treeTypes = {"simpletree":treebuilders.getTreeBuilder("simpletree"),
             "DOM":treebuilders.getTreeBuilder("dom")}

#Try whatever etree implementations are avaliable from a list that are
#"supposed" to work
try:
    import xml.etree.ElementTree as ElementTree
    treeTypes['ElementTree'] = treebuilders.getTreeBuilder("etree", ElementTree, fullTree=True)
except ImportError:
    try:
        import elementtree.ElementTree as ElementTree
        treeTypes['ElementTree'] = treebuilders.getTreeBuilder("etree", ElementTree, fullTree=True)
    except ImportError:
        pass

try:
    import xml.etree.cElementTree as cElementTree
    treeTypes['cElementTree'] = treebuilders.getTreeBuilder("etree", cElementTree, fullTree=True)
except ImportError:
    try:
        import cElementTree
        treeTypes['cElementTree'] = treebuilders.getTreeBuilder("etree", cElementTree, fullTree=True)
    except ImportError:
        pass

try:
    try:
        import lxml.html as lxml
    except ImportError:
        import lxml.etree as lxml
    treeTypes['lxml'] = treebuilders.getTreeBuilder("lxml", lxml, fullTree=True)
except ImportError:
    pass

try:
    import BeautifulSoup
    treeTypes["beautifulsoup"] = treebuilders.getTreeBuilder("beautifulsoup", fullTree=True)
except ImportError:
    pass

#Try whatever dom implementations are avaliable from a list that are
#"supposed" to work
try:
    import pxdom
    treeTypes["pxdom"] = treebuilders.getTreeBuilder("dom", pxdom)
except ImportError:
    pass

#Run the parse error checks
checkParseErrors = False

#XXX - There should just be one function here but for some reason the testcase
#format differs from the treedump format by a single space character
def convertTreeDump(data):
    return "\n".join(convert(3)(data).split("\n")[1:])

import re
attrlist = re.compile(r"^(\s+)\w+=.*(\n\1\w+=.*)+",re.M)
def sortattrs(x):
  lines = x.group(0).split("\n")
  lines.sort()
  return "\n".join(lines)

class TestCase(unittest.TestCase):
    def runParserTest(self, innerHTML, input, expected, errors, treeClass):
        #XXX - move this out into the setup function
        #concatenate all consecutive character tokens into a single token
        p = html5parser.HTMLParser(tree = treeClass)

        try:
            if innerHTML:
                document = p.parseFragment(StringIO.StringIO(input), innerHTML)
            else:
                try:
                    document = p.parse(StringIO.StringIO(input))
                except constants.DataLossWarning:
                    sys.stderr.write("Test input causes known dataloss, skipping")
                    return
        except:
            errorMsg = "\n".join(["\n\nInput:", input, "\nExpected:", expected,
                                  "\nTraceback:", traceback.format_exc()])
            self.assertTrue(False, errorMsg)

        output = convertTreeDump(p.tree.testSerializer(document))
        output = attrlist.sub(sortattrs, output)

        expected = convertExpected(expected)
        expected = attrlist.sub(sortattrs, expected)
        errorMsg = "\n".join(["\n\nInput:", input, "\nExpected:", expected,
                              "\nReceived:", output])
        self.assertEquals(expected, output, errorMsg)
        errStr = ["Line: %i Col: %i %s %s"%(line, col,
                                         constants.E[errorcode], datavars) for
                  ((line,col), errorcode, datavars) in p.errors]
        errorMsg2 = "\n".join(["\n\nInput:", input,
                               "\nExpected errors (" + str(len(errors)) + "):\n" + "\n".join(errors),
                               "\nActual errors (" + str(len(p.errors)) + "):\n" + "\n".join(errStr)])
        if checkParseErrors:
            self.assertEquals(len(p.errors), len(errors), errorMsg2)

def buildTestSuite():
    sys.stdout.write('Testing tree builders '+ " ".join(treeTypes.keys()) + "\n")

    for treeName, treeCls in treeTypes.iteritems():
        files = html5lib_test_files('tree-construction')
        files = [f for f in files if
                 not f.split(".")[-2][-2:] in ("s9", "10", "11", "12")] #skip namespace tests for now
        for filename in files:
            testName = os.path.basename(filename).replace(".dat","")

            tests = TestData(filename, "data")

            for index, test in enumerate(tests):
                input, errors, innerHTML, expected = [test[key] for key in
                                                      'data', 'errors',
                                                      'document-fragment',
                                                      'document']
                if errors:
                    errors = errors.split("\n")
                def testFunc(self, innerHTML=innerHTML, input=input,
                    expected=expected, errors=errors, treeCls=treeCls):
                    return self.runParserTest(innerHTML, input, expected, errors, treeCls)
                testFunc.__name__ = "test_%s_%d_%s" % (testName,index+1,treeName)
                setattr(TestCase, testFunc.__name__,
                     testFunc)

    return unittest.TestLoader().loadTestsFromTestCase(TestCase)

def main():
    # the following is temporary while the unit tests for parse errors are
    # still in flux
    if '-p' in sys.argv: # suppress check for parse errors
        sys.argv.remove('-p')
        global checkParseErrors
        checkParseErrors = False
    buildTestSuite()
    try:
        unittest.main()
    except SystemExit:
	    pass

    # f = open("graph.dot", "w")
    # f.write(str(g))

    # print g.nodes.keys()

if __name__ == "__main__":
    print sys.argv
    main()
<|endoftext|>


<commit_before>import os
import ast
import pytest
import logging
from structifytext import parser
from cStringIO import StringIO

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

@pytest.fixture(scope='module')
def mock_struct(request):
    return {
        'tables': [
            {
                'id': '\[TABLE (\d{1,2})\]',
                'flows': [
                    {
                        'id': '\[FLOW_ID(\d+)\]',
                        'timestamp': 'Timestamp\s+=\s+(.+)'
                    }
                ]
            }
        ]
    }

@pytest.fixture(scope='module')
def mock_group_struct(request):
    return {
        'groups': [
            {
                'id': 'Group id:\s+(\d+)',
                'ref_count': 'Reference count:\s+(\d+)',
                'packet_count': 'Packet count:\s+(\d+)',
                'byte_count': 'Byte count:\s+(\d+)',
                'bucket': [
                    {
                        'id': 'Bucket\s+(\d+)',
                        'packet_count': 'Packet count:\s+(\d+)',
                        'byte_count': 'Byte count:\s+(\d+)',
                    }
                ]
            }
        ]
    }

def read(filename):
    output_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), filename)
    with open(output_file, 'r') as fin:
        return fin.read()

def test_simple_struct():
    struct = { 'message': '(.*)' }
    lines = ["Hello World"]
    expected_output = { 'message': "Hello World" }
    parsed = parser.parse_struct(lines, struct)
    assert parsed == expected_output

def test_simple_list():
    struct = { 'count': ['(\d)'] }
    lines = ["The count says: 1", "The count says: 2", "The count says: 3",
             "The count says: 4", "The count says: 5"]
    expected_output = { 'count': [1,2,3,4,5] }
    parsed = parser.parse_struct(lines, struct)
    # For now we need to convert to type after parsing
    parsed = {'count': map(int, parsed['count'])}
    assert parsed == expected_output

def test_flows(mock_struct):
    lines = StringIO(read('./flow_output.txt')).readlines()
    expected_output = ast.literal_eval(read('./flow_output_parsed.txt'))
    parsed = parser.parse_struct(lines, mock_struct)
    assert parsed == expected_output

def test_groups(mock_group_struct):
    lines = StringIO(read('./group_output.txt')).readlines()
    expected_output = ast.literal_eval(read('./group_output_parsed.txt'))
    parsed = parser.parse_struct(lines, mock_group_struct)
    assert parsed == expected_output

def test_value_not_regex_string_raises_exception():
    struct = {'message': 123}
    lines = ["Hello World"]
    with pytest.raises(TypeError):
        parser.parse_struct(lines, struct)

def test_value_without_group_raises_exception():
    struct = {'message': 'ab'}
    lines = ["Hello World"]
    with pytest.raises(ValueError):
        parser.parse_struct(lines, struct)

def test_value_with_two_groups_raises_warning():
    struct = {'message': '(.*)\S+(.*)'}
    lines = ["Hello World"]
    with pytest.raises(UserWarning):
        parser.parse_struct(lines, struct)

def test_list_with_dict_no_id_raises_exception():
    struct = {'letter': [{'to': 'Dear\s+(\w+)', 'from': 'Regards,\s+(\w+)'}]}
    letter = "Dear Einstein,\r\n"
    letter += "I am a death destroyer of worlds.\r\n"
    letter += "And it's all your fault!\r\n"
    letter += "Regards, Openhiemer\r\n"
    lines = StringIO(letter).readlines()
    logger.debug(lines)
    with pytest.raises(KeyError):
        parser.parse_struct(lines, struct)
<commit_msg>Fixing up quote in test...<commit_after>import os
import ast
import pytest
import logging
from structifytext import parser
from cStringIO import StringIO

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

@pytest.fixture(scope='module')
def mock_struct(request):
    return {
        'tables': [
            {
                'id': '\[TABLE (\d{1,2})\]',
                'flows': [
                    {
                        'id': '\[FLOW_ID(\d+)\]',
                        'timestamp': 'Timestamp\s+=\s+(.+)'
                    }
                ]
            }
        ]
    }

@pytest.fixture(scope='module')
def mock_group_struct(request):
    return {
        'groups': [
            {
                'id': 'Group id:\s+(\d+)',
                'ref_count': 'Reference count:\s+(\d+)',
                'packet_count': 'Packet count:\s+(\d+)',
                'byte_count': 'Byte count:\s+(\d+)',
                'bucket': [
                    {
                        'id': 'Bucket\s+(\d+)',
                        'packet_count': 'Packet count:\s+(\d+)',
                        'byte_count': 'Byte count:\s+(\d+)',
                    }
                ]
            }
        ]
    }

def read(filename):
    output_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), filename)
    with open(output_file, 'r') as fin:
        return fin.read()

def test_simple_struct():
    struct = { 'message': '(.*)' }
    lines = ["Hello World"]
    expected_output = { 'message': "Hello World" }
    parsed = parser.parse_struct(lines, struct)
    assert parsed == expected_output

def test_simple_list():
    struct = { 'count': ['(\d)'] }
    lines = ["The count says: 1", "The count says: 2", "The count says: 3",
             "The count says: 4", "The count says: 5"]
    expected_output = { 'count': [1,2,3,4,5] }
    parsed = parser.parse_struct(lines, struct)
    # For now we need to convert to type after parsing
    parsed = {'count': map(int, parsed['count'])}
    assert parsed == expected_output

def test_flows(mock_struct):
    lines = StringIO(read('./flow_output.txt')).readlines()
    expected_output = ast.literal_eval(read('./flow_output_parsed.txt'))
    parsed = parser.parse_struct(lines, mock_struct)
    assert parsed == expected_output

def test_groups(mock_group_struct):
    lines = StringIO(read('./group_output.txt')).readlines()
    expected_output = ast.literal_eval(read('./group_output_parsed.txt'))
    parsed = parser.parse_struct(lines, mock_group_struct)
    assert parsed == expected_output

def test_value_not_regex_string_raises_exception():
    struct = {'message': 123}
    lines = ["Hello World"]
    with pytest.raises(TypeError):
        parser.parse_struct(lines, struct)

def test_value_without_group_raises_exception():
    struct = {'message': 'ab'}
    lines = ["Hello World"]
    with pytest.raises(ValueError):
        parser.parse_struct(lines, struct)

def test_value_with_two_groups_raises_warning():
    struct = {'message': '(.*)\S+(.*)'}
    lines = ["Hello World"]
    with pytest.raises(UserWarning):
        parser.parse_struct(lines, struct)

def test_list_with_dict_no_id_raises_exception():
    struct = {'letter': [{'to': 'Dear\s+(\w+)', 'from': 'Regards,\s+(\w+)'}]}
    letter = "Dear Einstein,\r\n"
    letter += "I am become Death, the destroyer of worlds.\r\n"
    letter += "And it's all your fault!\r\n"
    letter += "Regards, Oppenheimer\r\n"
    lines = StringIO(letter).readlines()
    logger.debug(lines)
    with pytest.raises(KeyError):
        parser.parse_struct(lines, struct)
<|endoftext|>


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

"""
    cobra
    ~~~~~

    Implements cobra main

    :author:    BlBana <635373043@qq.com>
    :homepage:  https://github.com/wufeifei/cobra
    :license:   MIT, see LICENSE for more details.
    :copyright: Copyright (c) 2017 Feei. All rights reserved
"""
from cobra.config import project_directory
from cobra.core_engine.php.parser import anlysis_params
from cobra.core_engine.php.parser import scan_parser
from cobra.pretreatment import ast_object

files = [('.php', {'list': ["v_parser.php", "v.php"]})]
ast_object.init_pre(project_directory + '/tests/vulnerabilities/', files)
ast_object.pre_ast_all('php')

target_projects = project_directory + '/tests/vulnerabilities/v_parser.php'
target_projects2 = project_directory + '/tests/vulnerabilities/v.php'

with open(target_projects, 'r') as fi:
    code_contents = fi.read()
with open(target_projects, 'r') as fi2:
    code_contents2 = fi2.read()

sensitive_func = ['system']
lineno = 7

param = '$callback'
lineno2 = 10

def test_scan_parser():
    assert scan_parser(sensitive_func, lineno, target_projects)

def test_anlysis_params():
    assert anlysis_params(param, target_projects2, lineno2)
<commit_msg>fix a bug for py.test<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
    cobra
    ~~~~~

    Implements cobra main

    :author:    BlBana <635373043@qq.com>
    :homepage:  https://github.com/wufeifei/cobra
    :license:   MIT, see LICENSE for more details.
    :copyright: Copyright (c) 2017 Feei. All rights reserved
"""
from cobra.config import project_directory
from cobra.core_engine.php.parser import anlysis_params
from cobra.core_engine.php.parser import scan_parser
from cobra.pretreatment import ast_object

files = [('.php', {'list': ["v_parser.php", "v.php"]})]
ast_object.init_pre(project_directory + '/tests/vulnerabilities/', files)
ast_object.pre_ast_all(['php'])

target_projects = project_directory + '/tests/vulnerabilities/v_parser.php'
target_projects2 = project_directory + '/tests/vulnerabilities/v.php'

with open(target_projects, 'r') as fi:
    code_contents = fi.read()
with open(target_projects, 'r') as fi2:
    code_contents2 = fi2.read()

sensitive_func = ['system']
lineno = 7

param = '$callback'
lineno2 = 10

def test_scan_parser():
    assert scan_parser(sensitive_func, lineno, target_projects)

def test_anlysis_params():
    assert anlysis_params(param, target_projects2, lineno2)
<|endoftext|>


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

"""
test_pulsar
----------------------------------

Tests for `pulsar` module. Will eventually want to add tests
for time slicing, PINT integration and pickling.
"""

import sys
import os
import shutil
import unittest
import pickle
import pytest

import numpy as np

from enterprise.pulsar import Pulsar
from tests.enterprise_test_data import datadir
from pint.models import get_model_and_toas

class TestPulsar(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        """Setup the Pulsar object."""

        # initialize Pulsar class
        cls.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim")

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree("pickle_dir", ignore_errors=True)

    def test_residuals(self):
        """Check Residual shape."""

        msg = "Residuals shape incorrect"
        assert self.psr.residuals.shape == (4005,), msg

    def test_toaerrs(self):
        """Check TOA errors shape."""

        msg = "TOA errors shape incorrect"
        assert self.psr.toaerrs.shape == (4005,), msg

    def test_toas(self):
        """Check TOA shape."""

        msg = "TOA shape incorrect"
        assert self.psr.toas.shape == (4005,), msg

    def test_stoas(self):
        """Check STOA shape."""

        msg = "stoa shape incorrect"
        assert self.psr.stoas.shape == (4005,), msg

    def test_dm(self):
        """Check DM/DMX access."""

        msg = "dm value incorrect"
        assert self.psr.dm == 13.299393, msg

        msg = "dmx struct incorrect (spotcheck)"
        assert len(self.psr.dmx) == 72, msg
        assert self.psr.dmx["DMX_0001"]["DMX"] == 0.015161863, msg
        assert self.psr.dmx["DMX_0001"]["fit"], msg

    def test_freqs(self):
        """Check frequencies shape."""

        msg = "Frequencies shape incorrect"
        assert self.psr.freqs.shape == (4005,), msg

    def test_flags(self):
        """Check flags shape and content"""

        msg = "Flags shape incorrect"
        assert self.psr.flags["f"].shape == (4005,), msg

        msg = "Flag content or sorting incorrect"
        assert np.all(self.psr._flags["fe"][self.psr._isort] == self.psr.flags["fe"]), msg

    def test_backend_flags(self):
        """Check backend_flags shape and content"""

        msg = "Backend flags shape incorrect"
        assert self.psr.backend_flags.shape == (4005,), msg

        # for the test pulsar, backend should be the same as 'f'
        msg = "Flag content or sorting incorrect"
        assert np.all(self.psr._flags["f"][self.psr._isort] == self.psr.backend_flags), msg

    def test_sky(self):
        """Check Sky location."""

        sky = (1.4023093811712661, 4.9533700839400492)

        msg = "Incorrect sky location"
        assert np.allclose(self.psr.theta, sky[0]), msg
        assert np.allclose(self.psr.phi, sky[1]), msg

    def test_design_matrix(self):
        """Check design matrix shape."""

        msg = "Design matrix shape incorrect."
        assert self.psr.Mmat.shape == (4005, 91), msg

    def test_filter_data(self):
        """Place holder for filter_data tests."""
        assert self.psr.filter_data() is None

    def test_planetssb(self):
        """Place holder for filter_data tests."""
        assert hasattr(self.psr, "planetssb")

    def test_sunssb(self):
        """Place holder for filter_data tests."""
        assert hasattr(self.psr, "sunssb")

    def test_to_pickle(self):
        """Place holder for to_pickle tests."""
        self.psr.to_pickle()
        with open("B1855+09.pkl", "rb") as f:
            pkl_psr = pickle.load(f)

        os.remove("B1855+09.pkl")

        assert np.allclose(self.psr.residuals, pkl_psr.residuals, rtol=1e-10)

        self.psr.to_pickle("pickle_dir")
        with open("pickle_dir/B1855+09.pkl", "rb") as f:
            pkl_psr = pickle.load(f)

        assert np.allclose(self.psr.residuals, pkl_psr.residuals, rtol=1e-10)

    @pytest.mark.skipif(sys.version_info < (3, 8), reason="Requires Python >= 3.8")
    def test_deflate_inflate(self):
        psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim")

        dm = psr._designmatrix.copy()

        psr.deflate()
        psr.to_pickle()

        with open("B1855+09.pkl", "rb") as f:
            pkl_psr = pickle.load(f)
        pkl_psr.inflate()

        assert np.allclose(dm, pkl_psr._designmatrix)

        del pkl_psr

        psr.destroy()

        with open("B1855+09.pkl", "rb") as f:
            pkl_psr = pickle.load(f)

        with self.assertRaises(FileNotFoundError):
            pkl_psr.inflate()

    def test_wrong_input(self):
        """Test exception when incorrect par(tim) file given."""

        with self.assertRaises(IOError) as context:
            Pulsar("wrong.par", "wrong.tim")

            msg = "Cannot find parfile wrong.par or timfile wrong.tim!"
            self.assertTrue(msg in context.exception)

    def test_value_error(self):
        """Test exception when unknown argument is given"""

        with self.assertRaises(ValueError):
            Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.time")

class TestPulsarPint(TestPulsar):
    @classmethod
    def setUpClass(cls):
        """Setup the Pulsar object."""

        # initialize Pulsar class
        cls.psr = Pulsar(
            datadir + "/B1855+09_NANOGrav_9yv1.gls.par",
            datadir + "/B1855+09_NANOGrav_9yv1.tim",
            ephem="DE430",
            drop_pintpsr=False,
            timing_package="pint",
        )

    def test_deflate_inflate(self):
        pass

    def test_load_radec_psr(cls):
        """Setup the Pulsar object."""

        # initialize Pulsar class with RA DEC
        psr = Pulsar(
            datadir + "/J0030+0451_RADEC_wrong.par",
            datadir + "/J0030+0451_NANOGrav_9yv1.tim",
            ephem="DE430",
            drop_pintpsr=False,
            timing_package="pint",
        )
        assert "AstrometryEquatorial" in psr.model.components

    def test_no_planet(self):
        """Test exception when incorrect par(tim) file given."""

        with self.assertRaises(ValueError) as context:
            model, toas = get_model_and_toas(
                datadir + "/J0030+0451_NANOGrav_9yv1.gls.par", datadir + "/J0030+0451_NANOGrav_9yv1.tim", planets=False
            )
            Pulsar(model, toas, planets=True)
            msg = "obs_earth_pos is not in toas.table.colnames. Either "
            msg += "`planet` flag is not True in `toas` or further Pint "
            msg += "development to add additional planets is needed."
            self.assertTrue(msg in context.exception)
<commit_msg>Fix test of PINT exception<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
test_pulsar
----------------------------------

Tests for `pulsar` module. Will eventually want to add tests
for time slicing, PINT integration and pickling.
"""

import sys
import os
import shutil
import unittest
import pickle
import pytest

import numpy as np

from enterprise.pulsar import Pulsar
from tests.enterprise_test_data import datadir

import pint.models.timing_model
from pint.models import get_model_and_toas

class TestPulsar(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        """Setup the Pulsar object."""

        # initialize Pulsar class
        cls.psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim")

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree("pickle_dir", ignore_errors=True)

    def test_residuals(self):
        """Check Residual shape."""

        msg = "Residuals shape incorrect"
        assert self.psr.residuals.shape == (4005,), msg

    def test_toaerrs(self):
        """Check TOA errors shape."""

        msg = "TOA errors shape incorrect"
        assert self.psr.toaerrs.shape == (4005,), msg

    def test_toas(self):
        """Check TOA shape."""

        msg = "TOA shape incorrect"
        assert self.psr.toas.shape == (4005,), msg

    def test_stoas(self):
        """Check STOA shape."""

        msg = "stoa shape incorrect"
        assert self.psr.stoas.shape == (4005,), msg

    def test_dm(self):
        """Check DM/DMX access."""

        msg = "dm value incorrect"
        assert self.psr.dm == 13.299393, msg

        msg = "dmx struct incorrect (spotcheck)"
        assert len(self.psr.dmx) == 72, msg
        assert self.psr.dmx["DMX_0001"]["DMX"] == 0.015161863, msg
        assert self.psr.dmx["DMX_0001"]["fit"], msg

    def test_freqs(self):
        """Check frequencies shape."""

        msg = "Frequencies shape incorrect"
        assert self.psr.freqs.shape == (4005,), msg

    def test_flags(self):
        """Check flags shape and content"""

        msg = "Flags shape incorrect"
        assert self.psr.flags["f"].shape == (4005,), msg

        msg = "Flag content or sorting incorrect"
        assert np.all(self.psr._flags["fe"][self.psr._isort] == self.psr.flags["fe"]), msg

    def test_backend_flags(self):
        """Check backend_flags shape and content"""

        msg = "Backend flags shape incorrect"
        assert self.psr.backend_flags.shape == (4005,), msg

        # for the test pulsar, backend should be the same as 'f'
        msg = "Flag content or sorting incorrect"
        assert np.all(self.psr._flags["f"][self.psr._isort] == self.psr.backend_flags), msg

    def test_sky(self):
        """Check Sky location."""

        sky = (1.4023093811712661, 4.9533700839400492)

        msg = "Incorrect sky location"
        assert np.allclose(self.psr.theta, sky[0]), msg
        assert np.allclose(self.psr.phi, sky[1]), msg

    def test_design_matrix(self):
        """Check design matrix shape."""

        msg = "Design matrix shape incorrect."
        assert self.psr.Mmat.shape == (4005, 91), msg

    def test_filter_data(self):
        """Place holder for filter_data tests."""
        assert self.psr.filter_data() is None

    def test_planetssb(self):
        """Place holder for filter_data tests."""
        assert hasattr(self.psr, "planetssb")

    def test_sunssb(self):
        """Place holder for filter_data tests."""
        assert hasattr(self.psr, "sunssb")

    def test_to_pickle(self):
        """Place holder for to_pickle tests."""
        self.psr.to_pickle()
        with open("B1855+09.pkl", "rb") as f:
            pkl_psr = pickle.load(f)

        os.remove("B1855+09.pkl")

        assert np.allclose(self.psr.residuals, pkl_psr.residuals, rtol=1e-10)

        self.psr.to_pickle("pickle_dir")
        with open("pickle_dir/B1855+09.pkl", "rb") as f:
            pkl_psr = pickle.load(f)

        assert np.allclose(self.psr.residuals, pkl_psr.residuals, rtol=1e-10)

    @pytest.mark.skipif(sys.version_info < (3, 8), reason="Requires Python >= 3.8")
    def test_deflate_inflate(self):
        psr = Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.tim")

        dm = psr._designmatrix.copy()

        psr.deflate()
        psr.to_pickle()

        with open("B1855+09.pkl", "rb") as f:
            pkl_psr = pickle.load(f)
        pkl_psr.inflate()

        assert np.allclose(dm, pkl_psr._designmatrix)

        del pkl_psr

        psr.destroy()

        with open("B1855+09.pkl", "rb") as f:
            pkl_psr = pickle.load(f)

        with self.assertRaises(FileNotFoundError):
            pkl_psr.inflate()

    def test_wrong_input(self):
        """Test exception when incorrect par(tim) file given."""

        with self.assertRaises(IOError) as context:
            Pulsar("wrong.par", "wrong.tim")

            msg = "Cannot find parfile wrong.par or timfile wrong.tim!"
            self.assertTrue(msg in context.exception)

    def test_value_error(self):
        """Test exception when unknown argument is given"""

        with self.assertRaises(ValueError):
            Pulsar(datadir + "/B1855+09_NANOGrav_9yv1.gls.par", datadir + "/B1855+09_NANOGrav_9yv1.time")

class TestPulsarPint(TestPulsar):
    @classmethod
    def setUpClass(cls):
        """Setup the Pulsar object."""

        # initialize Pulsar class
        cls.psr = Pulsar(
            datadir + "/B1855+09_NANOGrav_9yv1.gls.par",
            datadir + "/B1855+09_NANOGrav_9yv1.tim",
            ephem="DE430",
            drop_pintpsr=False,
            timing_package="pint",
        )

    def test_deflate_inflate(self):
        pass

    def test_load_radec_psr(self):
        """Setup the Pulsar object."""

        with self.assertRaises(pint.models.timing_model.TimingModelError):
            # initialize Pulsar class with RAJ DECJ and PMLAMBDA, PMBETA
            psr = Pulsar(
                datadir + "/J0030+0451_RADEC_wrong.par",
                datadir + "/J0030+0451_NANOGrav_9yv1.tim",
                ephem="DE430",
                drop_pintpsr=False,
                timing_package="pint",
            )

    def test_no_planet(self):
        """Test exception when incorrect par(tim) file given."""

        with self.assertRaises(ValueError) as context:
            model, toas = get_model_and_toas(
                datadir + "/J0030+0451_NANOGrav_9yv1.gls.par", datadir + "/J0030+0451_NANOGrav_9yv1.tim", planets=False
            )
            Pulsar(model, toas, planets=True)
            msg = "obs_earth_pos is not in toas.table.colnames. Either "
            msg += "`planet` flag is not True in `toas` or further Pint "
            msg += "development to add additional planets is needed."
            self.assertTrue(msg in context.exception)
<|endoftext|>