"""
.. module:: MalConTestSuite
	:synopsis: MalConv basic unit tests

.. moduleauthor:: Luca Demetrio <luca.demetrio@dibris.unige.it>
"""

import unittest
import torch
from secml.array import CArray

from secml_malware.models import MalConv
from secml_malware.attack.whitebox.tests.malware_test_base import End2EndBaseTests


class MalConvTestSuite(End2EndBaseTests):
	def setUp(self):
		super(MalConvTestSuite, self).setUp()
		self.m = MalConv()

	def test_empty_baseline_is_good_ware(self):
		self.m.load_simplified_model(self.ember_path)
		y_pred = self.m(torch.Tensor(self.baseline).long())
		if y_pred.is_cuda:
			y_pred = y_pred.cpu()
		y = int(y_pred.detach().numpy()[0][0] > 0.5)
		self.assertEqual(y, 0)

	def test_malconv_load_ember_no_path(self):
		self.classifier.load_pretrained_model()
		y = self.classifier.predict(
			CArray(torch.Tensor(self.baseline).long()), return_decision_function=False
		)
		self.assertEqual(y, 0)

	def test_empty_baseline_pytorch_classifier_wrapper(self):
		self.classifier.load_pretrained_model(self.ember_path)
		y = self.classifier.predict(
			CArray(torch.Tensor(self.baseline).long()), return_decision_function=False
		)
		self.assertEqual(y, 0)


if __name__ == "__main__":
	unittest.main()
