From 27e2307d0c05ad134df166f264d1b993d823930e Mon Sep 17 00:00:00 2001
From: Rasmus Wriedt Larsen <rasmuswl@github.com>
Date: Fri, 17 Feb 2023 15:48:18 +0100
Subject: [PATCH] Python: Add import regression for refined variable

---
 .../experimental/import-resolution/importflow.ql     |  3 +++
 .../ql/test/experimental/import-resolution/main.py   |  4 ++++
 .../test/experimental/import-resolution/refined.py   | 12 ++++++++++++
 3 files changed, 19 insertions(+)
 create mode 100644 python/ql/test/experimental/import-resolution/refined.py

diff --git a/python/ql/test/experimental/import-resolution/importflow.ql b/python/ql/test/experimental/import-resolution/importflow.ql
index e6e51afa963..94f5e4b0ec8 100644
--- a/python/ql/test/experimental/import-resolution/importflow.ql
+++ b/python/ql/test/experimental/import-resolution/importflow.ql
@@ -11,6 +11,9 @@ private class SourceString extends DataFlow::Node {
   SourceString() {
     this.asExpr().(StrConst).getText() = contents and
     this.asExpr().getParent() instanceof Assign
+    or
+    this.asExpr().(ClassExpr).getInnerScope().getName() = "SOURCE" and
+    contents = "SOURCE"
   }
 
   string getContents() { result = contents }
diff --git a/python/ql/test/experimental/import-resolution/main.py b/python/ql/test/experimental/import-resolution/main.py
index 840da95a03c..6ff5506cab4 100644
--- a/python/ql/test/experimental/import-resolution/main.py
+++ b/python/ql/test/experimental/import-resolution/main.py
@@ -93,6 +93,10 @@ check("subpackage2_attr", subpackage2_attr, "subpackage2_attr", globals()) #$ pr
 from if_then_else import if_then_else_defined
 check("if_then_else_defined", if_then_else_defined, "if_defined", globals()) #$ prints=if_defined prints=else_defined_1 prints=else_defined_2
 
+# check that refined definitions are handled correctly
+import refined # $ imports=refined as=refined
+check("refined.SOURCE", refined.SOURCE, refined.SOURCE, globals()) #$ MISSING: prints=SOURCE
+
 exit(__file__)
 
 print()
diff --git a/python/ql/test/experimental/import-resolution/refined.py b/python/ql/test/experimental/import-resolution/refined.py
new file mode 100644
index 00000000000..a687ce15655
--- /dev/null
+++ b/python/ql/test/experimental/import-resolution/refined.py
@@ -0,0 +1,12 @@
+from trace import *
+enter(__file__)
+
+class SOURCE(object): pass
+
+check("SOURCE", SOURCE, SOURCE, globals()) #$ prints=SOURCE
+
+SOURCE.foo = 42
+
+check("SOURCE", SOURCE, SOURCE, globals()) #$ prints=SOURCE
+
+exit(__file__)