import tensorflow as tf

class MobilenetV1Test(tf.test.TestCase):

    def testBuildClassificationNetwork(self):
        batch_size = 5
        height, width = 224, 224
        num_classes = 1000

        inputs = tf.random_uniform((batch_size, height, width, 3))
        logits, end_points = mobilenet_v1.mobilenet_v1(inputs, num_classes)
        self.assertTrue(logits.op.name.startswith(
            'MobilenetV1/Logits/SpatialSqueeze'))
        self.assertListEqual(logits.get_shape().as_list(),
                             [batch_size, num_classes])
        self.assertTrue('Predictions' in end_points)
