Merge "Check tests dir in PythonBinaryHostTest"
diff --git a/test_framework/com/android/tradefed/testtype/python/PythonBinaryHostTest.java b/test_framework/com/android/tradefed/testtype/python/PythonBinaryHostTest.java
index 1d23bfe..91893ab 100644
--- a/test_framework/com/android/tradefed/testtype/python/PythonBinaryHostTest.java
+++ b/test_framework/com/android/tradefed/testtype/python/PythonBinaryHostTest.java
@@ -178,15 +178,18 @@
     public final void run(TestInformation testInfo, ITestInvocationListener listener)
             throws DeviceNotAvailableException {
         mTestInfo = testInfo;
-        File hostTestDir = mTestInfo.executionFiles().get(FilesKey.HOST_TESTS_DIRECTORY);
-        if (hostTestDir.exists()) {
-            File libDir = new File(hostTestDir, "lib");
+        File testDir = mTestInfo.executionFiles().get(FilesKey.HOST_TESTS_DIRECTORY);
+        if (testDir == null || !testDir.exists()) {
+            testDir = mTestInfo.executionFiles().get(FilesKey.TESTS_DIRECTORY);
+        }
+        if (testDir != null && testDir.exists()) {
+            File libDir = new File(testDir, "lib");
             List<String> ldLibraryPath = new ArrayList<>();
             if (libDir.exists()) {
                 ldLibraryPath.add(libDir.getAbsolutePath());
             }
 
-            File lib64Dir = new File(hostTestDir, "lib64");
+            File lib64Dir = new File(testDir, "lib64");
             if (lib64Dir.exists()) {
                 ldLibraryPath.add(lib64Dir.getAbsolutePath());
             }
diff --git a/tests/src/com/android/tradefed/testtype/python/PythonBinaryHostTestTest.java b/tests/src/com/android/tradefed/testtype/python/PythonBinaryHostTestTest.java
index 8cf07d0..b0f8d09 100644
--- a/tests/src/com/android/tradefed/testtype/python/PythonBinaryHostTestTest.java
+++ b/tests/src/com/android/tradefed/testtype/python/PythonBinaryHostTestTest.java
@@ -348,7 +348,7 @@
 
     /** Test running the python tests when shared lib is available in HOST_TESTS_DIRECTORY. */
     @Test
-    public void testRun_withSharedLib() throws Exception {
+    public void testRun_withSharedLibInHostTestsDir() throws Exception {
         File hostTestsDir = FileUtil.createTempDir("host-test-cases");
         mTestInfo.executionFiles().put(FilesKey.HOST_TESTS_DIRECTORY, hostTestsDir);
         File binary = FileUtil.createTempFile("python-dir", "", hostTestsDir);
@@ -391,6 +391,51 @@
         }
     }
 
+    /** Test running the python tests when shared lib is available in TESTS_DIRECTORY. */
+    @Test
+    public void testRun_withSharedLib() throws Exception {
+        File testsDir = FileUtil.createTempDir("host-test-cases");
+        mTestInfo.executionFiles().put(FilesKey.TESTS_DIRECTORY, testsDir);
+        File binary = FileUtil.createTempFile("python-dir", "", testsDir);
+        File lib = new File(testsDir, "lib");
+        lib.mkdirs();
+        File lib64 = new File(testsDir, "lib64");
+        lib64.mkdirs();
+
+        try {
+            OptionSetter setter = new OptionSetter(mTest);
+            setter.setOptionValue("python-binaries", binary.getAbsolutePath());
+            mMockRunUtil.setEnvVariable(
+                    PythonBinaryHostTest.LD_LIBRARY_PATH,
+                    lib.getAbsolutePath() + ":" + lib64.getAbsolutePath());
+            expectedAdbPath(mFakeAdb);
+
+            CommandResult res = new CommandResult();
+            res.setStatus(CommandStatus.SUCCESS);
+            res.setStderr("TEST_RUN_STARTED {\"testCount\": 5, \"runName\": \"TestSuite\"}");
+            EasyMock.expect(
+                            mMockRunUtil.runTimedCmd(
+                                    EasyMock.anyLong(), EasyMock.eq(binary.getAbsolutePath())))
+                    .andReturn(res);
+            mMockListener.testRunStarted(
+                    EasyMock.eq(binary.getName()),
+                    EasyMock.eq(5),
+                    EasyMock.eq(0),
+                    EasyMock.anyLong());
+            mMockListener.testLog(
+                    EasyMock.eq(binary.getName() + "-stderr"),
+                    EasyMock.eq(LogDataType.TEXT),
+                    EasyMock.anyObject());
+            EasyMock.expect(mMockDevice.getIDevice()).andReturn(new StubDevice("serial"));
+
+            EasyMock.replay(mMockRunUtil, mMockBuildInfo, mMockListener, mMockDevice);
+            mTest.run(mTestInfo, mMockListener);
+            EasyMock.verify(mMockRunUtil, mMockBuildInfo, mMockListener, mMockDevice);
+        } finally {
+            FileUtil.recursiveDelete(testsDir);
+        }
+    }
+
     /**
      * If the binary returns an exception status, we should throw a runtime exception since
      * something went wrong with the binary setup.