Refactor server handshake test

mtglib/internal/obfuscated2/server_handshake_test.go

 
 type ServerHandshakeTestSuite struct {
 	suite.Suite
+
+	connMock  *testlib.NetConnMock
+	proxyConn *obfuscated2.Conn
+	encryptor cipher.Stream
+	decryptor cipher.Stream
 }
 
-func (suite *ServerHandshakeTestSuite) TestOk() {
+func (suite *ServerHandshakeTestSuite) SetupTest() {
 	buf := &bytes.Buffer{}
-	connMock := &testlib.NetConnMock{}
+	suite.connMock = &testlib.NetConnMock{}
 
 	encryptor, decryptor, err := obfuscated2.ServerHandshake(buf)
-	suite.NotNil(encryptor)
-	suite.NotNil(decryptor)
 	suite.NoError(err)
 
-	proxyConn := &obfuscated2.Conn{
-		Conn:      connMock,
+	suite.proxyConn = &obfuscated2.Conn{
+		Conn:      suite.connMock,
 		Encryptor: encryptor,
 		Decryptor: decryptor,
 	}
 	serverEncrypted := buf.Bytes()
 
 	decBlock, _ := aes.NewCipher(serverEncrypted[8 : 8+32])
-	serverDecryptor := cipher.NewCTR(decBlock, serverEncrypted[8+32:8+32+16])
+	suite.decryptor = cipher.NewCTR(decBlock, serverEncrypted[8+32:8+32+16])
+
 	serverDecrypted := make([]byte, len(serverEncrypted))
-	serverDecryptor.XORKeyStream(serverDecrypted, serverEncrypted)
+	suite.decryptor.XORKeyStream(serverDecrypted, serverEncrypted)
 
 	suite.Equal("3d3d3Q",
 		base64.RawStdEncoding.EncodeToString(serverDecrypted[8+32+16:8+32+16+4]))
 	}
 
 	encBlock, _ := aes.NewCipher(serverEncryptedReverted[8 : 8+32])
-	serverEncryptor := cipher.NewCTR(encBlock, serverEncryptedReverted[8+32:8+32+16])
+	suite.encryptor = cipher.NewCTR(encBlock, serverEncryptedReverted[8+32:8+32+16])
+}
+
+func (suite *ServerHandshakeTestSuite) TearDownTest() {
+	suite.connMock.AssertExpectations(suite.T())
+}
 
-	messageFromTelegram := []byte{1, 2, 3, 4, 5}
-	// messageToTelegram := []byte{10, 11, 13, 14}
-	bufferToRead := make([]byte, 5)
+func (suite *ServerHandshakeTestSuite) TestSendToTelegram() {
+	messageToTelegram := []byte{10, 11, 12, 13, 14, 'a'}
 
-	connMock.
-		On("Read", mock.Anything).
-		Return(5, nil).
+	suite.connMock.
+		On("Write", mock.Anything).
+		Return(len(messageToTelegram), nil).
 		Once().
 		Run(func(args mock.Arguments) {
-			messageToRead := make([]byte, len(messageFromTelegram))
-			serverEncryptor.XORKeyStream(messageToRead, messageFromTelegram)
-			copy(args.Get(0).([]byte), messageToRead)
+			message := make([]byte, len(messageToTelegram))
+			suite.decryptor.XORKeyStream(message, args.Get(0).([]byte))
+			suite.Equal(messageToTelegram, message)
 		})
 
-	n, err := proxyConn.Read(bufferToRead)
-	suite.EqualValues(5, n)
+	n, err := suite.proxyConn.Write(messageToTelegram)
+	suite.EqualValues(len(messageToTelegram), n)
 	suite.NoError(err)
-	suite.Equal(messageFromTelegram, bufferToRead)
+}
 
-	messageToTelegram := []byte{10, 11, 12, 13, 14}
+func (suite *ServerHandshakeTestSuite) TestRecieveFromTelegram() {
+	messageFromTelegram := []byte{10, 11, 12, 13, 14, 'a'}
+	buffer := make([]byte, len(messageFromTelegram))
 
-	connMock.
-		On("Write", mock.Anything).
-		Return(5, nil).
+	suite.connMock.
+		On("Read", mock.Anything).
+		Return(len(messageFromTelegram), nil).
 		Once().
 		Run(func(args mock.Arguments) {
-			message := make([]byte, len(messageToTelegram))
-			serverDecryptor.XORKeyStream(message, args.Get(0).([]byte))
-			suite.Equal(messageToTelegram, message)
+			message := make([]byte, len(messageFromTelegram))
+			suite.encryptor.XORKeyStream(message, messageFromTelegram)
+			copy(args.Get(0).([]byte), message)
 		})
 
-	n, err = proxyConn.Write(messageToTelegram)
-	suite.EqualValues(5, n)
+	n, err := suite.proxyConn.Read(buffer)
+	suite.EqualValues(len(messageFromTelegram), n)
 	suite.NoError(err)
-
-	connMock.AssertExpectations(suite.T())
+	suite.Equal(messageFromTelegram, buffer)
 }
 
 func TestServerHandshake(t *testing.T) {