# coding=utf-8
# Copyright 2017 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for google3.third_party.py.tensor2tensor.data_generators.desc2code."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
from tensor2tensor.data_generators import desc2code
from google3.testing.pybase import googletest


CODE_CPP_IN = """
  #include <iostream>

void main() {  // This comment will be removed
  // This too.
  //
  /* Not    this     one     */
\t
\t
  int a \t\n  =   3;//
//
}

"""

CODE_CPP_OUT = "#include <iostream> void main() { /* Not this one */ int a = 3; }"  # pylint: disable=line-too-loong


class Desc2codeTest(googletest.TestCase):

  def testCppPreprocess(self):
    """Check that the file correctly preprocess the code source."""
    cpp_pb = desc2code.Desc2CodeCppProblem()

    self.assertEqual(  # Add space beween two lines
        cpp_pb.preprocess_target("firstline//comm1\nsecondline//comm2\n"),
        "firstline secondline"
    )
    # Checking for boths comments and spaces
    self.assertEqual(cpp_pb.preprocess_target(CODE_CPP_IN), CODE_CPP_OUT)
    self.assertEqual(
        cpp_pb.preprocess_target("  not removed //abcd  "),
        "not removed //abcd"
    )


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

@@ -616,61 +616,8 @@
 for 
-google3.third_party.py.tensor2tensor.data_generators.
 desc
@@ -814,54 +814,32 @@
 ode%0A
-from google3.testing.pybase import googletest%0A
+%0Aimport tensorflow as tf
 %0A%0ACO
@@ -1031,16 +1031,17 @@
 P_OUT = 
+(
 %22#includ
@@ -1098,47 +1098,33 @@
 a = 
-3; %7D%22  # pylint: disable=line-too-loong
+%22%0A                %223; %7D%22)
 %0A%0A%0Ac
@@ -1142,22 +1142,19 @@
 odeTest(
-google
+tf.
 test.Tes
@@ -1464,21 +1464,16 @@
 ondline%22
-%0A    
 )%0A    # 
@@ -1695,21 +1695,16 @@
  //abcd%22
-%0A    
 )%0A%0A%0Aif _
@@ -1732,14 +1732,11 @@
 :%0A  
-google
+tf.
 test

fix some open-source imports