<commit_before>from core.db.query_rewriter import SQLQueryRewriter

from django.db.models import signals
from django.test import TestCase

import factory
import sqlparse
from mock import patch


class QueryRewriter(TestCase):

    """Tests all the query rewriter operations in query_rewriter.py."""

    @factory.django.mute_signals(signals.pre_save)
    def setUp(self):
        self.repo_base = "test_repobase"
        self.user = "test_user"
        self.query_rewriter = SQLQueryRewriter(self.repo_base, self.user)

        self.mock_connection = self.create_patch(
            'core.db.manager.DataHubConnection')

    def create_patch(self, name):
        # helper method for creating patches
        patcher = patch(name)
        thing = patcher.start()
        self.addCleanup(patcher.stop)
        return thing

    def test_extract_table_info(self):
        valid_table_token = "repo.table"
        expected_result = ["repo", "table", None]
        self.assertEqual(
            self.query_rewriter.extract_table_info(valid_table_token),
            expected_result)

        valid_table_token = "repobase.repo.table"
        expected_result = ["repo", "table", "repobase"]
        self.assertEqual(
            self.query_rewriter.extract_table_info(valid_table_token),
            expected_result)

        invalid_table_token = "testtable"
        exception_raised = False
        try:
            self.query_rewriter.extract_table_info(invalid_table_token)
        except Exception:
            exception_raised = True
        self.assertEquals(exception_raised, True)

        invalid_table_token = "table1.table2.table3.table4"
        exception_raised = False
        try:
            self.query_rewriter.extract_table_info(invalid_table_token)
        except Exception:
            exception_raised = True
        self.assertEquals(exception_raised, True)

    def test_extract_table_token(self):
        query = "SELECT * from repo1.table1 as tbl1"
        token = sqlparse.parse(query)[0].tokens[6]
        expected_result = [(["repo1", "table1", None], "as tbl1")]
        self.assertEqual(
            self.query_rewriter.extract_table_token(token), expected_result)

        query = ("SELECT * from repo1.table1 as tbl1, repo2.table2 as tbl2 "
                 "where ... ")
        token = sqlparse.parse(query)[0].tokens[6]
        expected_result = [(["repo1", "table1", None], "as tbl1"),
                           (["repo2", "table2", None], "as tbl2")]
        self.assertEqual(
            self.query_rewriter.extract_table_token(token), expected_result)

        query = "SELECT * from repo1.table1 tbl1, repo2.table2 tbl2 where ... "
        token = sqlparse.parse(query)[0].tokens[6]
        expected_result = [(["repo1", "table1", None], "tbl1"),
                           (["repo2", "table2", None], "tbl2")]
        self.assertEqual(
            self.query_rewriter.extract_table_token(token), expected_result)

    def test_extract_table_string(self):
        valid_table_string = "repo.table"
        expected_result = (["repo", "table", None], '')
        self.assertEqual(
            self.query_rewriter.extract_table_string(valid_table_string),
            expected_result)

        valid_table_string = "repo.table test"
        expected_result = (["repo", "table", None], 'test')
        self.assertEqual(
            self.query_rewriter.extract_table_string(valid_table_string),<commit_msg>Update query rewriter test for new sqlparse.<commit_after>from core.db.query_rewriter import SQLQueryRewriter

from django.db.models import signals
from django.test import TestCase

import factory
import sqlparse
from mock import patch


class QueryRewriter(TestCase):
    """Tests all the query rewriter operations in query_rewriter.py."""

    @factory.django.mute_signals(signals.pre_save)
    def setUp(self):
        self.repo_base = "test_repobase"
        self.user = "test_user"
        self.query_rewriter = SQLQueryRewriter(self.repo_base, self.user)

        self.mock_connection = self.create_patch(
            'core.db.manager.DataHubConnection')

    def create_patch(self, name):
        # helper method for creating patches
        patcher = patch(name)
        thing = patcher.start()
        self.addCleanup(patcher.stop)
        return thing

    def test_extract_table_info(self):
        valid_table_token = "<PASSWORD>"
        expected_result = ["repo", "table", None]
        self.assertEqual(
            self.query_rewriter.extract_table_info(valid_table_token),
            expected_result)

        valid_table_token = "<PASSWORD>"
        expected_result = ["repo", "table", "repobase"]
        self.assertEqual(
            self.query_rewriter.extract_table_info(valid_table_token),
            expected_result)

        invalid_table_token = "<PASSWORD>"
        exception_raised = False
        try:
            self.query_rewriter.extract_table_info(invalid_table_token)
        except Exception:
            exception_raised = True
        self.assertEquals(exception_raised, True)

        invalid_table_token = "table1.table2.table3.table4"
        exception_raised = False
        try:
            self.query_rewriter.extract_table_info(invalid_table_token)
        except Exception:
            exception_raised = True
        self.assertEquals(exception_raised, True)

    def test_extract_table_token(self):
        query = "SELECT * from repo1.table1 as tbl1"
        token = sqlparse.parse(query)[0].tokens[6]
        expected_result = [(["repo1", "table1", None], "as tbl1")]
        self.assertEqual(
            self.query_rewriter.extract_table_token(token), expected_result)

        query = ("SELECT * from repo1.table1 as tbl1, repo2.table2 as tbl2 "
                 "where ... ")
        token = sqlparse.parse(query)[0].tokens[6]
        expected_result = [(["repo1", "table1", None], "as tbl1"),
                           (["repo2", "table2", None], "as tbl2 where ...")]
        self.assertEqual(
            self.query_rewriter.extract_table_token(token), expected_result)

        query = "SELECT * from repo1.table1 tbl1, repo2.table2 tbl2 where ... "
        token = sqlparse.parse(query)[0].tokens[6]
        expected_result = [(["repo1", "table1", None], "tbl1"),
                           (["repo2", "table2", None], "tbl2")]
        self.assertEqual(
            self.query_rewriter.extract_table_token(token), expected_result)

    def test_extract_table_string(self):
        valid_table_string = "repo.table"
        expected_result = (["repo", "table", None], '')
        self.assertEqual(
            self.query_rewriter.extract_table_string(valid_table_string),
            expected_result)

        valid_table_string = "repo.table test"
        expected_result = (["repo", "table", None], 'test')
        self.assertEqual(
            self.query_rewriter.extract_table_string(valid_table_string),
{'fraction_non_alphanumeric': 0.10613958876339415, 'fraction_numerical': 0.011584129742253113, 'mean_word_length': 2.644854881266491, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2960661', 'n_tokens_mistral': 2033, 'n_tokens_neox': 1966, 'n_words': 386}