# encoding: utf8
from __future__ import unicode_literals
import pytest
from flexmock import flexmock
from hypothesis import given, strategies
import abc

from .... import vec2vec
from ....ops import NumpyOps


@pytest.fixture
def model_with_no_args():
    model = vec2vec.Affine(ops=NumpyOps())
    return model


def test_Affine_default_name(model_with_no_args):
    assert model_with_no_args.name == 'affine'


def test_Affine_defaults_to_cpu(model_with_no_args):
    assert isinstance(model_with_no_args.ops, NumpyOps)


def test_Affine_defaults_to_no_layers(model_with_no_args):
    assert model_with_no_args.layers == []


def test_Affine_defaults_to_param_descriptions(model_with_no_args):
    W_desc, b_desc = model_with_no_args.describe_params
    xavier_init = model_with_no_args.ops.xavier_uniform_init
    assert W_desc == ('W-affine', (None, None), xavier_init)
    assert b_desc == ('b-affine', (None,), None)


def test_Model_defaults_to_no_output_shape(model_with_no_args):
    assert model_with_no_args.output_shape == None
 

def test_Model_defaults_to_no_input_shape(model_with_no_args):
    assert model_with_no_args.input_shape == None


def test_Model_defaults_to_0_size(model_with_no_args):
    assert model_with_no_args.size == None

@@ -69,20 +69,16 @@
 st%0Afrom 
-flex
 mock imp
@@ -85,16 +85,19 @@
 ort 
-flexmock
+Mock, patch
 %0Afro
@@ -155,24 +155,39 @@
 rom ....
+_classes.affine
  import 
 vec2vec%0A
@@ -178,23 +178,22 @@
  import 
-vec2vec
+Affine
 %0Afrom ..
@@ -245,398 +245,543 @@
 odel
-_with_no_args():%0A    model = vec2vec.Affine(ops=NumpyOps())%0A    return model%0A%0A%0Adef test_
+():%0A    orig_desc = dict(Affine.descriptions)%0A    orig_on_init = list(
 Affine
-_default_name(model_with_no_args):%0A    assert model_with_no_args.name == 'affine'%0A%0A%0Adef test_Affine_defaults_to_cpu(model_with_no_args):%0A    assert isinstance(model_with_no_args.ops, NumpyOps)%0A%0A%0Adef test_Affine_defaults_to_no_layers(model_with_no_args):%0A    assert model_with_no_args.layers == %5B%5D
+.on_init_hooks)%0A    Affine.descriptions = %7B%0A        name: Mock(desc) for (name, desc) in Affine.descriptions.items()%0A    %7D%0A    Affine.on_init_hooks = %5BMock(hook) for hook in Affine.on_init_hooks%5D%0A    model = Affine()%0A    for attr in model.descriptions:%0A        setattr(model, attr, None)%0A    Affine.descriptions = dict(orig_desc)%0A    Affine.on_init_hooks = orig_on_init%0A    return model%0A%0A%0Adef test_Affine_default_name(model):%0A    assert model.name == 'affine'
 %0A%0A%0Ad
@@ -799,25 +799,21 @@
 ine_
+calls_
 default
-s_to_param
 _des
@@ -831,483 +831,317 @@
 odel
-_with_no_args):%0A    W_desc, b_desc = model_with_no_args.describe_params%0A    xavier_init = model_with_no_args.ops.xavier_uniform_init%0A    assert W_desc == ('W-affine', (None, None), xavier_init)%0A    assert b_desc == ('b-affine', (None,), None)%0A%0A%0Adef test_M
+):%0A    assert len(model.descriptions) == 5%0A    for name, desc in model.descriptions.items():%0A        desc.assert_called()%0A    assert 'nB' in m
 odel
-_defaults_to_no_output_shape(model_with_no_args):%0A    assert model_with_no_args.output_shape == None%0A %0A%0Adef test_M
+.descriptions%0A    assert 'nI' in model.descriptions%0A    assert 'nO' in model.descriptions%0A    assert 'W' in m
 odel
-_defaults_to_no_input_shape(model_with_no_args):%0A    assert model_with_no_args.input_shape == None
+.descriptions%0A    assert 'b' in model.descriptions
 %0A%0A%0Ad
@@ -1152,93 +1152,102 @@
 est_
-Model_defaults_to_0_size(model_with_no_args):%0A    assert model_with_no_args.size == None
+Affine_calls_init_hooks(model):%0A    for hook in model.on_init_hooks:%0A        hook.assert_called()
 %0A

Test Affine init calls hooks correctly, and sets descriptors